1 files changed, 91 insertions, 53 deletions

diff --git a/src/Network/Tox/Crypto/Handlers.hs b/src/Network/Tox/Crypto/Handlers.hs
index bde286fe..0027e414 100644
--- a/src/Network/Tox/Crypto/Handlers.hs
+++ b/src/Network/Tox/Crypto/Handlers.hs
@@ -225,7 +225,7 @@ data NetCryptoSession = NCrypto
     , ncHandShake          :: TVar (UponHandshake (Handshake Encrypted))
     , ncCookie             :: TVar (UponCookie Cookie) -- ^ Cookie issued by remote peer
     , ncTheirDHTKey        :: UponDHTKey PublicKey
-    , ncTheirSessionPublic :: UponHandshake PublicKey
+    , ncTheirSessionPublic :: TVar (UponHandshake PublicKey)
     , ncSessionSecret      :: SecretKey
     , ncSockAddr           :: UponDHTKey SockAddr
     -- The remaining fields correspond to implementation specific state --
@@ -550,22 +550,14 @@ freshCryptoSession  sessions
     mbpktoq
      <- case mbtheirSessionKey of
          Nothing -> return NeedHandshake
-         Just theirSessionKey -> do
-            let toWireIO = do
-                    f <- lookupNonceFunction crypto newsession theirSessionKey
-                    atomically $ do
-                        n24 <- readTVar ncMyPacketNonce0
-                        let n24plus1 = incrementNonce24 n24
-                        trace ("ncMyPacketNonce+1=" ++ show n24plus1) $ writeTVar ncMyPacketNonce0 n24plus1
-                        return (return (f n24, n24, ncOutgoingIdMap0))
-            pktoq <- atomically $ PQ.newOutGoing pktq ncToWire toWireIO 0 (outboundQueueCapacity sessions) 0
-            return (HaveHandshake pktoq)
+         Just theirSessionKey -> createNetCryptoOutQueue sessions newsession theirSessionKey pktq ncMyPacketNonce0 ncOutgoingIdMap0
     lastNQ <- atomically (CB.new 10 0 :: STM (CyclicBuffer (Bool,(ViewSnapshot,InOrOut CryptoMessage))))
     listeners <- atomically $ newTVar IntMap.empty
     msgNum <- atomically $ newTVar 0
     dropNum <- atomically $ newTVar 0
     theirbasenonce <- atomically $ readTVar ncTheirBaseNonce0
     dput XNetCrypto $ "freshCryptoSession: Session ncTheirBaseNonce=" ++ show theirbasenonce
+    ncTheirSessionPublic0 <- atomically $ newTVar (frmMaybe mbtheirSessionKey)
     let netCryptoSession0 =
             NCrypto { ncState              = ncState0
                     , ncMyPublicKey        = toPublic key
@@ -576,7 +568,7 @@ freshCryptoSession  sessions
                     , ncHandShake          = ncHandShake0
                     , ncCookie             = cookie0
                     , ncTheirDHTKey        = HaveDHTKey remoteDhtPublicKey
-                    , ncTheirSessionPublic = frmMaybe mbtheirSessionKey
+                    , ncTheirSessionPublic = ncTheirSessionPublic0
                     , ncSessionSecret      = newsession
                     , ncSockAddr           = HaveDHTKey addr
                     , ncHooks              = ncHooks0
@@ -594,7 +586,40 @@ freshCryptoSession  sessions
                     , ncLastNMsgs          = lastNQ
                     , ncListeners          = listeners
                     }
+    case mbpktoq of
+        NeedHandshake -> return ()
+        HaveHandshake pktoq -> runUponHandshake netCryptoSession0 addr pktoq
+
+type NetCryptoOutQueue = PQ.PacketOutQueue (State,Nonce24,RangeMap TArray Word8 TVar)
+                                           CryptoMessage
+                                           (CryptoPacket Encrypted)
+                                           CryptoData
+
+createNetCryptoOutQueue :: NetCryptoSessions -> SecretKey -> PublicKey -> PacketQueue CryptoData
+                        -> TVar Nonce24    -> RangeMap TArray Word8 TVar     -> IO (UponHandshake NetCryptoOutQueue)
+createNetCryptoOutQueue sessions newsession theirSessionKey pktq ncMyPacketNonce0 ncOutgoingIdMap0 = do
+            let crypto = transportCrypto sessions
+            let toWireIO = do
+                    f <- lookupNonceFunction crypto newsession theirSessionKey
+                    atomically $ do
+                        n24 <- readTVar ncMyPacketNonce0
+                        let n24plus1 = incrementNonce24 n24
+                        trace ("ncMyPacketNonce+1=" ++ show n24plus1) $ writeTVar ncMyPacketNonce0 n24plus1
+                        return (return (f n24, n24, ncOutgoingIdMap0))
+            pktoq <- atomically $ PQ.newOutGoing pktq ncToWire toWireIO 0 (outboundQueueCapacity sessions) 0
+            return (HaveHandshake pktoq)
+
+
+runUponHandshake :: NetCryptoSession -> SockAddr -> NetCryptoOutQueue -> IO ()
+runUponHandshake netCryptoSession0 addr pktoq = do
+    let sessions = ncAllSessions netCryptoSession0
+        pktq = ncPacketQueue netCryptoSession0
+        remotePublicKey = ncTheirPublicKey netCryptoSession0
+        crypto = transportCrypto sessions
+        allsessions = netCryptoSessions sessions
+        allsessionsByKey = netCryptoSessionsByKey sessions
     -- launch dequeue thread
