1 files changed, 44 insertions, 5 deletions

diff --git a/dht/src/Network/Tox/TCP.hs b/dht/src/Network/Tox/TCP.hs
index 932b4ab3..a37c0310 100644
--- a/dht/src/Network/Tox/TCP.hs
+++ b/dht/src/Network/Tox/TCP.hs
@@ -223,11 +223,53 @@ getTCPNodes tcp seeking dst = do
 getUDPNodes :: TCPClient err Nonce8 -> NodeId -> UDP.NodeInfo -> IO (Maybe ([UDP.NodeInfo], [UDP.NodeInfo], Maybe ()))
 getUDPNodes tcp seeking dst = fmap fst . resultToMaybe <$> getUDPNodes' tcp seeking dst
 
+
 getUDPNodes' :: TCPClient err Nonce8 -> NodeId -> UDP.NodeInfo -> IO (QR.Result (([UDP.NodeInfo], [UDP.NodeInfo], Maybe ()), NodeInfo))
 getUDPNodes' tcp seeking dst0 = do
+    goGetUDPNodes tcp seeking dst0 (return Canceled) $ \meth gateway dst -> do
+        r <- sendQuery (tcpClient tcp) meth (AnnounceRequest zeros32 seeking UDP.zeroID) gateway
+        forM r $ \response -> do
+            let (ns,_,mb) = either (const ([],[],Nothing)) (unwrapAnnounceResponse Nothing dst) $ response
+            return ( (ns,ns, const () <$> mb), gateway )
+
+-- Failure case, currently not treated as special.
+-- The current searchQuery type demands a valid Nonce8 is returned
+-- even if we were unable to send a query.
+fixmeNonce :: Nonce8
+fixmeNonce = Nonce8 0
+
+asyncUDPNodes :: TCPClient err Nonce8
+                       -> NodeId
+                       -> UDP.NodeInfo
+                       -> (Nonce8
+                           -> QR.Result (([UDP.NodeInfo], [UDP.NodeInfo], Maybe ()), NodeInfo)
+                           -> IO ())
+                       -> IO Nonce8
+asyncUDPNodes tcp seeking dst0 withResult =
+    goGetUDPNodes tcp seeking dst0 (return fixmeNonce) $ \meth gateway dst -> do
+        asyncQuery (tcpClient tcp) meth (AnnounceRequest zeros32 seeking UDP.zeroID) gateway $
+            \qid response -> do
+                let wut response =
+                        let (ns,_,mb) = either (const ([],[],Nothing)) (unwrapAnnounceResponse Nothing dst) $ response
+                        in ( (ns,ns, const () <$> mb), gateway )
+                withResult qid $ fmap wut response
+
+type Meth x = MethodSerializer
+                        Nonce8
+                        x -- NodeInfo
+                        (Bool, RelayPacket)
+                        PacketNumber
+                        AnnounceRequest
+                        (Either String AnnounceResponse)
+
+goGetUDPNodes :: TCPClient err Nonce8 -> NodeId -> UDP.NodeInfo
+                 -> IO a
+                 -> (Meth x -> NodeInfo -> UDP.NodeInfo -> IO a)
+                 -> IO a
+goGetUDPNodes tcp seeking dst0 fail go = do
     mgateway <- atomically $ tcpGetGateway tcp dst0
     case mgateway of
-      Nothing -> return Canceled
+      Nothing -> fail
       Just gateway -> do
         (b,c,n24) <- atomically $ do
             b <- transportNewKey (tcpCrypto tcp)
@@ -267,10 +309,7 @@ getUDPNodes' tcp seeking dst0 = do
                             -> decrypt (wrap0 n24') r >>= decodePlain
                         x   -> Left $ "getUDPNodes: unwrapResponse fail " ++ show x
                     }
-        r <- sendQuery (tcpClient tcp) meth (AnnounceRequest zeros32 seeking UDP.zeroID) gateway
-        forM r $ \response -> do
-            let (ns,_,mb) = either (const ([],[],Nothing)) (unwrapAnnounceResponse Nothing dst) $ response
-            return ( (ns,ns, const () <$> mb), gateway )
+        go meth gateway dst
 
 
 handleOOB :: PublicKey -> ByteString -> NodeInfo -> NodeInfo -> IO (Maybe (x -> x))