4 files changed, 123 insertions, 5 deletions

diff --git a/src/Network/BitTorrent/DHT/Message.hs b/src/Network/BitTorrent/DHT/Message.hs
index 1f835fa6..44dc9b2f 100644
--- a/src/Network/BitTorrent/DHT/Message.hs
+++ b/src/Network/BitTorrent/DHT/Message.hs
@@ -55,6 +55,7 @@
 --   For Kamelia messages see original Kademlia paper:
 --   <http://pdos.csail.mit.edu/~petar/papers/maymounkov-kademlia-lncs.pdf>
 --
+{-# LANGUAGE CPP #-}
 {-# LANGUAGE DeriveDataTypeable #-}
 {-# LANGUAGE FlexibleInstances  #-}
 {-# LANGUAGE MultiParamTypeClasses #-}
@@ -73,6 +74,8 @@ module Network.BitTorrent.DHT.Message
        , FindNode (..)
        , NodeFound (..)
 
+
+#ifdef VERSION_bencoding
          -- ** get_peers
        , PeerList
        , GetPeers (..)
@@ -81,12 +84,23 @@ module Network.BitTorrent.DHT.Message
          -- ** announce_peer
        , Announce  (..)
        , Announced (..)
+#endif
        ) where
 
 import Control.Applicative
 import Data.Bool
+#ifdef VERSION_bencoding
 import Data.BEncode as BE
 import Data.BEncode.BDict as BDict
+import Network.BitTorrent.Address
+#else
+import qualified Data.Tox as Tox
+import Data.Tox (NodeId)
+import Data.Word
+import Control.Monad
+import Network.KRPC.Method
+import Network.BitTorrent.Address hiding (NodeId)
+#endif
 import Data.ByteString (ByteString)
 import Data.List as L
 import Data.Monoid
@@ -97,7 +111,6 @@ import Network.KRPC
 import Data.Maybe
 
 import Data.Torrent (InfoHash)
-import Network.BitTorrent.Address
 import Network.BitTorrent.DHT.Token
 import Network.KRPC ()
 
@@ -105,6 +118,10 @@ import Network.KRPC ()
 -- envelopes
 -----------------------------------------------------------------------}
 
+#ifndef VERSION_bencoding
+type BKey = ByteString
+#endif
+
 node_id_key :: BKey
 node_id_key = "id"
 
@@ -112,6 +129,7 @@ read_only_key :: BKey
 read_only_key = "ro"
 
 
+#ifdef VERSION_bencoding
 -- | All queries have an \"id\" key and value containing the node ID
 -- of the querying node.
 data Query a = Query
@@ -134,7 +152,11 @@ instance BEncode a => BEncode (Query a) where
     Query <$> fromDict (field (req node_id_key)) v
           <*> fromDict (fromMaybe False <$>? read_only_key) v
           <*> fromBEncode v
+#else
+data Query a = Query a
+#endif
 
+#ifdef VERSION_bencoding
 -- | All responses have an \"id\" key and value containing the node ID
 -- of the responding node.
 data Response a = Response
@@ -150,7 +172,9 @@ instance BEncode a => BEncode (Response a) where
   fromBEncode b = fromQuery <$> fromBEncode b
     where
       fromQuery (Query nid _ a) = Response nid a
-
+#else
+data Response a = Response a
+#endif
 
 {-----------------------------------------------------------------------
 -- ping method
@@ -158,16 +182,45 @@ instance BEncode a => BEncode (Response a) where
 
 -- | The most basic query is a ping. Ping query is used to check if a
 -- quered node is still alive.
+#ifdef VERSION_bencoding
 data Ping = Ping
+#else
+data Ping = Ping Tox.Nonce8
+#endif
   deriving (Show, Eq, Typeable)
 
+#ifdef VERSION_bencoding
 instance BEncode Ping where
   toBEncode Ping = toDict endDict
   fromBEncode _  = pure Ping
+#else
+instance Serialize (Query Ping) where
+    get = do
+        b <- get
+        when ( (b::Word8) /= 0) $ fail "Bad ping request"
+        nonce <- get
+        return $ Query (Ping nonce)
+    put (Query (Ping nonce)) = do
+        put (0 :: Word8)
+        put nonce
+instance Serialize (Response Ping) where
+    get = do
+        b <- get
+        when ( (b::Word8) /= 1) $ fail "Bad ping response"
+        nonce <- get
+        return $ Response (Ping nonce)
+    put (Response (Ping nonce)) = do
+        put (1 :: Word8)
+        put nonce
+#endif
 
 -- | \"q\" = \"ping\"
 instance KRPC (Query Ping) (Response Ping) where
+#ifdef VERSION_bencoding
   method = "ping"
+#else
+  method = Method Tox.Ping -- response: Tox.Pong
+#endif
 
 {-----------------------------------------------------------------------
 -- find_node method
@@ -175,21 +228,41 @@ instance KRPC (Query Ping) (Response Ping) where
 
 -- | Find node is used to find the contact information for a node
 -- given its ID.
+#ifdef VERSION_bencoding
 newtype FindNode = FindNode NodeId
+#else
+data FindNode = FindNode NodeId Tox.Nonce8 -- Tox: Get Nodes
+#endif
   deriving (Show, Eq, Typeable)
 
 target_key :: BKey
 target_key = "target"
 
+#ifdef VERSION_bencoding
 instance BEncode FindNode where
   toBEncode (FindNode nid) = toDict   $ target_key .=! nid .: endDict
   fromBEncode              = fromDict $ FindNode  <$>! target_key
+#else
+instance Serialize (Query FindNode) where
+    get = do
+        nid <- get
+        nonce <- get
+        return $ Query (FindNode nid nonce)
+    put (Query (FindNode nid nonce)) = do
+        put nid
+        put nonce
+#endif
 
 -- | When a node receives a 'FindNode' query, it should respond with a
 -- the compact node info for the target node or the K (8) closest good
 -- nodes in its own routing table.
 --
+#ifdef VERSION_bencoding
 newtype NodeFound ip = NodeFound [NodeInfo ip]
+#else
+data NodeFound ip = NodeFound [Tox.NodeFormat] Tox.Nonce8
+#endif
+-- Tox: send_nodes
   deriving (Show, Eq, Typeable)
 
 nodes_key :: BKey
@@ -200,6 +273,7 @@ from4 :: forall s. Address s => NodeInfo IPv4 -> Either String (NodeInfo s)
 from4 n = maybe (Left "Error converting IPv4") Right
                     $ traverse (fromAddr :: IPv4 -> Maybe s) n
 
+#ifdef VERSION_bencoding
 binary :: Serialize a => BKey -> BE.Get [a]
 binary k = field (req k) >>= either (fail . format) return .
     runGet (many get)
@@ -213,12 +287,31 @@ instance Address ip => BEncode (NodeFound ip) where
 
   -- TODO: handle IPv6 by reading the "nodes6" key (see bep 32)
   fromBEncode bval = NodeFound <$> (traverse from4 =<< fromDict (binary nodes_key) bval)
+#else
+instance Serialize (Response (NodeFound ip)) where
+    get = do
+        count <- get :: Get Word8
+        nodes <- sequence $ replicate (fromIntegral count) get
+        nonce <- get :: Get Tox.Nonce8
+        return $ Response $ NodeFound nodes nonce
+
+    put (Response (NodeFound nodes nonce)) = do
+        put (fromIntegral (length nodes) :: Word8)
+        mapM_ put nodes
+        put nonce
+
+#endif
 
 -- | \"q\" == \"find_node\"
 instance (Address ip, Typeable ip)
       => KRPC (Query FindNode) (Response (NodeFound ip)) where
+#ifdef VERSION_bencoding
   method = "find_node"
+#else
+  method = Method Tox.GetNodes -- response: Tox.SendNodes
+#endif
 
+#ifdef VERSION_bencoding
 {-----------------------------------------------------------------------
 -- get_peers method
 -----------------------------------------------------------------------}
@@ -354,3 +447,6 @@ instance BEncode Announced where
 -- | \"q" = \"announce\"
 instance KRPC (Query Announce) (Response Announced) where
   method = "announce_peer"
+
+-- endif VERSION_bencoding
+#endif
diff --git a/src/Network/BitTorrent/DHT/Query.hs b/src/Network/BitTorrent/DHT/Query.hs
index 5345f8b1..c7e48920 100644
--- a/src/Network/BitTorrent/DHT/Query.hs
+++ b/src/Network/BitTorrent/DHT/Query.hs
@@ -96,7 +96,10 @@ import qualified Network.BitTorrent.DHT.Search as Search
 
 nodeHandler :: Address ip => KRPC (Query a) (Response b)
            => (NodeAddr ip -> a -> DHT ip b) -> NodeHandler ip
-nodeHandler action = handler $ \ sockAddr (Query remoteId read_only q) -> do
+nodeHandler action = handler $ \ sockAddr qry -> do
+  let remoteId  = queringNodeId qry
+      read_only = queryIsReadOnly qry
+      q         = queryParams qry
   case fromSockAddr sockAddr of
     Nothing    -> throwIO BadAddress
     Just naddr -> do
@@ -119,6 +122,7 @@ findNodeH :: Address ip => NodeHandler ip
 findNodeH = nodeHandler $ \ _ (FindNode nid) -> do
   NodeFound <$> getClosest nid
 
+#ifdef VERSION_bencoding
 -- | Default 'GetPeers' handler.
 getPeersH :: Ord ip => Address ip => NodeHandler ip
 getPeersH = nodeHandler $ \ naddr (GetPeers ih) -> do
@@ -141,6 +145,11 @@ announceH = nodeHandler $ \ naddr @ NodeAddr {..} (Announce {..}) -> do
 -- | Includes all default query handlers.
 defaultHandlers :: Ord ip => Address ip => [NodeHandler ip]
 defaultHandlers = [pingH, findNodeH, getPeersH, announceH]
+#else
+-- | Includes all default query handlers.
+defaultHandlers :: Ord ip => Address ip => [NodeHandler ip]
+defaultHandlers = [pingH, findNodeH]
+#endif
 
 {-----------------------------------------------------------------------
 --  Basic queries
@@ -165,6 +174,7 @@ findNodeQ key NodeInfo {..} = do
     <> T.pack (L.unlines $ L.map ((' ' :) . show . pPrint) closest)
   return $ Right closest
 
+#ifdef VERSION_bencoding
 getPeersQ :: Address ip => InfoHash -> Iteration ip PeerAddr
 getPeersQ topic NodeInfo {..} = do
   GotPeers {..} <- GetPeers topic <@> nodeAddr
@@ -184,6 +194,7 @@ announceQ ih p NodeInfo {..} = do
     Right _ -> do -- TODO *probably* add to peer cache
       Announced <- Announce False ih Nothing p grantedToken <@> nodeAddr
       return (Right [nodeAddr])
+#endif
 
 {-----------------------------------------------------------------------
 --  Iterative queries
diff --git a/src/Network/BitTorrent/DHT/Session.hs b/src/Network/BitTorrent/DHT/Session.hs
index bad783a5..82926b28 100644
--- a/src/Network/BitTorrent/DHT/Session.hs
+++ b/src/Network/BitTorrent/DHT/Session.hs
@@ -55,6 +55,7 @@ module Network.BitTorrent.DHT.Session
        , getTable
        , getClosest
 
+#ifdef VERSION_bencoding
          -- ** Peer storage
        , insertPeer
        , getPeerList
@@ -64,6 +65,7 @@ module Network.BitTorrent.DHT.Session
        , savePeerStore
        , mergeSavedPeers
        , allPeers
+#endif
 
          -- ** Messaging
        , queryParallel
@@ -482,6 +484,7 @@ getTimestamp = do
   return $ utcTimeToPOSIXSeconds utcTime
 
 
+#ifdef VERSION_bencoding
 -- | Prepare result for 'get_peers' query.
 --
 --   This operation use 'getClosest' as failback so it may block.
@@ -503,6 +506,8 @@ deleteTopic ih p = do
   var <- asks announceInfo
   liftIO $ atomically $ modifyTVar' var (S.delete (ih, p))
 
+#endif
+
 {-----------------------------------------------------------------------
 -- Messaging
 -----------------------------------------------------------------------}
diff --git a/src/Network/BitTorrent/DHT/Token.hs b/src/Network/BitTorrent/DHT/Token.hs
index 3f71aabe..4c930cbc 100644
--- a/src/Network/BitTorrent/DHT/Token.hs
+++ b/src/Network/BitTorrent/DHT/Token.hs
@@ -17,7 +17,7 @@
 -- must be accepted for a reasonable amount of time after they have
 -- been distributed.
 --
-{-# LANGUAGE GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE GeneralizedNewtypeDeriving, CPP #-}
 module Network.BitTorrent.DHT.Token
        ( -- * Token
          Token
@@ -38,7 +38,9 @@ module Network.BitTorrent.DHT.Token
        ) where
 
 import Control.Monad.State
+#ifdef VERSION_bencoding
 import Data.BEncode (BEncode)
+#endif
 import Data.ByteString as BS
 import Data.ByteString.Char8 as B8
 import Data.ByteString.Lazy as BL
@@ -57,7 +59,11 @@ import Network.BitTorrent.Address
 
 -- | An opaque value.
 newtype Token = Token BS.ByteString
-  deriving (Eq, BEncode, IsString)
+  deriving ( Eq, IsString
+#ifdef VERSION_bencoding
+           , BEncode
+#endif
+           )
 
 instance Show Token where
     show (Token bs) = B8.unpack $ Base16.encode bs