1 files changed, 35 insertions, 28 deletions

diff --git a/src/Network/BitTorrent/DHT/Query.hs b/src/Network/BitTorrent/DHT/Query.hs
index 820db8ba..4b386cdc 100644
--- a/src/Network/BitTorrent/DHT/Query.hs
+++ b/src/Network/BitTorrent/DHT/Query.hs
@@ -80,7 +80,7 @@ import Text.PrettyPrint.HughesPJClass hiding ((<>),($$))
 import Data.Time
 import Data.Time.Clock.POSIX
 
-import Network.KRPC hiding (Options, def)
+import Network.KRPC as KRPC hiding (Options, def)
 import Network.KRPC.Message (ReflectedIP(..))
 import Network.KRPC.Manager (QueryFailure(..))
 import Data.Torrent
@@ -90,14 +90,15 @@ import Network.BitTorrent.DHT.Session
 import Control.Concurrent.STM
 import qualified Network.BitTorrent.DHT.Search as Search
 #ifdef VERSION_bencoding
-import Network.BitTorrent.Address
 import Data.BEncode (BValue)
-import Network.DHT.Mainline
+import Network.DHT.Mainline hiding (NodeId)
+import Network.KRPC.Message (KMessageOf)
 #else
-import Network.BitTorrent.Address hiding (NodeId)
 import Data.ByteString (ByteString)
 import Data.Tox
 #endif
+import Network.BitTorrent.Address hiding (NodeId)
+import Network.RPC as RPC hiding (Query,Response)
 
 {-----------------------------------------------------------------------
 --  Handlers
@@ -106,18 +107,17 @@ import Data.Tox
 nodeHandler :: ( Address ip
                , KRPC (Query a) (Response b)
 #ifdef VERSION_bencoding
-               , Envelope (Query a) (Response b) ~ BValue )
+               , KRPC.Envelope (Query a) (Response b) ~ BValue )
 #else
-               , Envelope (Query a) (Response b) ~ ByteString )
+               , KPRC.Envelope (Query a) (Response b) ~ ByteString )
 #endif
            => (NodeAddr ip -> a -> DHT ip b) -> NodeHandler ip
-#ifdef VERSION_bencoding
 nodeHandler action = handler $ \ sockAddr qry -> do
+#ifdef VERSION_bencoding
   let remoteId  = queringNodeId qry
       read_only = queryIsReadOnly qry
       q         = queryParams qry
 #else
-nodeHandler action = handler (error "TODO TOX Messaging") $ \ sockAddr qry -> do
   let remoteId  = msgClient qry
       read_only = False
       q         = msgPayload qry
@@ -125,7 +125,7 @@ nodeHandler action = handler (error "TODO TOX Messaging") $ \ sockAddr qry -> do
   case fromSockAddr sockAddr of
     Nothing    -> throwIO BadAddress
     Just naddr -> do
-      let ni = NodeInfo remoteId naddr
+      let ni = NodeInfo remoteId naddr ()
       -- Do not route read-only nodes. (bep 43)
       if read_only
         then $(logWarnS) "nodeHandler" $ "READ-ONLY " <> T.pack (show $ pPrint ni)
@@ -136,8 +136,11 @@ nodeHandler action = handler (error "TODO TOX Messaging") $ \ sockAddr qry -> do
 
 -- | Default 'Ping' handler.
 pingH :: Address ip => NodeHandler ip
-pingH = nodeHandler $ \ _ Ping -> do
-  return Ping
+#ifdef VERSION_bencoding
+pingH = nodeHandler $ \ _ Ping -> return Ping
+#else
+pingH = nodeHandler $ \ _ p@PingPayload{} -> return p { isPong = True }
+#endif
 
 -- | Default 'FindNode' handler.
 findNodeH :: Address ip => NodeHandler ip
@@ -177,19 +180,23 @@ defaultHandlers = [pingH, findNodeH]
 --  Basic queries
 -----------------------------------------------------------------------}
 
