1 files changed, 45 insertions, 13 deletions

diff --git a/dht/src/Network/Tox/DHT/Handlers.hs b/dht/src/Network/Tox/DHT/Handlers.hs
index d132da88..f4563a3b 100644
--- a/dht/src/Network/Tox/DHT/Handlers.hs
+++ b/dht/src/Network/Tox/DHT/Handlers.hs
@@ -133,8 +133,8 @@ data Routing = Routing
     { tentativeId     :: NodeInfo
     , committee4      :: TriadCommittee NodeId SockAddr
     , committee6      :: TriadCommittee NodeId SockAddr
-    , refresher4      :: BucketRefresher NodeId NodeInfo
-    , refresher6      :: BucketRefresher NodeId NodeInfo
+    , refresher4      :: BucketRefresher NodeId NodeInfo TransactionId
+    , refresher6      :: BucketRefresher NodeId NodeInfo TransactionId
     , nodesOfInterest :: TVar (HashMap NodeId [NodeInfoCallback])
     }
 
@@ -172,6 +172,20 @@ routing4 Routing { refresher4 = BucketRefresher { refreshBuckets } } = refreshBu
 routing6 :: Routing -> TVar (R.BucketList NodeInfo)
 routing6 Routing { refresher6 = BucketRefresher { refreshBuckets } } = refreshBuckets
 
+nullTransactionId :: TransactionId
+nullTransactionId = TransactionId (Nonce8 0) (Nonce24 zeros24)
+
+nullSearch :: Search NodeId (IP, PortNumber) tok NodeInfo r TransactionId
+nullSearch = Search
+    { searchSpace       = toxSpace
+    , searchNodeAddress = nodeIP &&& nodePort
+    , searchQuery       = \_ _ f -> f nullTransactionId Canceled >> return nullTransactionId
+    , searchQueryCancel = \_ _ -> return ()
+    , searchAlpha       = 1
+    , searchK           = 2
+    }
+
+
 newRouting :: SockAddr -> TransportCrypto
                        -> (TVar (R.BucketList NodeInfo) -> SockAddr -> STM ()) -- ^ invoked on IPv4 change
                        -> (TVar (R.BucketList NodeInfo) -> SockAddr -> STM ()) -- ^ invoked on IPv6 change
@@ -195,13 +209,6 @@ newRouting addr crypto update4 update6 = do
         -- We defer initializing the refreshSearch and refreshPing until we
         -- have a client to send queries with.
         let nullPing = const $ return False
-            nullSearch = Search
-                { searchSpace       = toxSpace
-                , searchNodeAddress = nodeIP &&& nodePort
-                , searchQuery       = \_ _ -> return Canceled
-                , searchAlpha       = 1
-                , searchK           = 2
-                }
         tbl4 <- newTVar $ R.nullTable (comparing nodeId) (\s -> hashWithSalt s . nodeId) tentative_info4 R.defaultBucketCount
         tbl6 <- newTVar $ R.nullTable (comparing nodeId) (\s -> hashWithSalt s . nodeId) tentative_info6 R.defaultBucketCount
         refresher4 <- newBucketRefresher tbl4 nullSearch nullPing
@@ -432,6 +439,30 @@ getNodesUDP :: Client -> TVar (HashMap NodeId [NodeInfoCallback]) -> NodeId -> N
                       -> IO (QR.Result ([NodeInfo],[NodeInfo],Maybe ()))
 getNodesUDP client cbvar nid addr = getNodes client cbvar nid (Multi.UDP ==> addr)
 
+asyncGetNodes :: Client -> TVar (HashMap NodeId [NodeInfoCallback]) -> NodeId -> Multi.NodeInfo
+                   -> (TransactionId -> QR.Result ([NodeInfo],[NodeInfo],Maybe ()) -> IO ())
+                   -> IO TransactionId
+asyncGetNodes client cbvar nid addr withResult = do
+    QR.asyncQuery client (serializer GetNodesType DHTGetNodes unsendNodes) (GetNodes nid) addr $
+        \qid reply -> do
+            forM_ (join $ resultToMaybe reply) $ \(SendNodes ns) ->
+                forM_ ns $ \n -> do
+                    now <- getPOSIXTime
+                    atomically $ do
+                        mcbs <- HashMap.lookup (nodeId . udpNodeInfo $ n) <$> readTVar cbvar
+                        forM_ mcbs $ \cbs -> do
+                            forM_ cbs $ \cb -> do
+                                rumoredAddress cb now addr (udpNodeInfo n)
+            withResult qid $ case reply of
+                Success x -> maybe Canceled (Success . unwrapNodes) x
+                _         -> fmap (error "Network.Tox.DHT.Handlers.getNodes: the impossible happened!") reply
+
+asyncGetNodesUDP :: Client -> TVar (HashMap NodeId [NodeInfoCallback]) -> NodeId -> NodeInfo
+                      -> (TransactionId -> QR.Result ([NodeInfo],[NodeInfo],Maybe ()) -> IO ())
+                      -> IO TransactionId
+asyncGetNodesUDP client cbvar nid addr go = asyncGetNodes client cbvar nid (Multi.UDP ==> addr) go
+
+
 updateRouting :: Client -> Routing
                         -> (TVar (R.BucketList NodeInfo) -> RoutingTransition NodeInfo -> STM ())
                         -> Multi.NodeInfo
@@ -462,7 +493,7 @@ updateTable client routing orouter naddr = do
         Want_Both -> do dput XMisc "BUG:unreachable"
                         error $ "BUG:unreachable at " ++ __FILE__ ++ ":" ++ show __LINE__
  where
-    go :: TriadCommittee NodeId SockAddr -> BucketRefresher NodeId NodeInfo -> IO ()
+    go :: TriadCommittee NodeId SockAddr -> BucketRefresher NodeId NodeInfo TransactionId -> IO ()
     go committee refresher = do
         self <- atomically $ R.thisNode <$> readTVar (refreshBuckets refresher)
         -- dput XMisc $ "(tox)updateRouting: " ++ show (nodeIP self, nodeIP naddr)
@@ -473,7 +504,7 @@ updateTable client routing orouter naddr = do
 toxKademlia :: Client
                 -> (TVar (R.BucketList NodeInfo) -> RoutingTransition NodeInfo -> STM ())
                 -> TriadCommittee NodeId SockAddr
-                -> BucketRefresher NodeId NodeInfo
+                -> BucketRefresher NodeId NodeInfo TransactionId
                 -> Kademlia NodeId NodeInfo
 toxKademlia client orouter committee refresher
     = Kademlia quietInsertions
@@ -541,11 +572,12 @@ handlers crypto _       CookieRequestType = Just $ MethodHandler (isCookieReques
 handlers _      _       DHTRequestType    = Just $ NoReply (isDHTRequest snd) $ dhtRequestH
 handlers _      _       typ               = error $ "TODO DHT handlers " ++ show typ
 
-nodeSearch :: Client -> TVar (HashMap NodeId [NodeInfoCallback]) -> Search NodeId (IP,PortNumber) () NodeInfo NodeInfo
+nodeSearch :: Client -> TVar (HashMap NodeId [NodeInfoCallback]) -> Search NodeId (IP,PortNumber) () NodeInfo NodeInfo TransactionId
 nodeSearch client cbvar = Search
     { searchSpace       = toxSpace
     , searchNodeAddress = nodeIP &&& nodePort
-    , searchQuery       = getNodesUDP client cbvar
+    , searchQuery       = asyncGetNodesUDP client cbvar
+    , searchQueryCancel = cancelQuery client
     , searchAlpha       = 8
     , searchK           = 16
     }