2 files changed, 23 insertions, 13 deletions

diff --git a/src/Network/BitTorrent/DHT.hs b/src/Network/BitTorrent/DHT.hs
index 41a76886..5cf7468e 100644
--- a/src/Network/BitTorrent/DHT.hs
+++ b/src/Network/BitTorrent/DHT.hs
@@ -35,7 +35,6 @@ import Data.List as L
 import Data.Monoid
 import Data.Text as T
 import Network.Socket (PortNumber)
-import System.Timeout.Lifted
 import Text.PrettyPrint as PP hiding ((<>))
 import Text.PrettyPrint.Class
 
@@ -103,16 +102,12 @@ bootstrap startNodes = do
   where
     insertClosest addr = do
       nid <- getNodeId
-      result <- try $ timeout 1000000 $ FindNode nid <@> addr
+      result <- try $ FindNode nid <@> addr
       case result of
         Left                    e -> do
           $(logWarnS) "bootstrap" $ T.pack $ show (e :: IOError)
 
-        Right Nothing -> do
-          $(logWarnS) "bootstrap" $ "not responding @ "
-                                 <> T.pack (show (pretty  addr))
-
-        Right (Just (NodeFound closest)) -> do
+        Right (NodeFound closest) -> do
           $(logDebug) ("Get a list of closest nodes: " <>
                        T.pack (PP.render (pretty closest)))
           forM_ (L.take 2 closest) $ \ info @ NodeInfo {..} -> do
diff --git a/src/Network/BitTorrent/DHT/Session.hs b/src/Network/BitTorrent/DHT/Session.hs
index ac285bdc..a76880d7 100644
--- a/src/Network/BitTorrent/DHT/Session.hs
+++ b/src/Network/BitTorrent/DHT/Session.hs
@@ -48,6 +48,7 @@ import Control.Monad.Reader
 import Control.Monad.Trans.Control
 import Control.Monad.Trans.Resource
 import Data.Default
+import Data.Fixed
 import Data.Hashable
 import Data.List as L
 import Data.Monoid
@@ -57,6 +58,7 @@ import Data.Time
 import Data.Time.Clock.POSIX
 import System.Log.FastLogger
 import System.Random (randomIO)
+import System.Timeout.Lifted
 import Text.PrettyPrint as PP hiding ((<>))
 import Text.PrettyPrint.Class
 
@@ -101,6 +103,11 @@ instance Default Options where
     , optTimeout = 5 -- seconds
     }
 
+milliseconds :: NominalDiffTime -> Int
+milliseconds dt = fromEnum millis
+  where
+    millis = realToFrac dt :: Milli
+
 {-----------------------------------------------------------------------
 -- Tokens policy
 -----------------------------------------------------------------------}
@@ -185,11 +192,9 @@ runDHT handlers opts naddr action = runResourceT $ do
 --  Routing
 -----------------------------------------------------------------------}
 
---  TODO fork?
 routing :: Address ip => Routing ip a -> DHT ip (Maybe a)
 routing = runRouting ping refreshNodes getTimestamp
 
--- TODO add timeout
 ping :: Address ip => NodeAddr ip -> DHT ip Bool
 ping addr = do
   $(logDebugS) "routing.questionable_node" (T.pack (render (pretty addr)))
@@ -250,6 +255,7 @@ getTable = do
   var <- asks routingTable
   liftIO (readMVar var)
 
+-- FIXME no blocking
 getNodeId :: DHT ip NodeId
 getNodeId = thisId <$> getTable
 
@@ -305,6 +311,7 @@ getPeerList ih = do
 -- Messaging
 -----------------------------------------------------------------------}
 
+-- | Throws exception if node is not responding.
 (<@>) :: forall a b ip. Address ip => KRPC (Query a) (Response b)
       => a -> NodeAddr ip -> DHT ip b
 q <@> addr = do
@@ -313,11 +320,19 @@ q <@> addr = do
   let Method name = method :: Method (Query a) (Response b)
   let signature = T.decodeUtf8 name <> " @ " <> T.pack (render (pretty addr))
   $(logDebugS) "queryNode" $ "Query sent | " <> signature
-  Response remoteId r <- query (toSockAddr addr) (Query nid q)
-  $(logDebugS) "queryNode" $ "Query recv | " <> signature
 
-  insertNode (NodeInfo remoteId addr)
-  return r
+  interval <- asks (milliseconds . optTimeout . options)
+  result   <- timeout interval  $ query (toSockAddr addr) (Query nid q)
+  case result of
+    Nothing -> do
+      $(logWarnS) "queryNode" $ "not responding @ "
+                   <> T.pack (show (pretty  addr))
+      throwIO $ KError GenericError "timeout expired" ""
+
+    Just (Response remoteId r) -> do
+      $(logDebugS) "queryNode" $ "Query recv | " <> signature
+      insertNode (NodeInfo remoteId addr)
+      return r
 
 type NodeHandler ip = Handler (DHT ip)