-type Iteration ip o = NodeInfo ip -> DHT ip (Either [NodeInfo ip] [o ip])
+type Iteration ip o = NodeInfo KMessageOf ip () -> DHT ip (Either [NodeInfo KMessageOf ip ()] [o ip])
 
 -- | The most basic query. May be used to check if the given node is
 -- alive or get its 'NodeId'.
-pingQ :: Address ip => NodeAddr ip -> DHT ip (NodeInfo ip, Maybe ReflectedIP)
+pingQ :: Address ip => NodeAddr ip -> DHT ip (NodeInfo KMessageOf ip (), Maybe ReflectedIP)
 pingQ addr = do
+#ifdef VERSION_bencoding
   (nid, Ping, mip) <- queryNode' addr Ping
-  return (NodeInfo nid addr, mip)
+#else
+  (nid, PingPayload{}, mip) <- queryNode' addr PingPayload {isPong=False, pingId=pid}
+#endif
+  return (NodeInfo nid addr (), mip)
 
 -- TODO [robustness] match range of returned node ids with the
 -- expected range and either filter bad nodes or discard response at
 -- all throwing an exception
-findNodeQ :: Address ip => TableKey key => key -> Iteration ip NodeInfo
+-- findNodeQ :: Address ip => TableKey key => key -> IterationI ip NodeInfo
 findNodeQ key NodeInfo {..} = do
   NodeFound closest <- FindNode (toNodeId key) <@> nodeAddr
   $(logInfoS) "findNodeQ" $ "NodeFound\n"
@@ -223,7 +230,7 @@ announceQ ih p NodeInfo {..} = do
 -----------------------------------------------------------------------}
 
 
-ioGetPeers :: Address ip => InfoHash -> DHT ip (NodeInfo ip -> IO ([NodeInfo ip], [PeerAddr ip]))
+ioGetPeers :: Address ip => InfoHash -> DHT ip (NodeInfo KMessageOf ip () -> IO ([NodeInfo KMessageOf ip ()], [PeerAddr ip]))
 ioGetPeers ih = do
     session <- ask
     return $ \ni -> runDHT session $ do
@@ -232,7 +239,7 @@ ioGetPeers ih = do
             Right e -> return $ either (,[]) ([],) e
             Left e -> let _ = e :: QueryFailure in return ([],[])
 
-ioFindNode :: Address ip => InfoHash -> DHT ip (NodeInfo ip -> IO ([NodeInfo ip], [NodeInfo ip]))
+ioFindNode :: Address ip => InfoHash -> DHT ip (NodeInfo KMessageOf ip () -> IO ([NodeInfo KMessageOf ip ()], [NodeInfo KMessageOf ip ()]))
 ioFindNode ih = do
     session <- ask
     return $ \ni -> runDHT session $ do
@@ -240,7 +247,7 @@ ioFindNode ih = do
         return $ L.partition (\n -> nodeId n /= toNodeId ih) ns
 
 isearch :: (Ord r, Ord ip) =>
-            (InfoHash -> DHT ip (NodeInfo ip -> IO ([NodeInfo ip], [r])))
+            (InfoHash -> DHT ip (NodeInfo KMessageOf ip () -> IO ([NodeInfo KMessageOf ip ()], [r])))
             -> InfoHash
             -> DHT ip (ThreadId, Search.IterativeSearch ip r)
 isearch f ih = do
@@ -255,10 +262,10 @@ isearch f ih = do
                 return (a, s)
 
 
-type Search    ip o = Conduit [NodeInfo ip] (DHT ip) [o ip]
+type Search    ip o = Conduit [NodeInfo KMessageOf ip ()] (DHT ip) [o KMessageOf ip ()]
 
 -- TODO: use reorder and filter (Traversal option) leftovers
-search :: k -> Iteration ip o -> Search ip o
+-- search :: k -> IterationI ip o -> Search ip o
 search _ action = do
   awaitForever $ \ batch -> unless (L.null batch) $ do
     $(logWarnS) "search" "start query"
@@ -285,15 +292,15 @@ probeNode addr = do
 
 
 -- FIXME do not use getClosest sinse we should /refresh/ them
