1 files changed, 13 insertions, 6 deletions

diff --git a/src/Network/Tox/DHT/Handlers.hs b/src/Network/Tox/DHT/Handlers.hs
index 2bce382a..25244a9b 100644
--- a/src/Network/Tox/DHT/Handlers.hs
+++ b/src/Network/Tox/DHT/Handlers.hs
@@ -179,7 +179,7 @@ newRouting addr crypto update4 update6 = do
         cbvar <- newTVar HashMap.empty
         return $ \client ->
             -- Now we have a client, so tell the BucketRefresher how to search and ping.
-            let updIO r = updateRefresherIO (nodeSearch client) (ping client) r
+            let updIO r = updateRefresherIO (nodeSearch client cbvar) (ping client) r
             in Routing { tentativeId     = tentative_info
                        , committee4      = committee4
                        , committee6      = committee6
@@ -376,11 +376,18 @@ unsendNodes _                    = Nothing
 unwrapNodes :: SendNodes -> ( [NodeInfo], [NodeInfo], Maybe () )
 unwrapNodes (SendNodes ns) = (ns,ns,Just ())
 
-getNodes :: Client -> NodeId -> NodeInfo -> IO (Maybe ([NodeInfo],[NodeInfo],Maybe ()))
-getNodes client nid addr = do
+getNodes :: Client -> TVar (HashMap NodeId [NodeInfoCallback]) -> NodeId -> NodeInfo -> IO (Maybe ([NodeInfo],[NodeInfo],Maybe ()))
+getNodes client cbvar nid addr = do
     -- hPutStrLn stderr $ show addr ++ " <-- getnodes " ++ show nid
     reply <- QR.sendQuery client (serializer GetNodesType DHTGetNodes unsendNodes) (GetNodes nid) addr
     -- hPutStrLn stderr $ show addr ++ " -sendnodes-> " ++ show reply
+    forM_ (join reply) $ \(SendNodes ns) -> 
+        forM_ ns $ \n -> do
+            atomically $ do
+                mcbs <- HashMap.lookup (nodeId n) <$> readTVar cbvar
+                forM_ mcbs $ \cbs -> do
+                    forM_ cbs $ \cb -> do
+                        rumoredAddress cb (nodeAddr addr) n
     return $ fmap unwrapNodes $ join reply
 
 updateRouting :: Client -> Routing -> OnionRouter -> NodeInfo -> Message -> IO ()
@@ -474,9 +481,9 @@ handlers crypto _       CookieRequestType = Just $ MethodHandler (isCookieReques
 handlers _      _       DHTRequestType    = Just $ NoReply (isDHTRequest snd) $ dhtRequestH
 handlers _      _       typ               = error $ "TODO DHT handlers " ++ show typ
 
-nodeSearch :: Client -> Search NodeId (IP,PortNumber) () NodeInfo NodeInfo
-nodeSearch client  = Search
+nodeSearch :: Client -> TVar (HashMap NodeId [NodeInfoCallback]) -> Search NodeId (IP,PortNumber) () NodeInfo NodeInfo
+nodeSearch client cbvar = Search
     { searchSpace       = toxSpace
     , searchNodeAddress = nodeIP &&& nodePort
-    , searchQuery       = getNodes client
+    , searchQuery       = getNodes client cbvar
     }