+    -- (In terms of data dependency, this thread could be launched prior to handshake)
     threadid <- forkIO $ do
         tid <- myThreadId
         labelThread tid ("NetCryptoDequeue." ++ show (key2id remotePublicKey))
@@ -602,44 +627,41 @@ freshCryptoSession  sessions
             cd <- atomically $ PQ.dequeue pktq
             _ <- runCryptoHook (netCryptoSession0 {ncDequeueThread=Just tid}) (bufferData cd)
             loop
-    case mbpktoq of
-        NeedHandshake -> return ()
-        HaveHandshake pktoq -> do
-            -- launch dequeueOutgoing thread
-            threadidOutgoing <- forkIO $ do
-                tid <- myThreadId
-                labelThread tid ("NetCryptoDequeueOutgoing." ++ show (key2id remotePublicKey))
-                fix $ \loop -> do
-                    (_,pkt) <- atomically $ PQ.dequeueOutgoing pktoq
-                    dput XNetCrypto "NetCryptoDequeueOutgoing thread... Sending encrypted Packet"
-                    sendMessage (sessionTransport sessions) addr (NetCrypto pkt)
-                    loop
-            -- launch ping thread
-            fuzz <- randomRIO (0,2000)
-            pingMachine <- forkPingMachine ("NetCrypto." ++ show (key2id remotePublicKey)) (15000 + fuzz) 2000
-            -- 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))
-            -- run announceNewSessionHooks
-            hooks <- atomically $ readTVar (announceNewSessionHooks sessions)
-            flip fix (hooks,netCryptoSession) $ \loop (hooks,session) ->
-                case hooks of
-                    [] -> return ()
-                    (h:hs) -> do
-                        r <- h Nothing session
-                        case r of
-                          Just f -> loop (hs, f session)
-                          Nothing -> return ()
+    -- launch dequeueOutgoing thread
+    threadidOutgoing <- forkIO $ do
+        tid <- myThreadId
+        labelThread tid ("NetCryptoDequeueOutgoing." ++ show (key2id remotePublicKey))
+        fix $ \loop -> do
+            (_,pkt) <- atomically $ PQ.dequeueOutgoing pktoq
+            dput XNetCrypto "NetCryptoDequeueOutgoing thread... Sending encrypted Packet"
+            sendMessage (sessionTransport sessions) addr (NetCrypto pkt)
+            loop
+    -- launch ping thread
+    fuzz <- randomRIO (0,2000)
+    pingMachine <- forkPingMachine ("NetCrypto." ++ show (key2id remotePublicKey)) (15000 + fuzz) 2000
+    -- 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))
+    -- run announceNewSessionHooks
+    hooks <- atomically $ readTVar (announceNewSessionHooks sessions)
+    flip fix (hooks,netCryptoSession) $ \loop (hooks,session) ->
+        case hooks of
+            [] -> return ()
+            (h:hs) -> do
+                r <- h Nothing session
+                case r of
+                  Just f -> loop (hs, f session)
+                  Nothing -> return ()
 
 -- | Called when we get a handshake, but there's already a session entry.
-updateCryptoSession :: NetCryptoSessions -> SockAddr -> HandshakeParams -> NetCryptoSession -> IO ()
-updateCryptoSession sessions addr hp session = do
+updateCryptoSession :: NetCryptoSessions -> SockAddr -> HandshakeParams -> NetCryptoSession -> Handshake Encrypted -> IO ()
+updateCryptoSession sessions addr hp session handshake = do
     ncState0 <- atomically $ readTVar (ncState session)
     ncTheirBaseNonce0 <- atomically $ readTVar (ncTheirBaseNonce session)
     if (ncState0 >= {-Accepted-}InProgress AwaitingSessionPacket)
