1 files changed, 103 insertions, 107 deletions

diff --git a/src/Network/Tox/Crypto/Handlers.hs b/src/Network/Tox/Crypto/Handlers.hs
index eabbc9b0..722d8507 100644
--- a/src/Network/Tox/Crypto/Handlers.hs
+++ b/src/Network/Tox/Crypto/Handlers.hs
@@ -7,7 +7,7 @@ module Network.Tox.Crypto.Handlers where
 import Network.Tox.NodeId
 import Network.Tox.Crypto.Transport
 import Network.Tox.DHT.Transport (Cookie(..),CookieData(..), CookieRequest(..), NoSpam(..))
-import Network.Tox.DHT.Handlers (Client, cookieRequest, createCookie )
+import Network.Tox.DHT.Handlers (Client, cookieRequest, createCookieSTM )
 import Crypto.Tox
 import Control.Arrow
 import Control.Concurrent.STM
@@ -395,25 +395,18 @@ data HandshakeParams
             , hpCookieRemoteDhtkey    :: PublicKey
             }
 
-newHandShakeData :: TransportCrypto -> Nonce24 -> HandshakeParams -> SockAddr -> PublicKey -> IO (Maybe HandshakeData)
-newHandShakeData crypto basenonce (HParam {hpOtherCookie,hpMySecretKey,hpCookieRemotePubkey,hpCookieRemoteDhtkey}) addr mySessionPublic
+newHandShakeData :: POSIXTime -> TransportCrypto -> Nonce24 -> HandshakeParams -> SockAddr -> PublicKey -> STM (Maybe HandshakeData)
+newHandShakeData timestamp crypto basenonce (HParam {hpOtherCookie,hpMySecretKey,hpCookieRemotePubkey,hpCookieRemoteDhtkey}) addr mySessionPublic
     = do
     freshCookie
       <- case  nodeInfo (key2id hpCookieRemoteDhtkey) addr  of
-            Right nodeinfo -> Just <$> createCookie crypto nodeinfo hpCookieRemotePubkey
+            Right nodeinfo -> Just <$> createCookieSTM timestamp crypto nodeinfo hpCookieRemotePubkey
             Left er        -> return Nothing
     let hinit                = hashInit
         Cookie n24 encrypted = hpOtherCookie
         hctx                 = hashUpdate hinit n24
         hctx'                = hashUpdate hctx encrypted
         digest               = hashFinalize hctx'
