1 files changed, 23 insertions, 17 deletions

diff --git a/src/Network/Tox/DHT/Handlers.hs b/src/Network/Tox/DHT/Handlers.hs
index d010f36d..4576fc85 100644
--- a/src/Network/Tox/DHT/Handlers.hs
+++ b/src/Network/Tox/DHT/Handlers.hs
@@ -259,26 +259,31 @@ ping client addr = do
     hPutStrLn stderr $ show addr ++ " -pong-> " ++ show reply
     maybe (return False) (\Pong -> return True) $ join reply
 
+
+saveCookieKey :: TVar [(SockAddr, (Int, PublicKey))] -> SockAddr -> PublicKey -> STM ()
+saveCookieKey var saddr pk = do
+    cookiekeys <- readTVar var
+    case break (\(stored,_) -> stored == saddr) cookiekeys of
+        (xs,[])                               -> writeTVar var $ (saddr, (1  ,pk)) : xs
+        (xs,(_,(c,stored)):ys) | stored == pk -> writeTVar var $ (saddr, (c+1,pk)) : xs ++ ys
+        _                                     -> retry -- Wait for requests to this address
+                                                       -- under a different key to time out
+                                                       -- before we try this key.
+
+loseCookieKey :: TVar [(SockAddr, (Int, PublicKey))] -> SockAddr -> PublicKey -> STM ()
+loseCookieKey var saddr pk = do
+    cookiekeys <- readTVar var
+    case break (\(stored,_) -> stored == saddr) cookiekeys of
+        (xs,(_,(1,stored)):ys) | stored == pk -> writeTVar var $                     xs ++ ys
+        (xs,(_,(c,stored)):ys) | stored == pk -> writeTVar var $ (saddr, (c-1,pk)) : xs ++ ys
+        _                                     -> return () -- unreachable?
+
+
 cookieRequest :: TransportCrypto -> Client -> PublicKey -> NodeInfo -> IO (Maybe Cookie)
 cookieRequest crypto client localUserKey addr = do
     let sockAddr = nodeAddr addr
         nid = id2key $ nodeId addr
-    let incAddr sockMap
-            = case partition ((==sockAddr) . fst) sockMap of
-                    ([],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' == 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 (pendingCookies crypto)
-            writeTVar (pendingCookies crypto) mp
-            return mp
-    let cookieSerializer
+        cookieSerializer
             = MethodSerializer
                    { methodTimeout  = \tid addr -> return (addr, 5000000)
                    , method    = CookieRequestType
@@ -286,9 +291,10 @@ cookieRequest crypto client localUserKey addr = do
                    , unwrapResponse = fmap snd . unCookie
                    }
         cookieRequest = CookieRequest localUserKey
+    atomically $ saveCookieKey (pendingCookies crypto) sockAddr nid
     hPutStrLn stderr $ show addr ++ " <-- cookieRequest"
     reply <- QR.sendQuery client cookieSerializer cookieRequest addr
-    atomically $ modifyTVar (pendingCookies crypto) decAddr
+    atomically $ loseCookieKey (pendingCookies crypto) sockAddr nid
     hPutStrLn stderr $ show addr ++ " -cookieResponse-> " ++ show reply
     return $ join reply