1 files changed, 34 insertions, 8 deletions

diff --git a/src/Network/BitTorrent/DHT/Session.hs b/src/Network/BitTorrent/DHT/Session.hs
index aa6ee396..bec2dabc 100644
--- a/src/Network/BitTorrent/DHT/Session.hs
+++ b/src/Network/BitTorrent/DHT/Session.hs
@@ -106,9 +106,11 @@ import Data.Serialize as S
 import Data.Torrent as Torrent
 import Network.KRPC as KRPC hiding (Options, def)
 import qualified Network.KRPC as KRPC (def)
-import Network.KRPC.Message (KMessageOf)
 #ifdef VERSION_bencoding
 import Data.BEncode (BValue)
+import Network.KRPC.Message (KMessageOf)
+#else
+import Data.Tox as Tox
 #endif
 import Network.BitTorrent.Address
 import Network.BitTorrent.DHT.ContactInfo (PeerStore)
@@ -257,11 +259,19 @@ data Node ip = Node
 
     -- | Pseudo-unique self-assigned session identifier. This value is
     -- constant during DHT session and (optionally) between sessions.
-  , tentativeNodeId    :: !NodeId
+#ifdef VERSION_bencoding
+  , tentativeNodeId    :: !(NodeId KMessageOf)
+#else
+  , tentativeNodeId    :: !(NodeId Tox.Message)
+#endif
 
   , resources     :: !InternalState
   , manager       :: !(Manager (DHT       ip     )) -- ^ RPC manager;
-  , routingInfo   :: !(TVar    (Maybe (R.Info ip))) -- ^ search table;
+#ifdef VERSION_bencoding
+  , routingInfo   :: !(TVar    (Maybe (R.Info KMessageOf ip ()))) -- ^ search table;
+#else
+  , routingInfo   :: !(TVar    (Maybe (R.Info Tox.Message ip Bool))) -- ^ search table;
+#endif
   , contactInfo   :: !(TVar    (PeerStore ip     )) -- ^ published by other nodes;
   , announceInfo  :: !(TVar     AnnounceSet       ) -- ^ to publish by this node;
   , sessionTokens :: !(TVar     SessionTokens     ) -- ^ query session IDs.
@@ -319,7 +329,7 @@ instance MonadLogger (DHT ip) where
 #ifdef VERSION_bencoding
 type NodeHandler ip = Handler (DHT ip) KMessageOf BValue
 #else
-type NodeHandler ip = Handler (DHT ip) KMessageOf ByteString
+type NodeHandler ip = Handler (DHT ip) Tox.Message ByteString
 #endif
 
 -- | Run DHT session. You /must/ properly close session using
@@ -330,7 +340,11 @@ newNode :: Address ip
           -> Options          -- ^ various dht options;
           -> NodeAddr ip      -- ^ node address to bind;
           -> LogFun           -- ^
-          -> Maybe NodeId     -- ^ use this NodeId, if not given a new one is generated.
+#ifdef VERSION_bencoding
+          -> Maybe (NodeId KMessageOf)     -- ^ use this NodeId, if not given a new one is generated.
+#else
+          -> Maybe (NodeId Tox.Message)     -- ^ use this NodeId, if not given a new one is generated.
+#endif
           -> IO (Node ip)     -- ^ a new DHT node running at given address.
 newNode hs opts naddr logger mbid = do
     s <- createInternalState
@@ -406,7 +420,11 @@ routableAddress = do
     return $ myAddress <$> info
 
 -- | The current NodeId that the given remote node should know us by.
-myNodeIdAccordingTo :: NodeAddr ip -> DHT ip NodeId
+#ifdef VERSION_bencoding
+myNodeIdAccordingTo :: NodeAddr ip -> DHT ip (NodeId KMessageOf)
+#else
+myNodeIdAccordingTo :: NodeAddr ip -> DHT ip (NodeId Tox.Message)
+#endif
 myNodeIdAccordingTo _ = do
   info <- asks routingInfo >>= liftIO . atomically . readTVar
   maybe (asks tentativeNodeId)
@@ -415,7 +433,11 @@ myNodeIdAccordingTo _ = do
 
 -- | Get current routing table. Normally you don't need to use this
 -- function, but it can be usefull for debugging and profiling purposes.
-getTable :: Eq ip => DHT ip (Table ip)
+#ifdef VERSION_bencoding
+getTable :: Eq ip => DHT ip (Table KMessageOf ip ())
+#else
+getTable :: Eq ip => DHT ip (Table Tox.Message ip Bool)
+#endif
 getTable = do
   Node { tentativeNodeId = myId
        , routingInfo = var
@@ -452,7 +474,11 @@ allPeers ih = do
 --
 --   This operation used for 'find_nodes' query.
 --
-getClosest :: Eq ip => TableKey k => k -> DHT ip [NodeInfo KMessageOf ip ()]
+#ifdef VERSION_bencoding
+getClosest :: Eq ip => TableKey KMessageOf k => k -> DHT ip [NodeInfo KMessageOf ip ()]
+#else
+getClosest :: Eq ip => TableKey Tox.Message k => k -> DHT ip [NodeInfo Tox.Message ip Bool]
+#endif
 getClosest node = do
   k <- asks (optK . options)
   kclosest k node <$> getTable