--- parameters addr {--> SockAddr -}
---    mbcookie <- case hpOtherCookie of
---                    Nothing -> case (nodeInfo hpCookieRemoteDhtkey addr) of
---                                Right nodeinfo -> cookieRequest crypto netCryptoDHTClient (toPublic hpMySecretKey) nodeinfo
---                                Left er -> return Nothing
---                    Just c -> return (Just c)
-
     return $
       fmap (\freshCookie' ->
             HandshakeData
@@ -489,9 +482,11 @@ ncToWire getState seqno bufend pktno msg = do
 -- then increment it regardless. (Keep addr in MinMaxPSQ in Roster.Contact)
 --
 -- This function sends a handshake response packet.
-freshCryptoSession :: NetCryptoSessions -> SockAddr -> HandshakeParams -> IO ()
+freshCryptoSession :: NetCryptoSessions -> SockAddr -> SecretKey -> POSIXTime -> HandshakeParams -> STM (Maybe (Handshake Encrypted),IO ())
 freshCryptoSession  sessions
                     addr
+                    newsession
+                    timestamp
                     hp@(HParam
                               { hpTheirBaseNonce = mbtheirBaseNonce
                               , hpOtherCookie    = otherCookie
@@ -503,20 +498,20 @@ freshCryptoSession  sessions
     let crypto = transportCrypto sessions
         allsessions = netCryptoSessions sessions
         allsessionsByKey = netCryptoSessionsByKey sessions
-    sessionId <- atomically $ do
+        dmsg msg = trace msg (return ())
+    sessionId <- do
         x <- readTVar (nextSessionId sessions)
         modifyTVar    (nextSessionId sessions) (+1)
         return x
-    -- ncState0 <- atomically $ newTVar Accepted -- (InProgress AwaitingSessionPacket)
-    ncState0 <- atomically $ newTVar (if isJust mbtheirBaseNonce
+    -- ncState0 <- newTVar Accepted -- (InProgress AwaitingSessionPacket)
+    ncState0 <- newTVar (if isJust mbtheirBaseNonce
                                         then InProgress AwaitingSessionPacket
                                         else InProgress AwaitingHandshake)
-    ncTheirBaseNonce0 <- atomically $ newTVar (frmMaybe mbtheirBaseNonce)
-    n24 <- atomically $ transportNewNonce crypto
-    state <- lookupSharedSecret crypto key remotePublicKey n24
-    newBaseNonce <- atomically $ transportNewNonce crypto
-    newsession <- generateSecretKey
-    mbMyhandshakeData <- newHandShakeData crypto newBaseNonce hp addr (toPublic newsession)
+    ncTheirBaseNonce0 <- newTVar (frmMaybe mbtheirBaseNonce)
+    n24 <- transportNewNonce crypto
+    state <- ($ n24) <$> lookupNonceFunctionSTM timestamp crypto key remotePublicKey
+    newBaseNonce <- transportNewNonce crypto
+    mbMyhandshakeData <- newHandShakeData timestamp crypto newBaseNonce hp addr (toPublic newsession)
     let encodeHandshake myhandshakeData = let plain = encodePlain myhandshakeData
     --                                        state = computeSharedSecret key remoteDhtPublicKey n24
                                               encrypted = encrypt state plain
@@ -525,16 +520,14 @@ freshCryptoSession  sessions
                                                            , handshakeData = encrypted
                                                            }
     let myhandshake= encodeHandshake <$> mbMyhandshakeData
-    ncHandShake0 <- atomically $ newTVar (frmMaybe myhandshake)
-    forM myhandshake $ \response_handshake -> do
-        sendHandshake sessions addr response_handshake
-    ncMyPacketNonce0 <- atomically $ newTVar newBaseNonce
-    cookie0 <- atomically $ newTVar (HaveCookie otherCookie)
-    ncHooks0 <- atomically $ newTVar (defaultHooks sessions)
-    ncUnrecognizedHook0 <- atomically $ newTVar (defaultUnrecognizedHook sessions)
-    ncIncomingTypeArray0 <- atomically $ newTVar (msgTypeArray sessions)
+    ncHandShake0 <- newTVar (frmMaybe myhandshake)
+    ncMyPacketNonce0 <- newTVar newBaseNonce
+    cookie0 <- newTVar (HaveCookie otherCookie)
+    ncHooks0 <- newTVar (defaultHooks sessions)
+    ncUnrecognizedHook0 <- newTVar (defaultUnrecognizedHook sessions)
+    ncIncomingTypeArray0 <- newTVar (msgTypeArray sessions)
     let idMap = foldl (\mp (x,y) -> W64.insert x y mp) W64.empty (zip [0..255] [0..255])
-    (ncOutgoingIdMap0,lossyEscapeIdMap,losslessEscapeIdMap) <- atomically $ do
+    (ncOutgoingIdMap0,lossyEscapeIdMap,losslessEscapeIdMap) <- do
         idmap <- emptySTMRangeMap
         insertArrayAt idmap  0 (A.listArray (0,255) [0 .. 255])
         -- the 2 escape ranges are adjacent, so put them in one array:
@@ -546,21 +539,21 @@ freshCryptoSession  sessions
         lossyEsc    <- newTVar $ A.listArray (0,255) [0 .. 255]
         losslessEsc <- newTVar $ A.listArray (0,255) [0 .. 255]
         return (idmap,lossyEsc,losslessEsc)
-    ncView0 <- atomically $ newTVar (sessionView sessions)
-    pktq <- atomically $ PQ.new (inboundQueueCapacity sessions) 0
-    bufstart <- atomically $ newTVar 0
+    ncView0 <- newTVar (sessionView sessions)
+    pktq <- PQ.new (inboundQueueCapacity sessions) 0
+    bufstart <- newTVar 0
     mbpktoq
      <- case mbtheirSessionKey of
          Nothing -> return NeedHandshake
          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
-    dput XNetCrypto $ "freshCryptoSession: My Session Public =" ++ show (key2id $ toPublic newsession)
-    ncTheirSessionPublic0 <- atomically $ newTVar (frmMaybe mbtheirSessionKey)
+    lastNQ <- CB.new 10 0 :: STM (CyclicBuffer (Bool,(ViewSnapshot,InOrOut CryptoMessage)))
+    listeners <- newTVar IntMap.empty
+    msgNum <- newTVar 0
+    dropNum <- newTVar 0
+    theirbasenonce <- readTVar ncTheirBaseNonce0
+    dmsg $ "freshCryptoSession: Session ncTheirBaseNonce=" ++ show theirbasenonce
+    dmsg $ "freshCryptoSession: My Session Public =" ++ show (key2id $ toPublic newsession)
+    ncTheirSessionPublic0 <- newTVar (frmMaybe mbtheirSessionKey)
     let netCryptoSession0 =
             NCrypto { ncState              = ncState0
                     , ncMyPublicKey        = toPublic key
@@ -590,9 +583,11 @@ freshCryptoSession  sessions
                     , ncListeners          = listeners
                     }
     addSessionToMapIfNotThere sessions addr netCryptoSession0
-    case mbpktoq of
-        NeedHandshake -> return ()
-        HaveHandshake pktoq -> runUponHandshake netCryptoSession0 addr pktoq
+    maybeLaunchMissles
+        <- case mbpktoq of
+            NeedHandshake -> return (return ())
+            HaveHandshake pktoq -> return (runUponHandshake netCryptoSession0 addr pktoq)
+    return (myhandshake,maybeLaunchMissles)
 
 type NetCryptoOutQueue = PQ.PacketOutQueue (State,Nonce24,RangeMap TArray Word8 TVar)
                                            CryptoMessage
@@ -600,7 +595,7 @@ type NetCryptoOutQueue = PQ.PacketOutQueue (State,Nonce24,RangeMap TArray Word8
                                            CryptoData
 
 createNetCryptoOutQueue :: NetCryptoSessions -> SecretKey -> PublicKey -> PacketQueue CryptoData
-                        -> TVar Nonce24    -> RangeMap TArray Word8 TVar     -> IO (UponHandshake NetCryptoOutQueue)
+                        -> TVar Nonce24    -> RangeMap TArray Word8 TVar     -> STM (UponHandshake NetCryptoOutQueue)
 createNetCryptoOutQueue sessions newsession theirSessionKey pktq ncMyPacketNonce0 ncOutgoingIdMap0 = do
             let crypto = transportCrypto sessions
             let toWireIO = do
@@ -613,33 +608,33 @@ createNetCryptoOutQueue sessions newsession theirSessionKey pktq ncMyPacketNonce
                             ++ "\n toWireIO: my public session key = " ++ show (key2id (toPublic newsession))
                               ) $ writeTVar ncMyPacketNonce0 n24plus1
                         return (return (f n24, n24, ncOutgoingIdMap0))
-            pktoq <- atomically $ PQ.newOutGoing pktq ncToWire toWireIO 0 (outboundQueueCapacity sessions) 0
+            pktoq <- 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 :: NetCryptoSessions -> SockAddr -> NetCryptoSession -> STM ()
 addSessionToMapIfNotThere sessions addr netCryptoSession = do
-    dput XNetCrypto $ "addSessionToMapIfNotThere sockaddr = " ++ show addr ++ ", sessionid = " ++ show (ncSessionId netCryptoSession)
-    atomically $ do
-        let remotePublicKey = ncTheirPublicKey netCryptoSession
-            allsessions     = netCryptoSessions sessions
-            allsessionsByKey= netCryptoSessionsByKey sessions
-        byAddrResult <- readTVar allsessions >>= return . Map.lookup addr
-        case byAddrResult of
-            Just (NCrypto { ncSessionId = staleId }) -> do
-                -- manually remove the stale session from the by-key map
-                modifyTVar allsessionsByKey (Map.map (filter ((/=staleId) . ncSessionId)))
-            Nothing -> return () -- nothing to remove
-        -- write session to by-addr map regardless of whether one is in there,
-        -- it should overwrite on match
-        modifyTVar allsessions (Map.insert addr netCryptoSession)
-        -- Now insert new session into by-key map
-        byKeyResult <- readTVar allsessionsByKey >>= return . Map.lookup remotePublicKey
-        case byKeyResult of
-            Nothing -> modifyTVar allsessionsByKey (Map.insert remotePublicKey [netCryptoSession])
-            Just xs -> do
-                -- in case we're using the same long term key on different IPs ...
-                modifyTVar allsessionsByKey (Map.insert remotePublicKey (netCryptoSession:xs))
+    let dmsg msg = trace msg (return ())
+    dmsg $ "addSessionToMapIfNotThere sockaddr = " ++ show addr ++ ", sessionid = " ++ show (ncSessionId netCryptoSession)
+    let remotePublicKey = ncTheirPublicKey netCryptoSession
+        allsessions     = netCryptoSessions sessions
+        allsessionsByKey= netCryptoSessionsByKey sessions
+    byAddrResult <- readTVar allsessions >>= return . Map.lookup addr
+    case byAddrResult of
+        Just (NCrypto { ncSessionId = staleId }) -> do
+            -- manually remove the stale session from the by-key map
+            modifyTVar allsessionsByKey (Map.map (filter ((/=staleId) . ncSessionId)))
+        Nothing -> return () -- nothing to remove
+    -- write session to by-addr map regardless of whether one is in there,
+    -- it should overwrite on match
+    modifyTVar allsessions (Map.insert addr netCryptoSession)
+    -- Now insert new session into by-key map
+    byKeyResult <- readTVar allsessionsByKey >>= return . Map.lookup remotePublicKey
+    case byKeyResult of
+        Nothing -> modifyTVar allsessionsByKey (Map.insert remotePublicKey [netCryptoSession])
+        Just xs -> do
+            -- in case we're using the same long term key on different IPs ...
+            modifyTVar allsessionsByKey (Map.insert remotePublicKey (netCryptoSession:xs))
 
 runUponHandshake :: NetCryptoSession -> SockAddr -> NetCryptoOutQueue -> IO ()
 runUponHandshake netCryptoSession0 addr pktoq = do
@@ -673,7 +668,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
-    addSessionToMapIfNotThere sessions addr netCryptoSession
+    -- atomically $ addSessionToMapIfNotThere sessions addr netCryptoSession
     -- run announceNewSessionHooks
     hooks <- atomically $ readTVar (announceNewSessionHooks sessions)
     flip fix (hooks,netCryptoSession) $ \loop (hooks,session) ->
@@ -691,10 +686,12 @@ runUponHandshake netCryptoSession0 addr pktoq = do
 --     2) handshake for new session (old session is lost?)
 
 --     3) we initiated, this a response
-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)
+updateCryptoSession :: NetCryptoSessions -> SockAddr -> SecretKey -> POSIXTime -> HandshakeParams
+                    -> NetCryptoSession -> Handshake Encrypted -> STM (Maybe (Handshake Encrypted), IO ())
+updateCryptoSession sessions addr newsession timestamp hp session handshake = do
+    let dmsg msg = trace msg (return ())
+    ncState0 <- readTVar (ncState session)
+    ncTheirBaseNonce0 <- readTVar (ncTheirBaseNonce session)
     if (ncState0 >= {-Accepted-}InProgress AwaitingSessionPacket)
         -- 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
