1 files changed, 98 insertions, 2 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