HandshakeCache remembers sent handshake data.

1 files changed, 130 insertions, 0 deletions

diff --git a/HandshakeCache.hs b/HandshakeCache.hs
new file mode 100644
index 00000000..6f9d466f
--- /dev/null
+++ b/HandshakeCache.hs
@@ -0,0 +1,130 @@
+module HandshakeCache where
+
+import Control.Concurrent.STM
+import Control.Monad
+import Data.Functor.Identity
+import qualified Data.Map.Strict as Map
+         ;import Data.Map.Strict (Map)
+import Data.Time.Clock.POSIX
+import Network.Socket
+import Data.Bool
+
+import Crypto.Hash
+
+import Crypto.Tox
+import qualified Data.MinMaxPSQ     as MM
+         ;import Data.MinMaxPSQ     (MinMaxPSQ')
+import DPut
+import Network.Tox.Crypto.Transport (Handshake, HandshakeData (..))
+import Network.Tox.DHT.Handlers     (createCookieSTM)
+import Network.Tox.DHT.Transport    (Cookie (..), CookieData (..), NodeInfo,
+                                     key2id, nodeInfo)
+import Network.Tox.Handshake
+
+data HandshakeCache = HandshakeCache
+    { -- Note that currently we are storing sent handshakes keyed by the
+      -- remotely issued cookie.  This probably means that it's possible for
+      -- one your contacts that you are trying to open a session with to
+      -- prevent you from opening a session with another contact if they know
+      -- the cookie that person issued you.
+      hscTable  :: TVar (MinMaxPSQ' (Cookie Encrypted) POSIXTime (SecretKey,HandshakeData))
+    , hscSend   :: SockAddr -> Handshake Encrypted -> IO ()
+    , hscCrypto :: TransportCrypto
+    , hscPendingCookies :: TVar (Map (PublicKey,PublicKey) ())
+    }
+
+
+newHandshakeCache :: TransportCrypto -> (SockAddr -> Handshake Encrypted -> IO ()) -> IO HandshakeCache
+newHandshakeCache crypto send = atomically $ do
+    tbl <- newTVar MM.empty
+    pcs <- newTVar Map.empty
+    return HandshakeCache
+        { hscTable  = tbl
+        , hscSend   = send
+        , hscCrypto = crypto
+        , hscPendingCookies = pcs
+        }
+
+getSentHandshake :: HandshakeCache
+                    -> SecretKey
+                    -> SockAddr
+                    -> Cookie Identity -- locally issued
+                    -> Cookie Encrypted -- remotely issued
+                    -> IO (Maybe (SecretKey, HandshakeData))
+getSentHandshake hscache me their_addr (Cookie _ (Identity cd)) ecookie = do
+    now <- getPOSIXTime
+    io <- atomically $ do
+        m <- checkExpiry now . MM.lookup' ecookie <$> readTVar (hscTable hscache)
+        case m of
+            Just s  -> return $ return $ Just s
+            Nothing -> do
+                let them = longTermKey cd
+                case nodeInfo (key2id $ dhtKey cd) their_addr of
+                    Left _ -> return $ return Nothing -- Non-internet address.
+                    Right their_node -> do
+                        (s,hs) <- cacheHandshakeSTM hscache me them their_node ecookie now
+                        return $ do
+                            hscSend hscache their_addr hs
+                            return $ Just s
+    io
+
+
+checkExpiry :: POSIXTime -> Maybe (POSIXTime,r) -> Maybe r
+checkExpiry now m = do
+    (tm,s) <- m
+    guard $ tm + 5 {- seconds -} > now
+    return s
+
+hashCookie :: HashAlgorithm a => Cookie Encrypted -> Digest a
+hashCookie (Cookie n24 encrypted)
+    = hashFinalize $ hashUpdate (hashUpdate hashInit n24) encrypted
+
+cacheHandshakeSTM :: HandshakeCache
+                    -> SecretKey        -- ^ my ToxID key
+                    -> PublicKey        -- ^ them
+                    -> NodeInfo         -- ^ their DHT node
+                    -> Cookie Encrypted -- ^ issued to me by them
+                    -> POSIXTime        -- ^ current time
+                    -> STM ((SecretKey,HandshakeData), Handshake Encrypted)
+cacheHandshakeSTM hscache me them their_node ecookie timestamp = do
+    newsession <- transportNewKey (hscCrypto hscache)
+    freshCookie <- createCookieSTM timestamp (hscCrypto hscache) their_node them
+    n24 <- transportNewNonce (hscCrypto hscache)
+    let hsdata = HandshakeData
+            { baseNonce   = n24
+            , sessionKey  = toPublic newsession
+            , cookieHash  = hashCookie ecookie
+            , otherCookie = freshCookie
+            }
+    hs <- encodeHandshake timestamp (hscCrypto hscache) me them ecookie hsdata
+    modifyTVar' (hscTable hscache) $ MM.insertTake' 20 ecookie (newsession,hsdata) timestamp
+    return ((newsession,hsdata),hs)
+
+cacheHandshake :: HandshakeCache
+                    -> SecretKey
+                    -> PublicKey
+                    -> NodeInfo
+                    -> Cookie Encrypted
+                    -> IO (Handshake Encrypted)
+cacheHandshake hscache me them their_node ecookie = do
+    timestamp <- getPOSIXTime
+    dput XNetCrypto $ "cacheHandshake " ++ show (key2id them,ecookie)
+    atomically $ snd <$> cacheHandshakeSTM hscache me them their_node ecookie timestamp
+
+haveCachedCookie :: HandshakeCache
+                    -> PublicKey
+                    -> PublicKey
+                    -> STM Bool
+haveCachedCookie hscache me them = do
+    m <- Map.lookup (me,them) <$> readTVar (hscPendingCookies hscache)
+    return $ maybe True (const False) m
+
+
+setPendingCookie :: HandshakeCache
+                    -> PublicKey
+                    -> PublicKey
+                    -> Bool
+                    -> STM ()
+setPendingCookie hscache me them pending = do
+    modifyTVar' (hscPendingCookies hscache) $ Map.alter (const $ bool Nothing (Just ()) pending)
+                                                        (me,them)