Modify kademlia search to distinguish a Canceled from timed-out query.

1 files changed, 8 insertions, 5 deletions

diff --git a/dht/TCPProber.hs b/dht/TCPProber.hs
index faf8b35c..17b68f64 100644
--- a/dht/TCPProber.hs
+++ b/dht/TCPProber.hs
@@ -26,6 +26,7 @@ import Data.Wrapper.PSQ as PSQ
 import Network.Kademlia.Search
 import Network.Tox.NodeId
 import qualified Network.Tox.TCP as TCP
+import Network.QueryResponse as QR
 
 -- Probe TCP ports in a staggered fashion to up the odds of discovering
 -- a higher priority port like 443.
@@ -156,7 +157,7 @@ runProbeQueue prober client maxjobs = do
         loop
 
 
-getNodes :: TCPProber -> TCP.TCPClient err Nonce8 -> NodeId -> TCP.NodeInfo -> IO (Maybe ([TCP.NodeInfo],[TCP.NodeInfo],Maybe ()))
+getNodes :: TCPProber -> TCP.TCPClient err Nonce8 -> NodeId -> TCP.NodeInfo -> IO (Result ([TCP.NodeInfo],[TCP.NodeInfo],Maybe ()))
 getNodes prober tcp seeking dst = do
     r <- TCP.getUDPNodes' tcp seeking (TCP.udpNodeInfo dst)
     dput XTCP $ "Got via TCP nodes: " ++ show r
@@ -164,14 +165,16 @@ getNodes prober tcp seeking dst = do
          where ns' = do
                     n <- ns
                     [ TCP.NodeInfo n 0 ]
-    fmap join $ forM r $ \(ns,gw) -> do
+    case r of
+      Success (ns,gw) -> do
         let ts = tcps ns
         if TCP.nodeId gw == TCP.nodeId dst
-            then return $ Just ts
+            then return $ Success ts
             else do
                 enqueueProbe prober (TCP.udpNodeInfo dst)
-                return $ Just ts
-        return $ Just ts
+                return $ Success ts
+        return $ Success ts
+      _ -> return $ fmap (const $ error "TCPProber.getNodes: The impossible happened!") r
 
 nodeSearch :: TCPProber -> TCP.TCPClient err Nonce8 -> Search NodeId (IP, PortNumber) () TCP.NodeInfo TCP.NodeInfo
 nodeSearch prober tcp = Search