1 files changed, 7 insertions, 6 deletions

diff --git a/src/Network/Tox/Crypto/Handlers.hs b/src/Network/Tox/Crypto/Handlers.hs
index e8cdfcdd..b8e99d2d 100644
--- a/src/Network/Tox/Crypto/Handlers.hs
+++ b/src/Network/Tox/Crypto/Handlers.hs
@@ -38,6 +38,7 @@ data NetCryptoSession = NCrypto { ncState         :: TVar NetCryptoSessionStatus
                                 , ncMyPacketNonce :: TVar Nonce24 -- base nonce + packet number
                                 , ncHandShake     :: TVar (Maybe (Handshake Encrypted))
                                 , ncCookie        :: TVar (Maybe Cookie)
+                                , ncTheirDHTKey   :: PublicKey
                                 , ncTheirSessionPublic :: Maybe PublicKey
                                 , ncSessionSecret :: SecretKey
                                 , ncSockAddr      :: SockAddr
@@ -87,7 +88,7 @@ freshCryptoSession  sessions
                               , hpMySecretKey        = key
                               , hpCookieRemotePubkey = remotePublicKey
                               , hpCookieRemoteDhtkey = remoteDhtPublicKey
-                              }) -> do
+                              }) = do
     let crypto = transportCrypto sessions
         allsessions = netCryptoSessions sessions
     ncState0 <- atomically $ newTVar Accepted
@@ -114,6 +115,7 @@ freshCryptoSession  sessions
                     , ncMyPacketNonce   = ncMyPacketNonce0
                     , ncHandShake     = ncHandShake0
                     , ncCookie        = cookie0
+                    , ncTheirDHTKey   = remoteDhtPublicKey
                     , ncTheirSessionPublic = Just theirSessionKey
                     , ncSessionSecret = newsession
                     , ncSockAddr      = addr
@@ -129,19 +131,18 @@ updateCryptoSession :: NetCryptoSessions -> SockAddr -> HandshakeParams -> NetCr
 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
+     then when (  Just ncTheirBaseNonce0 /= hpTheirBaseNonce hp
+               || ncTheirDHTKey session /= hpCookieRemoteDhtkey hp
                ) $ freshCryptoSession sessions addr hp
-     else if ( ncTheirBaseNonce0 /= hpTheirBaseNonce)
+     else if ( Just ncTheirBaseNonce0 /= hpTheirBaseNonce hp)
             then freshCryptoSession sessions addr hp -- basenonce mismatch, trigger refresh
-            else atomically $ modifyTVar (ncState session) Accepted
+            else atomically $ writeTVar (ncState session) Accepted
 
 
 cryptoNetHandler :: NetCryptoSessions -> SockAddr -> NetCrypto -> IO (Maybe (NetCrypto -> NetCrypto))