1 files changed, 16 insertions, 7 deletions

diff --git a/src/Network/Tox/Crypto/Handlers.hs b/src/Network/Tox/Crypto/Handlers.hs
index 493ae925..b47aeac1 100644
--- a/src/Network/Tox/Crypto/Handlers.hs
+++ b/src/Network/Tox/Crypto/Handlers.hs
@@ -589,6 +589,7 @@ freshCryptoSession  sessions
                     , ncLastNMsgs          = lastNQ
                     , ncListeners          = listeners
                     }
+    addSessionToMapIfNotThere sessions addr netCryptoSession0
     case mbpktoq of
         NeedHandshake -> return ()
         HaveHandshake pktoq -> runUponHandshake netCryptoSession0 addr pktoq
@@ -610,11 +611,23 @@ createNetCryptoOutQueue sessions newsession theirSessionKey pktq ncMyPacketNonce
                         trace ("ncMyPacketNonce+1=" ++ show n24plus1
                             ++ "\n toWireIO: theirSessionKey = " ++ show (key2id theirSessionKey)
                             ++ "\n toWireIO: my public session key = " ++ show (key2id (toPublic newsession))
-                              ) $ writeTVar ncMyPacketNonce0 n24plus1
+                              ) $ writeTVar ncMyPacketNonce0 n24plus1
                         return (return (f n24, n24, ncOutgoingIdMap0))
             pktoq <- atomically $ PQ.newOutGoing pktq ncToWire toWireIO 0 (outboundQueueCapacity sessions) 0
             return (HaveHandshake pktoq)
 
+-- | add this session to the lookup maps, unless its already in them
+addSessionToMapIfNotThere :: NetCryptoSessions -> SockAddr -> NetCryptoSession -> IO ()
+addSessionToMapIfNotThere sessions addr netCryptoSession = do
+    atomically $ do
+        let remotePublicKey = ncTheirPublicKey netCryptoSession
+            allsessions     = netCryptoSessions sessions
+            allsessionsByKey= netCryptoSessionsByKey sessions
+        modifyTVar allsessions (Map.insert addr netCryptoSession)
+        byKeyResult <- readTVar allsessionsByKey >>= return . Map.lookup remotePublicKey
+        case byKeyResult of
+            Nothing -> modifyTVar allsessionsByKey (Map.insert remotePublicKey [netCryptoSession])
+            Just xs -> modifyTVar allsessionsByKey (Map.insert remotePublicKey (netCryptoSession:xs))
 
 runUponHandshake :: NetCryptoSession -> SockAddr -> NetCryptoOutQueue -> IO ()
 runUponHandshake netCryptoSession0 addr pktoq = do
@@ -648,12 +661,7 @@ runUponHandshake netCryptoSession0 addr pktoq = do
     -- update session with thread ids
     let netCryptoSession = netCryptoSession0 {ncDequeueThread=Just threadid, ncPingMachine=Just pingMachine}
     -- add this session to the lookup maps
-    atomically $ do
-        modifyTVar allsessions (Map.insert addr netCryptoSession)
-        byKeyResult <- readTVar allsessionsByKey >>= return . Map.lookup remotePublicKey
-        case byKeyResult of
-            Nothing -> modifyTVar allsessionsByKey (Map.insert remotePublicKey [netCryptoSession])
-            Just xs -> modifyTVar allsessionsByKey (Map.insert remotePublicKey (netCryptoSession:xs))
+    addSessionToMapIfNotThere sessions addr netCryptoSession
     -- run announceNewSessionHooks
     hooks <- atomically $ readTVar (announceNewSessionHooks sessions)
     flip fix (hooks,netCryptoSession) $ \loop (hooks,session) ->
@@ -762,6 +770,7 @@ decryptHandshake crypto hshake@(Handshake (Cookie n24 ecookie) nonce24 encrypted
                             , handshakeData   = pure hsdata
                             } )
 
+toHandshakeParams :: (SecretKey, Handshake Identity) -> HandshakeParams
 toHandshakeParams (key,hs)
     = let hd = runIdentity $ handshakeData hs
           Cookie _ cd0 = handshakeCookie hs