@@ -702,42 +699,31 @@ updateCryptoSession sessions addr hp session handshake = do
         -- refresh all state.
         --
      then do
-          dput XNetCrypto "updateCryptoSession already accepted.."
-          dput XNetCrypto ("   ncTheirBaseNonce0=" ++ show ncTheirBaseNonce0
+          dmsg "updateCryptoSession already accepted.."
+          dmsg ("   ncTheirBaseNonce0=" ++ show ncTheirBaseNonce0
                          ++ bool "(/=)" "(==)" (toMaybe ncTheirBaseNonce0 == hpTheirBaseNonce hp)
                          ++ "hpTheirBaseNonce=" ++ show (hpTheirBaseNonce hp))
-          dput XNetCrypto ("   ncTheirDHTKey=" ++ show (ncTheirDHTKey session)
+          dmsg ("   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
-             --  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
+             )
+            then freshCryptoSession sessions addr newsession timestamp hp
+            else return (Nothing,return ())
      else do
-          dput XNetCrypto "updateCryptoSession else clause"
-          dput XNetCrypto ("   ncTheirBaseNonce0=" ++ show ncTheirBaseNonce0
+          dmsg "updateCryptoSession else clause"
+          dmsg ("   ncTheirBaseNonce0=" ++ show ncTheirBaseNonce0
                          ++ bool "(/=)" "(==)" (toMaybe ncTheirBaseNonce0 == hpTheirBaseNonce hp)
                          ++ "hpTheirBaseNonce=" ++ show (hpTheirBaseNonce hp))
           if ( ncTheirBaseNonce0 /= frmMaybe (hpTheirBaseNonce hp))
             then do
-                 dput XNetCrypto "basenonce mismatch, trigger refresh"
-                 freshCryptoSession sessions addr hp -- basenonce mismatch, trigger refresh
-            else atomically $ writeTVar (ncState session) {-Accepted-}(InProgress AwaitingSessionPacket)
+                dmsg "basenonce mismatch, trigger refresh"
+                freshCryptoSession sessions addr newsession timestamp hp -- basenonce mismatch, trigger refresh
+            else do
+                writeTVar (ncState session) {-Accepted-}(InProgress AwaitingSessionPacket)
+                return (Nothing,return ())
 
 anyRight :: Monad m => a -> [t] -> (t -> m (Either b b1)) -> m (Either a b1)
 anyRight e []     f = return $ Left e
@@ -819,15 +805,25 @@ handshakeH sessions addr hshake@(Handshake (Cookie n24 ecookie) nonce24 encrypte
                       , hpCookieRemoteDhtkey    = remoteDhtPublicKey
                       }) -> do
             dput XNetCrypto ("(NetCrypto)handshakeH: hpTheirBaseNonce = " ++ show theirBaseNonce)
-            sessionsmap <- atomically $ readTVar allsessions
+            -- IO action to get a new session key in case we need it in transaction to come
+            newsession <- generateSecretKey
             -- Do a lookup, so we can handle the update case differently
-            case Map.lookup addr sessionsmap of
-                Nothing -> do
-                    dput XNetCrypto "sockaddr not in session map, so freshCryptoSession"
-                    freshCryptoSession sessions addr hp -- create new session
-                Just session -> do
-                    dput XNetCrypto "sockaddr ALREADY in session map, so updateCryptoSession"
-                    updateCryptoSession sessions addr hp session hshake -- update existing session
+            let dmsg msg = trace msg (return ())
+            timestamp <- getPOSIXTime
+            (myhandshake,launchThreads)
+              <- atomically $ do
+                    sessionsmap <- readTVar allsessions
+                    case Map.lookup addr sessionsmap of
+                        Nothing -> do
+                            dmsg "sockaddr not in session map, so freshCryptoSession"
+                            freshCryptoSession sessions addr newsession timestamp hp -- create new session
+                        Just session -> do
+                            dmsg "sockaddr ALREADY in session map, so updateCryptoSession"
+                            updateCryptoSession sessions addr newsession timestamp hp session hshake -- update existing session
+            launchThreads
+            forM myhandshake $ \response_handshake -> do
+                sendHandshake sessions addr response_handshake
+            return ()
     return Nothing
 
 sessionPacketH :: NetCryptoSessions -> SockAddr -> CryptoPacket Encrypted -> IO (Maybe (x -> x))