3 files changed, 107 insertions, 85 deletions

diff --git a/src/Network/BitTorrent/DHT/Query.hs b/src/Network/BitTorrent/DHT/Query.hs
index d51ab505..c5fcccb4 100644
--- a/src/Network/BitTorrent/DHT/Query.hs
+++ b/src/Network/BitTorrent/DHT/Query.hs
@@ -9,8 +9,9 @@
 --   Normally, you don't need to import this module, use
 --   "Network.BitTorrent.DHT" instead.
 --
-{-# LANGUAGE FlexibleContexts  #-}
-{-# LANGUAGE TemplateHaskell   #-}
+{-# LANGUAGE FlexibleContexts    #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TemplateHaskell     #-}
 module Network.BitTorrent.DHT.Query
        ( -- * Handler
          -- | To bind specific set of handlers you need to pass
@@ -38,6 +39,14 @@ module Network.BitTorrent.DHT.Query
        , Search
        , search
        , publish
+
+       -- ** Routing table
+       , insertNode
+       , refreshNodes
+
+       -- ** Messaging
+       , queryNode
+       , (<@>)
        ) where
 
 import Control.Applicative
@@ -54,12 +63,14 @@ import Data.Text as T
 import Network
 import Text.PrettyPrint as PP hiding ((<>), ($$))
 import Text.PrettyPrint.HughesPJClass hiding ((<>),($$))
+import Data.Time
+import Data.Time.Clock.POSIX
 
 import Network.KRPC hiding (Options, def)
 import Data.Torrent
 import Network.BitTorrent.Address
 import Network.BitTorrent.DHT.Message
-import Network.BitTorrent.DHT.Routing
+import Network.BitTorrent.DHT.Routing as R
 import Network.BitTorrent.DHT.Session
 
 {-----------------------------------------------------------------------
@@ -165,7 +176,7 @@ search _ action = do
     $(logWarnS) "search" "start query"
     responses <- lift $ queryParallel (action <$> batch)
     let (nodes, results) = partitionEithers responses
-    $(logWarnS) "search" ("done query more:" <> T.pack (show $ L.length nodes))
+    $(logWarnS) "search" ("done query more:" <> T.pack (show (L.length nodes, L.length results)))
     leftover $ L.concat nodes
     mapM_ yield results
 
@@ -180,3 +191,78 @@ republish :: DHT ip ThreadId
 republish = fork $ do
   i <- asks (optReannounce . options)
   error "DHT.republish: not implemented"
+
+routing :: Address ip => Routing ip a -> DHT ip (Maybe a)
+routing = runRouting probeNode refreshNodes getTimestamp
+
+getTimestamp :: DHT ip Timestamp
+getTimestamp = do
+  utcTime <- liftIO $ getCurrentTime
+  $(logDebugS) "routing.make_timestamp" (T.pack (render (pPrint utcTime)))
+  return $ utcTimeToPOSIXSeconds utcTime
+
+
+probeNode :: Address ip => NodeAddr ip -> DHT ip Bool
+probeNode addr = do
+  $(logDebugS) "routing.questionable_node" (T.pack (render (pPrint addr)))
+  result <- try $ Ping <@> addr
+  let _ = result :: Either SomeException Ping
+  return $ either (const False) (const True) result
+
+
+-- FIXME do not use getClosest sinse we should /refresh/ them
+refreshNodes :: Address ip => NodeId -> DHT ip () -- [NodeInfo 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) ()
+    -- 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."
+    queryParallel $ flip L.map (L.concat nss) $ \n -> do
+        $(logWarnS) "refreshNodes" $ "received node: " <> T.pack (show (pPrint n))
+        pingQ (nodeAddr n)
+        insertNode n
+    return ()
+  return () -- $ L.concat nss
+
+-- | This operation do not block but acquire exclusive access to
+--   routing table.
+insertNode :: Address ip => NodeInfo ip -> DHT ip ThreadId
+insertNode info = fork $ do
+  var <- asks routingTable
+  t <- takeMVar var
+  t' <- do -- modifyMVar_ var $ \ t -> do
+    result <- routing (R.insert info t)
+    case result of
+      Nothing -> do
+        $(logDebugS) "insertNode" $ "Routing table is full: "
+                   <> T.pack (show (pPrint t))
+        return t
+      Just t' -> do
+        let logMsg = "Routing table updated: "
+                  <> pPrint t <> " -> " <> pPrint t'
+        $(logDebugS) "insertNode" (T.pack (render logMsg))
+        return t'
+  putMVar var t'
+
+-- | 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)
+queryNode addr q = do
+  nid <- asks thisNodeId
+  let read_only = False -- TODO: check for NAT issues. (BEP 43)
+  Response remoteId r <- query (toSockAddr addr) (Query nid read_only q)
+  insertNode (NodeInfo remoteId addr)
+  return (remoteId, r)
+
+-- | Infix version of 'queryNode' function.
+(<@>) :: Address ip => KRPC (Query a) (Response b)
+      => a -> NodeAddr ip -> DHT ip b
+q <@> addr = snd <$> queryNode addr q
+{-# INLINE (<@>) #-}
+
diff --git a/src/Network/BitTorrent/DHT/Routing.hs b/src/Network/BitTorrent/DHT/Routing.hs
index 8705a5a2..fee52380 100644
--- a/src/Network/BitTorrent/DHT/Routing.hs
+++ b/src/Network/BitTorrent/DHT/Routing.hs
@@ -103,14 +103,14 @@ data Routing ip result
   | Done      result
   | GetTime                ( Timestamp    -> Routing ip result)
   | NeedPing (NodeAddr ip) ( Bool         -> Routing ip result)
-  | Refresh   NodeId       ([NodeInfo ip] -> Routing ip result)
+  | Refresh   NodeId       (Routing ip result)
 
 instance Functor (Routing ip) where
   fmap _  Full             = Full
   fmap f (Done          r) = Done          (     f   r)
   fmap f (GetTime       g) = GetTime       (fmap f . g)
   fmap f (NeedPing addr g) = NeedPing addr (fmap f . g)
-  fmap f (Refresh  nid  g) = Refresh  nid  (fmap f . g)
+  fmap f (Refresh  nid  g) = Refresh  nid  (fmap f g)
 
 instance Monad (Routing ip) where
   return = Done
@@ -119,7 +119,7 @@ instance Monad (Routing ip) where
   Done       r >>= m = m r
   GetTime    f >>= m = GetTime     $ \ t -> f t >>= m
   NeedPing a f >>= m = NeedPing a  $ \ p -> f p >>= m
-  Refresh  n f >>= m = Refresh  n  $ \ i -> f i >>= m
+  Refresh  n f >>= m = Refresh  n  $ f >>= m
 
 instance Applicative (Routing ip) where
   pure  = return
@@ -132,15 +132,15 @@ instance Alternative (Routing ip) where
   Done     a   <|> _ = Done a
   GetTime  f   <|> m = GetTime    $ \ t -> f t <|> m
   NeedPing a f <|> m = NeedPing a $ \ p -> f p <|> m
-  Refresh  n f <|> m = Refresh  n $ \ i -> f i <|> m
+  Refresh  n f <|> m = Refresh  n (f <|> m)
 
 -- | Run routing table operation.
 runRouting :: (Monad m, Eq ip)
-           => (NodeAddr ip -> m Bool)          -- ^ ping the specific node;
-           -> (NodeId      -> m [NodeInfo ip]) -- ^ get closest nodes;
-           -> m Timestamp                      -- ^ get current time;
-           -> Routing ip f                     -- ^ operation to run;
-           -> m (Maybe f)                      -- ^ operation result;
+           => (NodeAddr ip -> m Bool) -- ^ ping the specific node;
+           -> (NodeId      -> m ())   -- ^ refresh nodes;
+           -> m Timestamp             -- ^ get current time;
+           -> Routing ip f            -- ^ operation to run;
+           -> m (Maybe f)             -- ^ operation result;
 runRouting ping_node find_nodes timestamper = go
   where
     go  Full             = return (Nothing)
@@ -154,8 +154,8 @@ runRouting ping_node find_nodes timestamper = go
       go (f pong)
 
     go (Refresh nid f) = do
-      infos <- find_nodes nid
-      go (f infos)
+      find_nodes nid
+      go f
 
 getTime :: Routing ip Timestamp
 getTime = GetTime return
@@ -165,8 +165,8 @@ needPing :: NodeAddr ip -> Routing ip Bool
 needPing addr = NeedPing addr return
 {-# INLINE needPing #-}
 
-refresh :: NodeId -> Routing ip [NodeInfo ip]
-refresh nid = Refresh nid return
+refresh :: NodeId -> Routing ip ()
+refresh nid = Refresh nid (Done ())
 {-# INLINE refresh #-}
 
 {-----------------------------------------------------------------------
@@ -238,10 +238,8 @@ insertBucket curTime info bucket
   --  Buckets that have not been changed in 15 minutes should be "refreshed."
   | Just (NodeInfo {..} :-> lastSeen) <- lastChanged bucket
   , curTime - lastSeen > delta = do
-    infos   <- refresh nodeId
-    refTime <- getTime
-    let newBucket = L.foldr (\ x -> PSQ.insertWith max x refTime) bucket infos
-    insertBucket refTime info newBucket
+    refresh nodeId
+    insertBucket curTime info bucket
 
   -- If there are any questionable nodes in the bucket have not been
   -- seen in the last 15 minutes, the least recently seen node is
@@ -260,7 +258,7 @@ insertBucket curTime info bucket
     return $ PSQ.insert info curTime bucket
 
   -- When the bucket is full of good nodes, the new node is simply discarded.
-  | otherwise = A.empty
+  | otherwise = Full
 
 insertNode :: Eq ip => NodeInfo ip -> Bucket ip -> ip `Routing` Bucket ip
 insertNode info bucket = do
diff --git a/src/Network/BitTorrent/DHT/Session.hs b/src/Network/BitTorrent/DHT/Session.hs
index ffce47de..2bb3ce85 100644
--- a/src/Network/BitTorrent/DHT/Session.hs
+++ b/src/Network/BitTorrent/DHT/Session.hs
@@ -30,6 +30,7 @@ module Network.BitTorrent.DHT.Session
        , Node
        , options
        , thisNodeId
+       , routingTable
 
          -- ** Initialization
        , LogFun
@@ -50,7 +51,6 @@ module Network.BitTorrent.DHT.Session
          -- ** Routing table
        , getTable
        , getClosest
-       , insertNode
 
          -- ** Peer storage
        , insertPeer
@@ -59,9 +59,7 @@ module Network.BitTorrent.DHT.Session
        , deleteTopic
 
          -- ** Messaging
-       , queryNode
        , queryParallel
-       , (<@>)
        ) where
 
 import Prelude hiding (ioError)
@@ -347,35 +345,9 @@ runDHT node action = runReaderT (unDHT action) node
 --  Routing
 -----------------------------------------------------------------------}
 
-routing :: Address ip => Routing ip a -> DHT ip (Maybe a)
-routing = runRouting probeNode refreshNodes getTimestamp
-
-probeNode :: Address ip => NodeAddr ip -> DHT ip Bool
-probeNode addr = do
-  $(logDebugS) "routing.questionable_node" (T.pack (render (pPrint addr)))
-  result <- try $ Ping <@> addr
-  let _ = result :: Either SomeException Ping
-  return $ either (const False) (const True) result
-
 -- /pick a random ID/ in the range of the bucket and perform a
 -- find_nodes search on it.
 
--- FIXME do not use getClosest sinse we should /refresh/ them
-refreshNodes :: Address ip => NodeId -> DHT ip [NodeInfo ip]
-refreshNodes nid = do
-  $(logDebugS) "routing.refresh_bucket" (T.pack (render (pPrint nid)))
-  nodes <- getClosest nid
-  nss <- forM (nodeAddr <$> nodes) $ \ addr -> do
-    NodeFound ns <- FindNode nid <@> addr
-    return ns
-  return $ L.concat nss
-
-getTimestamp :: DHT ip Timestamp
-getTimestamp = do
-  utcTime <- liftIO $ getCurrentTime
-  $(logDebugS) "routing.make_timestamp" (T.pack (render (pPrint utcTime)))
-  return $ utcTimeToPOSIXSeconds utcTime
-
 {-----------------------------------------------------------------------
 -- Tokens
 -----------------------------------------------------------------------}
@@ -421,24 +393,6 @@ getClosest node = do
   k <- asks (optK . options)
   kclosest k node <$> getTable
 
--- | This operation do not block but acquire exclusive access to
---   routing table.
-insertNode :: Address ip => NodeInfo ip -> DHT ip ThreadId
-insertNode info = fork $ do
-  var <- asks routingTable
-  modifyMVar_ var $ \ t -> do
-    result <- routing (R.insert info t)
-    case result of
-      Nothing -> do
-        $(logDebugS) "insertNode" $ "Routing table is full: "
-                   <> T.pack (show (pPrint t))
-        return t
-      Just t' -> do
-        let logMsg = "Routing table updated: "
-                  <> pPrint t <> " -> " <> pPrint t'
-        $(logDebugS) "insertNode" (T.pack (render logMsg))
-        return t'
-
 {-----------------------------------------------------------------------
 -- Peer storage
 -----------------------------------------------------------------------}
@@ -487,22 +441,6 @@ deleteTopic ih p = do
 -- Messaging
 -----------------------------------------------------------------------}
 
--- | 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)
-queryNode addr q = do
-  nid <- asks thisNodeId
-  let read_only = False -- TODO: check for NAT issues. (BEP 43)
-  Response remoteId r <- query (toSockAddr addr) (Query nid read_only q)
-  insertNode (NodeInfo remoteId addr)
-  return (remoteId, r)
-
--- | Infix version of 'queryNode' function.
-(<@>) :: Address ip => KRPC (Query a) (Response b)
-      => a -> NodeAddr ip -> DHT ip b
-q <@> addr = snd <$> queryNode addr q
-{-# INLINE (<@>) #-}
-
 -- TODO: use alpha
 -- | Failed queries are ignored.
 queryParallel :: [DHT ip a] -> DHT ip [a]