2 files changed, 17 insertions, 8 deletions

diff --git a/src/Network/Tox.hs b/src/Network/Tox.hs
index 68714224..41deff42 100644
--- a/src/Network/Tox.hs
+++ b/src/Network/Tox.hs
@@ -449,7 +449,9 @@ newTox keydb addr mbSessionsState suppliedDHTKey = do
     dhtclient <- newClient drg dhtnet0 DHT.classify (myAddr tbl4 tbl6) (DHT.handlers crypto . mkrouting) id
                     $ \client net -> onInbound (DHT.updateRouting client (mkrouting client) orouter) net
 
-    orouter <- forkRouteBuilder orouter $ \nid ni -> fmap (\(_,ns,_)->ns) <$> DHT.getNodes dhtclient nid ni
+    orouter' <- forkRouteBuilder orouter
+        $ \nid ni -> fmap (\(_,ns,_)->ns)
+                     <$> DHT.getNodes dhtclient (DHT.nodesOfInterest $ mkrouting dhtclient) nid ni
 
     toks <- do
         nil <- nullSessionTokens
@@ -459,7 +461,7 @@ newTox keydb addr mbSessionsState suppliedDHTKey = do
     onionclient <- newClient oniondrg onionnet (const Onion.classify)
                     (getOnionAlias crypto $ R.thisNode <$> readTVar (DHT.routing4 $ mkrouting dhtclient))
                     (const $ Onion.handlers onionnet (mkrouting dhtclient) toks keydb)
-                    (hookQueries orouter DHT.transactionKey)
+                    (hookQueries orouter' DHT.transactionKey)
                     (const id)
 
     return Tox
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
     }