1 files changed, 26 insertions, 2 deletions

diff --git a/dht/TCPProber.hs b/dht/TCPProber.hs
index 17b68f64..ccdbd8d1 100644
--- a/dht/TCPProber.hs
+++ b/dht/TCPProber.hs
@@ -176,11 +176,35 @@ getNodes prober tcp seeking dst = do
         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
+asyncGetNodes :: TCPProber -> TCP.TCPClient err Nonce8 -> NodeId -> TCP.NodeInfo
+                           -> (Nonce8 -> Result ([TCP.NodeInfo],[TCP.NodeInfo],Maybe ()) -> IO ())
+                           -> IO Nonce8
+asyncGetNodes prober tcp seeking dst withResponse = do
+    TCP.asyncUDPNodes tcp seeking (TCP.udpNodeInfo dst) $ \qid r -> do
+        dput XTCP $ "Got via TCP nodes: " ++ show r
+        let tcps (ns,_,mb) = (ns',ns',mb)
+             where ns' = do
+                        n <- ns
+                        [ TCP.NodeInfo n 0 ]
+        r' <- case r of
+          Success (ns,gw) -> do
+            let ts = tcps ns
+            if TCP.nodeId gw == TCP.nodeId dst
+                then return $ Success ts
+                else do
+                    enqueueProbe prober (TCP.udpNodeInfo dst)
+                    return $ Success ts
+            return $ Success ts
+          _ -> return $ fmap (const $ error "TCPProber.getNodes: The impossible happened!") r
+        withResponse qid r'
+
+
+nodeSearch :: TCPProber -> TCP.TCPClient err Nonce8 -> Search NodeId (IP, PortNumber) () TCP.NodeInfo TCP.NodeInfo Nonce8
 nodeSearch prober tcp = Search
     { searchSpace       = TCP.tcpSpace
     , searchNodeAddress = TCP.nodeIP &&& TCP.tcpPort
-    , searchQuery       = getNodes prober tcp
+    , searchQuery       = asyncGetNodes prober tcp
+    , searchQueryCancel = cancelQuery (TCP.tcpClient tcp)
     , searchAlpha       = 8
     , searchK           = 16
     }