1 files changed, 77 insertions, 39 deletions

diff --git a/src/Network/Tox/Crypto/Handlers.hs b/src/Network/Tox/Crypto/Handlers.hs
index 34d2d87e..e8cdfcdd 100644
--- a/src/Network/Tox/Crypto/Handlers.hs
+++ b/src/Network/Tox/Crypto/Handlers.hs
@@ -64,7 +64,7 @@ newSessionsState crypto unrechook hooks = do
                       , defaultUnrecognizedHook = unrechook
                       }
 
-data HandshakeParams 
+data HandshakeParams
         = HParam
             { hpTheirBaseNonce :: Maybe Nonce24 -- ignore and generate your own
             , hpOtherCookie :: Maybe Cookie
@@ -76,6 +76,74 @@ data HandshakeParams
 newHandShakeData :: TransportCrypto -> HandshakeParams -> HandshakeData
 newHandShakeData = error "todo"
 
+-- | called when we recieve a crypto handshake with valid cookie
+freshCryptoSession :: NetCryptoSessions -> SockAddr -> HandshakeParams -> IO ()
+freshCryptoSession  sessions
+                    addr
+                    hp@(HParam
+                              { hpTheirBaseNonce = Just theirBaseNonce
+                              , hpOtherCookie    = Just otherCookie
+                              , hpTheirSessionKeyPublic = theirSessionKey
+                              , hpMySecretKey        = key
+                              , hpCookieRemotePubkey = remotePublicKey
+                              , hpCookieRemoteDhtkey = remoteDhtPublicKey
+                              }) -> do
+    let crypto = transportCrypto sessions
+        allsessions = netCryptoSessions sessions
+    ncState0 <- atomically $ newTVar Accepted
+    ncTheirBaseNonce0 <- atomically $ newTVar theirBaseNonce
+    n24 <- atomically $ transportNewNonce crypto
+    let myhandshakeData = newHandShakeData crypto hp
+        plain = encodePlain myhandshakeData
+        state = computeSharedSecret key remoteDhtPublicKey n24
+        encrypted = encrypt state plain
+        myhandshake = Handshake { handshakeCookie = otherCookie
+                                , handshakeNonce = n24
+                                , handshakeData = encrypted
+                                }
+    ncMyPacketNonce0 <- atomically $ newTVar (baseNonce myhandshakeData)
+    ncHandShake0 <- atomically $ newTVar (Just myhandshake)
+    cookie0 <- atomically $ newTVar (Just otherCookie)
+    newsession <- generateSecretKey
+    ncHooks0 <- atomically $ newTVar (defaultHooks sessions)
+    ncUnrecognizedHook0 <- atomically $ newTVar (defaultUnrecognizedHook sessions)
+    ncGroups0 <- atomically $ newTVar (Map.empty)
+    let netCryptoSession =
+            NCrypto { ncState         = ncState0
+                    , ncTheirBaseNonce= ncTheirBaseNonce0
+                    , ncMyPacketNonce   = ncMyPacketNonce0
+                    , ncHandShake     = ncHandShake0
+                    , ncCookie        = cookie0
+                    , ncTheirSessionPublic = Just theirSessionKey
+                    , ncSessionSecret = newsession
+                    , ncSockAddr      = addr
+                    , ncHooks         = ncHooks0
+                    , ncUnrecognizedHook = ncUnrecognizedHook0
+                    , ncAllSessions = sessions
+                    , ncGroups = ncGroups0
+                    }
+    atomically $ modifyTVar allsessions (Map.insert addr netCryptoSession)
+
+-- | 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
+    ncState0 <- atomically $ readTVar (ncState session)
+    ncTheirBaseNonce0 <- atomically $ readTVar (ncTheirBaseNonce session)
+    (Cookie _ _ presentCookieDHTKey )<- atomically $ readTVar (ncCookie session)
+    if (ncState0 >= Accepted)
+        -- If the nonce in the handshake and the dht key are both the same as
+        -- the ones we have saved, assume we already handled this and this is a
+        -- duplicate handshake packet, otherwise disregard everything, and
+        -- refresh all state.
+        --
+     then when (  ncTheirBaseNonce0 /= hpTheirBaseNonce hp
+               || presentCookieDHTKey /= hpCookieRemoteDhtkey hp
+               ) $ freshCryptoSession sessions addr hp
+     else if ( ncTheirBaseNonce0 /= hpTheirBaseNonce)
+            then freshCryptoSession sessions addr hp -- basenonce mismatch, trigger refresh
+            else atomically $ modifyTVar (ncState session) Accepted
+
+
 cryptoNetHandler :: NetCryptoSessions -> SockAddr -> NetCrypto -> IO (Maybe (NetCrypto -> NetCrypto))
 cryptoNetHandler sessions addr (NetHandshake (Handshake (Cookie n24 ecookie) nonce24 encrypted)) = do
     -- Handle Handshake Message
@@ -99,7 +167,7 @@ cryptoNetHandler sessions addr (NetHandshake (Handshake (Cookie n24 ecookie) non
             guard (cookieHash == digest)
             -- known friend?
             -- todo
-            return 
+            return
              HParam
                 { hpTheirBaseNonce = Just baseNonce
                 , hpOtherCookie    = Just otherCookie
@@ -119,43 +187,13 @@ cryptoNetHandler sessions addr (NetHandshake (Handshake (Cookie n24 ecookie) non
                       , hpCookieRemoteDhtkey = remoteDhtPublicKey
                       }) -> do
             sessionsmap <- atomically $ readTVar allsessions
-            -- Do a lookup, in case we decide to handle the update case differently
+            -- Do a lookup, so we can handle the update case differently
             case Map.lookup addr sessionsmap of
-                _ -> do -- create new session
-                    ncState0 <- atomically $ newTVar Accepted
-                    ncTheirBaseNonce0 <- atomically $ newTVar theirBaseNonce
-                    n24 <- atomically $ transportNewNonce crypto
-                    let myhandshakeData = newHandShakeData crypto hp
-                        plain = encodePlain myhandshakeData
-                        state = computeSharedSecret key remoteDhtPublicKey n24
-                        encrypted = encrypt state plain
-                        myhandshake = Handshake { handshakeCookie = otherCookie
-                                                , handshakeNonce = n24
-                                                , handshakeData = encrypted
-                                                }
-                    ncMyPacketNonce0 <- atomically $ newTVar (baseNonce myhandshakeData)
-                    ncHandShake0 <- atomically $ newTVar (Just myhandshake)
-                    cookie0 <- atomically $ newTVar (Just otherCookie)
-                    newsession <- generateSecretKey
-                    ncHooks0 <- atomically $ newTVar (defaultHooks sessions)
-                    ncUnrecognizedHook0 <- atomically $ newTVar (defaultUnrecognizedHook sessions)
-                    ncGroups0 <- atomically $ newTVar (Map.empty)
-                    let netCryptoSession = 
-                            NCrypto { ncState         = ncState0
-                                    , ncTheirBaseNonce= ncTheirBaseNonce0
-                                    , ncMyPacketNonce   = ncMyPacketNonce0
-                                    , ncHandShake     = ncHandShake0
-                                    , ncCookie        = cookie0
-                                    , ncTheirSessionPublic = Just theirSessionKey
-                                    , ncSessionSecret = newsession
-                                    , ncSockAddr      = addr
-                                    , ncHooks         = ncHooks0
-                                    , ncUnrecognizedHook = ncUnrecognizedHook0
-                                    , ncAllSessions = sessions
-                                    , ncGroups = ncGroups0
-                                    }
-                    atomically $ modifyTVar allsessions (Map.insert addr netCryptoSession)
-    return Nothing 
+                Nothing -> freshCryptoSession sessions addr hp -- create new session
+                Just session -> updateCryptoSession sessions addr hp session -- update existing session
+    return Nothing
+
+
 cryptoNetHandler sessions addr (NetCrypto (CryptoPacket nonce16 encrypted)) = do
     let crypto = transportCrypto sessions
         allsessions = netCryptoSessions sessions
@@ -221,7 +259,7 @@ cryptoNetHandler sessions addr (NetCrypto (CryptoPacket nonce16 encrypted)) = do
         last2Bytes (Nonce24 bs) = case S.decode (B.drop 22 bs) of
                                     Right n -> n
                                     _ -> error "unreachable-last2Bytes"
-        dATA_NUM_THRESHOLD = 21845 -- = 65535 / 3 
+        dATA_NUM_THRESHOLD = 21845 -- = 65535 / 3
 
 -- | handles nothing
 defaultCryptoDataHooks :: Map.Map MessageType [NetCryptoHook]