4 files changed, 26 insertions, 15 deletions

diff --git a/examples/dhtd.hs b/examples/dhtd.hs
index 36b8e294..3fb1c641 100644
--- a/examples/dhtd.hs
+++ b/examples/dhtd.hs
@@ -734,10 +734,15 @@ main = do
 
             let toxDHT bkts = DHT
                     { dhtBuckets = bkts (Tox.toxRouting tox)
-                    , dhtPing    = Map.singleton "ping" $ DHTPing
-                                    { pingQuery = fmap (bool Nothing (Just ())) . Tox.ping (Tox.toxDHT tox)
-                                    , pingShowResult = show
-                                    }
+                    , dhtPing    = Map.fromList
+                                    [ ("ping", DHTPing
+                                            { pingQuery = fmap (bool Nothing (Just ())) . Tox.ping (Tox.toxDHT tox)
+                                            , pingShowResult = show
+                                            })
+                                    , ("cookie", DHTPing
+                                            { pingQuery = Tox.cookieRequest (Tox.toxCryptoKeys tox) (Tox.toxDHT tox)
+                                            , pingShowResult = show
+                                            })]
                     , dhtQuery   = Map.fromList
                                     [ ("node", DHTQuery (Tox.nodeSearch $ Tox.toxDHT tox)
                                                         (\ni -> fmap Tox.unwrapNodes
diff --git a/src/Crypto/Tox.hs b/src/Crypto/Tox.hs
index d5727254..8a65dfb4 100644
--- a/src/Crypto/Tox.hs
+++ b/src/Crypto/Tox.hs
@@ -81,6 +81,7 @@ import System.Endian
 import qualified Data.ByteString.Internal
 import Control.Concurrent.STM
 import Crypto.Error.Types                     (CryptoFailable (..), throwCryptoError)
+import Network.Socket (SockAddr)
 
 -- | A 16-byte mac and an arbitrary-length encrypted stream.
 newtype Encrypted a = Encrypted ByteString
@@ -367,6 +368,7 @@ data TransportCrypto = TransportCrypto
     , transportSymmetric :: STM SymmetricKey
     , transportNewNonce :: STM Nonce24
     , userKeys :: TVar [(SecretKey,PublicKey)]
+    , pendingCookies :: TVar [(SockAddr, (Int, PublicKey))]
     }
 
 getPublicKey :: S.Get PublicKey
diff --git a/src/Network/Tox.hs b/src/Network/Tox.hs
index e9220fcb..d434360f 100644
--- a/src/Network/Tox.hs
+++ b/src/Network/Tox.hs
@@ -112,6 +112,7 @@ newCrypto = do
         drg0 <- getSystemDRG
         return $ randomBytesGenerate 32 drg0 :: IO (ByteString, SystemDRG)
     noncevar <- atomically $ newTVar $ fst $ withDRG drg drgNew
+    cookieKeys <- atomically $ newTVar []
     hPutStrLn stderr $ "secret(tox) = " ++ DHT.showHex secret
     hPutStrLn stderr $ "public(tox) = " ++ DHT.showHex pubkey
     hPutStrLn stderr $ "symmetric(tox) = " ++ DHT.showHex symkey
@@ -129,6 +130,7 @@ newCrypto = do
             writeTVar noncevar drg2
             return nonce
         , userKeys = ukeys
+        , pendingCookies = cookieKeys
         }
 
 updateIP :: TVar (R.BucketList NodeInfo) -> SockAddr -> STM ()
diff --git a/src/Network/Tox/DHT/Handlers.hs b/src/Network/Tox/DHT/Handlers.hs
index 840e2e6b..38763f9c 100644
--- a/src/Network/Tox/DHT/Handlers.hs
+++ b/src/Network/Tox/DHT/Handlers.hs
@@ -243,39 +243,41 @@ ping client addr = do
     hPutStrLn stderr $ show addr ++ " -pong-> " ++ show reply
     maybe (return False) (\Pong -> return True) $ join reply
 
-cookieRequest :: TVar [(SockAddr,(Int,NodeInfo))] -> PublicKey -> Client -> NodeInfo -> IO (Maybe Cookie)
-cookieRequest tvar myDhtKey client addr = do
+cookieRequest :: TransportCrypto -> Client -> NodeInfo -> IO (Maybe Cookie)
+cookieRequest crypto client addr = do
     let sockAddr = nodeAddr addr
+        nid = id2key $ nodeId addr
     let incAddr sockMap
             = case partition ((==sockAddr) . fst) sockMap of
-                    ([],xs)                              -> insert (sockAddr, (1  ,addr)) xs
-                    ([(_,(c,addr'))],xs) | addr' == addr -> insert (sockAddr, (c+1,addr)) xs
+                    ([],xs)                             -> (sockAddr, (1  ,nid)) : xs
+                    ([(_,(c,addr'))],xs) | addr' == nid -> (sockAddr, (c+1,nid)) : xs
                     anythingElse -> error $ "unreachable at " ++ __FILE__ ++ show __LINE__ ++ show anythingElse ++ show (sockAddr,addr)
         decAddr sockMap
             = case partition ((==sockAddr) . fst) sockMap of
-                    ([],xs)                              -> xs -- unreachable?
-                    ([(_,(1,addr'))],xs) | addr' == addr -> xs
-                    ([(_,(c,addr'))],xs) | addr' == addr -> insert (sockAddr,(c-1,addr)) xs
+                    ([],xs)                             -> xs -- unreachable?
+                    ([(_,(1,addr'))],xs) | addr' == nid -> xs
+                    ([(_,(c,addr'))],xs) | addr' == nid -> (sockAddr,(c-1,nid)) : xs
                     anythingElse -> error $ "unreachable at " ++ __FILE__ ++ show __LINE__ ++ show anythingElse ++ show (sockAddr,addr)
     sockMap <- atomically $ do
-            mp <- incAddr <$> readTVar tvar
-            writeTVar tvar mp
+            mp <- incAddr <$> readTVar (pendingCookies crypto)
+            writeTVar (pendingCookies crypto) mp
             return mp
     let cookieSerializer
             = MethodSerializer
                    { methodTimeout  = \tid addr -> do
-                        modifyTVar tvar decAddr
+                        modifyTVar (pendingCookies crypto) decAddr
                         return (addr, 5000000)
                    , method    = CookieRequestType
                    , wrapQuery = \tid src dst cr -> DHTCookieRequest $ wrapAsymm tid src dst (, cr)
                    , unwrapResponse = fmap snd . unCookie
                    }
-        cookieRequest = CookieRequest myDhtKey
+        cookieRequest = CookieRequest (transportPublic crypto)
     hPutStrLn stderr $ show addr ++ " <-- cookieRequest"
     reply <- QR.sendQuery client cookieSerializer cookieRequest addr
     hPutStrLn stderr $ show addr ++ " -cookieResponse-> " ++ show reply
     return $ join reply
 
+unCookie :: DHTMessage t -> Maybe (t Cookie)
 unCookie (DHTCookie n24 fcookie) = Just fcookie
 unCookie _ = Nothing