1 files changed, 39 insertions, 34 deletions

diff --git a/src/Network/BitTorrent/DHT/Query.hs b/src/Network/BitTorrent/DHT/Query.hs
index 68c67900..254b347c 100644
--- a/src/Network/BitTorrent/DHT/Query.hs
+++ b/src/Network/BitTorrent/DHT/Query.hs
@@ -80,6 +80,7 @@ import Text.PrettyPrint as PP hiding ((<>), ($$))
 import Text.PrettyPrint.HughesPJClass hiding ((<>),($$))
 import Data.Time
 import Data.Time.Clock.POSIX
+import Data.Hashable (Hashable)
 
 import Network.DatagramServer as KRPC hiding (Options, def)
 import Network.KRPC.Method    as KRPC
@@ -109,13 +110,9 @@ import Control.Monad.Trans.Control
 
 nodeHandler :: ( Address ip
                , KRPC (Query a) (Response b)
-#ifdef VERSION_bencoding
-               , KRPC.Envelope (Query a) (Response b) ~ BValue )
-#else
-               , KPRC.Envelope (Query a) (Response b) ~ ByteString )
-#endif
-           => QueryMethod KMessageOf -> (NodeAddr ip -> a -> DHT ip b) -> NodeHandler ip
-nodeHandler method action = handler method $ \ sockAddr qry -> do
+               )
+           => (NodeInfo KMessageOf ip () -> Maybe ReflectedIP -> IO ()) -> (NodeAddr ip -> IO (NodeId KMessageOf)) -> (Char -> String -> Text -> IO ()) -> QueryMethod KMessageOf -> (NodeAddr ip -> a -> IO b) -> NodeHandler
+nodeHandler insertNode myNodeIdAccordingTo logm method action = handler method $ \ sockAddr qry -> do
 #ifdef VERSION_bencoding
   let remoteId  = queringNodeId qry
       read_only = queryIsReadOnly qry
@@ -131,53 +128,55 @@ nodeHandler method action = handler method $ \ sockAddr qry -> do
       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)
+        then logm 'W' "nodeHandler" $ "READ-ONLY " <> T.pack (show $ pPrint ni)
         else insertNode ni Nothing >> return () -- TODO need to block. why?
       Response
         <$> myNodeIdAccordingTo naddr
         <*> action naddr q
 
 -- | Default 'Ping' handler.
-pingH :: Address ip => NodeHandler ip
-#ifdef VERSION_bencoding
-pingH = nodeHandler "ping" $ \ _ Ping -> return Ping
-#else
-pingH = nodeHandler $ \ _ p@PingPayload{} -> return p { isPong = True }
-#endif
+pingH :: NodeAddr ip -> Ping -> IO Ping
+pingH _ Ping = return Ping
+-- pingH = nodeHandler $ \ _ p@PingPayload{} -> return p { isPong = True }
 
 -- | Default 'FindNode' handler.
-findNodeH :: Address ip => NodeHandler ip
-findNodeH = nodeHandler "find-nodes" $ \ _ (FindNode nid) -> do
-  NodeFound <$> getClosest nid
+findNodeH :: (NodeId KMessageOf -> IO [NodeInfo KMessageOf ip ()]) -> NodeAddr ip -> FindNode ip -> IO (NodeFound ip)
+findNodeH getclosest _ (FindNode nid) = NodeFound <$> getclosest nid
 
-#ifdef VERSION_bencoding
 -- | Default 'GetPeers' handler.
-getPeersH :: Ord ip => Address ip => NodeHandler ip
-getPeersH = nodeHandler "get_peers" $ \ naddr (GetPeers ih) -> do
+getPeersH :: Hashable ip => (InfoHash -> IO (PeerList ip)) -> TVar SessionTokens -> NodeAddr ip -> GetPeers ip -> IO (GotPeers ip)
+getPeersH getPeerList toks naddr (GetPeers ih) = do
   ps <- getPeerList ih
-  tok <- grantToken naddr
+  tok <- grantToken toks naddr
   return $ GotPeers ps tok
 
 -- | Default 'Announce' handler.
-announceH :: Ord ip => Address ip => NodeHandler ip
-announceH = nodeHandler "announce_peer" $ \ naddr @ NodeAddr {..} (Announce {..}) -> do
-  valid <- checkToken naddr sessionToken
+announceH :: ( Ord ip, Hashable ip ) => TVar (PeerStore ip) -> TVar SessionTokens -> NodeAddr ip -> Announce -> IO Announced
+announceH peers toks naddr @ NodeAddr {..} (Announce {..}) = do
+  valid <- checkToken toks naddr sessionToken
   unless valid $ do
     throwIO $ InvalidParameter "token"
 
   let annPort  = if impliedPort then nodePort else port
       peerAddr = PeerAddr Nothing nodeHost annPort
-  insertPeer topic announcedName peerAddr
+  insertPeer peers topic announcedName peerAddr
   return Announced
 
 -- | 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
+defaultHandlers :: forall ip. (Eq ip, Ord ip, Address ip) => LogFun -> DHT ip [NodeHandler]
+defaultHandlers logger = do
+    groknode <- insertNode1
+    toks <- asks sessionTokens
+    getclosest <- getClosest1
+    mynid <- myNodeIdAccordingTo1
+    peers <- asks contactInfo
+    getpeers <- getPeerList1
+    let handler :: KRPC (Query a) (Response b) => QueryMethod KMessageOf -> (NodeAddr ip -> a -> IO b) -> NodeHandler
+        handler = nodeHandler groknode mynid (logt logger)
+    return [ handler "ping"          $ pingH
+           , handler "find-nodes"    $ findNodeH getclosest
+           , handler "get_peers"     $ getPeersH getpeers toks
+           , handler "announce_peer" $ announceH peers toks ]
 
 {-----------------------------------------------------------------------
 --  Basic queries
@@ -324,6 +323,11 @@ logc c   = $(logInfoS)  "insertNode" . T.pack . (c :) . (':' :)
 --   routing table.
 insertNode :: forall ip. Address ip => NodeInfo KMessageOf ip () -> Maybe ReflectedIP -> DHT ip ()
 insertNode info witnessed_ip0 = do
+    f <- insertNode1
+    liftIO $ f info witnessed_ip0
+
+insertNode1 :: forall ip. Address ip => DHT ip (NodeInfo KMessageOf ip () -> Maybe ReflectedIP -> IO ())
+insertNode1 = do
   bc <- optBucketCount <$> asks options
   nid <- asks tentativeNodeId
   logm0 <- embed_ (uncurry logc)
@@ -349,7 +353,7 @@ insertNode info witnessed_ip0 = do
         , grokNode    = DHT.insertNode params state
         , grokAddress = \_ _ -> return () -- :: Maybe SockAddr -> ReflectedIP -> IO ()
         }
-  liftIO $ DHT.insertNode params state info witnessed_ip0
+  return $ \info witnessed_ip0 -> DHT.insertNode params state info witnessed_ip0
 
 -- | Throws exception if node is not responding.
 queryNode :: forall a b ip. Address ip => KRPC (Query a) (Response b)
@@ -362,7 +366,8 @@ queryNode' addr q = do
   nid <- myNodeIdAccordingTo addr
   let read_only = False -- TODO: check for NAT issues. (BEP 43)
   let KRPC.Method name = KRPC.method :: KRPC.Method (Query a) (Response b)
-  (Response remoteId r, witnessed_ip) <- query' name (toSockAddr addr) (Query nid read_only q)
+  mgr <- asks manager
+  (Response remoteId r, witnessed_ip) <- liftIO $ query' mgr name (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