-refreshNodes :: Address ip => NodeId -> DHT ip () -- [NodeInfo ip]
+refreshNodes :: Address ip => NodeId KMessageOf -> DHT ip () -- [NodeInfo KMessageOf ip ()]
 refreshNodes nid = do
   $(logDebugS) "routing.refresh_bucket" (T.pack (render (pPrint nid)))
   nodes <- getClosest nid
   do
     -- forM (L.take 1 nodes) $ \ addr -> do
     -- NodeFound ns <- FindNode nid <@> addr
-    -- Expected type: ConduitM [NodeAddr ip] [NodeInfo ip] (DHT ip) ()
-    --   Actual type: ConduitM [NodeInfo ip] [NodeInfo ip] (DHT ip) ()
+    -- Expected type: ConduitM [NodeAddr ip] [NodeInfo KMessageOf ip ()] (DHT ip) ()
+    --   Actual type: ConduitM [NodeInfo KMessageOf ip ()] [NodeInfo KMessageOf ip ()] (DHT ip) ()
     -- nss <- sourceList [[addr]] $= search nid (findNodeQ nid) $$ C.consume
     nss <- sourceList [nodes] $= search nid (findNodeQ nid) $$ C.consume
     $(logWarnS) "refreshNodes" $ "received " <> T.pack (show (L.length (L.concat nss))) <> " nodes."
@@ -306,7 +313,7 @@ refreshNodes nid = do
 
 -- | This operation do not block but acquire exclusive access to
 --   routing table.
-insertNode :: forall ip. Address ip => NodeInfo ip -> Maybe ReflectedIP -> DHT ip ()
+insertNode :: forall ip. Address ip => NodeInfo KMessageOf ip () -> Maybe ReflectedIP -> DHT ip ()
 insertNode info witnessed_ip0 = do
   var <- asks routingInfo
   tm <- getTimestamp
@@ -315,7 +322,7 @@ insertNode info witnessed_ip0 = do
           let logMsg = "Routing table: " <> pPrint t
           $(logDebugS) "insertNode" (T.pack (render logMsg))
   let arrival0 = TryInsert info
-      arrival4 = TryInsert (fmap fromAddr info) :: Event (Maybe IPv4)
+      arrival4 = TryInsert (mapAddress fromAddr info) :: Event (Maybe IPv4)
   $(logDebugS) "insertNode" $ T.pack (show arrival4)
   maxbuckets <- asks (optBucketCount . options)
   fallbackid <- asks tentativeNodeId
@@ -380,18 +387,18 @@ insertNode info witnessed_ip0 = do
 
 -- | Throws exception if node is not responding.
 queryNode :: forall a b ip. Address ip => KRPC (Query a) (Response b)
-          => NodeAddr ip -> a -> DHT ip (NodeId, b)
+          => NodeAddr ip -> a -> DHT ip (NodeId KMessageOf, b)
 queryNode addr q = fmap (\(n,b,_) -> (n,b)) $ queryNode' addr q
 
 queryNode' :: forall a b ip. Address ip => KRPC (Query a) (Response b)
-          => NodeAddr ip -> a -> DHT ip (NodeId, b, Maybe ReflectedIP)
+          => NodeAddr ip -> a -> DHT ip (NodeId KMessageOf, b, Maybe ReflectedIP)
 queryNode' addr q = do
   nid <- myNodeIdAccordingTo addr
   let read_only = False -- TODO: check for NAT issues. (BEP 43)
   (Response remoteId r, witnessed_ip) <- query' (toSockAddr addr) (Query nid read_only q)
   -- $(logDebugS) "queryNode" $ "Witnessed address: " <> T.pack (show witnessed_ip)
   --                             <> " by " <> T.pack (show (toSockAddr addr))
-  _ <- insertNode (NodeInfo remoteId addr) witnessed_ip
+  _ <- insertNode (NodeInfo remoteId addr ()) witnessed_ip
   return (remoteId, r, witnessed_ip)
 
 -- | Infix version of 'queryNode' function.