@@ -656,10 +678,25 @@ updateCryptoSession sessions addr hp session = do
           dput XNetCrypto ("   ncTheirDHTKey=" ++ show (ncTheirDHTKey session)
                          ++ bool "{/=}" "{==}" (ncTheirDHTKey session == HaveDHTKey (hpCookieRemoteDhtkey hp))
                          ++ "hpCookieRemoteDhtkey=" ++ show (hpCookieRemoteDhtkey hp))
-          when (  -- Just ncTheirBaseNonce0 /= hpTheirBaseNonce hp -- XXX: Do we really want to compare base nonce here?
+          if   (  -- Just ncTheirBaseNonce0 /= hpTheirBaseNonce hp -- XXX: Do we really want to compare base nonce here?
                   -- ||
                ncTheirDHTKey session /= HaveDHTKey (hpCookieRemoteDhtkey hp)
-               ) $ freshCryptoSession sessions addr hp
+               ) then freshCryptoSession sessions addr hp
+                 else do
+                    atomically $ do
+                        writeTVar (ncTheirBaseNonce session)    (frmMaybe (hpTheirBaseNonce hp))
+                        writeTVar (ncTheirSessionPublic session) (frmMaybe (hpTheirSessionKeyPublic hp))
+                        writeTVar (ncHandShake session) (HaveHandshake handshake)
+                    case ncOutgoingQueue session of
+                        NeedHandshake -> do
+                            case hpTheirSessionKeyPublic hp of
+                                Just sessionpubkey -> do
+                                    pktoq <- createNetCryptoOutQueue sessions (ncSessionSecret session) sessionpubkey
+                                                                (ncPacketQueue session) (ncMyPacketNonce session) (ncOutgoingIdMap session)
+                                    case pktoq of
+                                        NeedHandshake -> dput XNetCrypto "Unexpectedly missing ncOutgoingQueue"
+                                        HaveHandshake pktoq -> runUponHandshake session addr pktoq
+                        HaveHandshake pktoq -> runUponHandshake session addr pktoq
      else do
           dput XNetCrypto "updateCryptoSession else clause"
           dput XNetCrypto ("   ncTheirBaseNonce0=" ++ show ncTheirBaseNonce0
@@ -674,7 +711,7 @@ cryptoNetHandler sessions addr (NetHandshake hs)  = handshakeH     sessions addr
 cryptoNetHandler sessions addr (NetCrypto    pkt) = sessionPacketH sessions addr pkt
 
 handshakeH :: NetCryptoSessions -> SockAddr -> Handshake Encrypted -> IO (Maybe a)
-handshakeH sessions addr (Handshake (Cookie n24 ecookie) nonce24 encrypted) = do
+handshakeH sessions addr hshake@(Handshake (Cookie n24 ecookie) nonce24 encrypted) = do
     dput XNetCrypto ("RECIEVED HANDSHAKE from " ++ show addr)
     -- Handle Handshake Message
     let crypto = transportCrypto sessions
@@ -746,7 +783,7 @@ handshakeH sessions addr (Handshake (Cookie n24 ecookie) nonce24 encrypted) = do
                     freshCryptoSession sessions addr hp -- create new session
                 Just session -> do
                     dput XNetCrypto "sockaddr ALREADY in session map, so updateCryptoSession"
-                    updateCryptoSession sessions addr hp session -- update existing session
+                    updateCryptoSession sessions addr hp session hshake -- update existing session
     return Nothing
 
 sessionPacketH :: NetCryptoSessions -> SockAddr -> CryptoPacket Encrypted -> IO (Maybe (x -> x))
@@ -771,8 +808,9 @@ sessionPacketH sessions addr (CryptoPacket nonce16 encrypted) = do
             let diff :: Word16
                 diff      = nonce16 - (last2Bytes theirBaseNonce) -- truncating to Word16
                 tempNonce = addtoNonce24 theirBaseNonce (fromIntegral diff)    -- expanding to Word
+            mbpublickey <- atomically (readTVar ncTheirSessionPublic)
             lr <- fmap join $ sequence $ do -- Either Monad --
-                        pubkey <- maybeToEither ncTheirSessionPublic
+                        pubkey <- maybeToEither mbpublickey
                         Right $ do -- IO Monad
                         dput XNetCrypto $ "cryptoNetHandler: pubkey =  " ++ show (key2id $ pubkey)
                         dput XNetCrypto $ "cryptoNetHandler: theirBaseNonce =  " ++ show theirBaseNonce