Tox: Shared secrets cache.

1 files changed, 79 insertions, 10 deletions

diff --git a/src/Crypto/Tox.hs b/src/Crypto/Tox.hs
index 37725aea..7cfbd193 100644
--- a/src/Crypto/Tox.hs
+++ b/src/Crypto/Tox.hs
@@ -8,6 +8,7 @@
 {-# LANGUAGE ExplicitNamespaces #-}
 {-# LANGUAGE TypeOperators #-}
 {-# LANGUAGE MagicHash, UnboxedTuples, BangPatterns #-}
+{-# LANGUAGE NamedFieldPuns #-}
 module Crypto.Tox
     ( PublicKey
     , publicKey
@@ -18,6 +19,8 @@ module Crypto.Tox
     , toPublic
     , SymmetricKey(..)
     , TransportCrypto(..)
+    , SecretsCache
+    , newSecretsCache
     , Encrypted
     , Encrypted8(..)
     , type (∘)(..)
@@ -89,6 +92,11 @@ import Network.Socket (SockAddr)
 import GHC.Exts (Word(..))
 import GHC.Prim
 import Data.Word64Map (fitsInInt)
+import Data.MinMaxPSQ (MinMaxPSQ')
+import qualified Data.MinMaxPSQ as MM
+import Data.Time.Clock.POSIX
+import Data.Hashable
+import System.IO.Unsafe (unsafePerformIO)
 
 -- | A 16-byte mac and an arbitrary-length encrypted stream.
 newtype Encrypted a = Encrypted ByteString
@@ -236,23 +244,73 @@ encrypt (State hash crypt) (Plain m) = Encrypted $ B.append (encode a) c
 
 -- (Poly1305.State, XSalsa.State)
 computeSharedSecret :: SecretKey -> PublicKey -> Nonce24 -> State
-computeSharedSecret sk recipient nonce = State hash crypt
+computeSharedSecret sk recipient = \nonce ->
+    let -- cipher state
+        st0 = XSalsa.initialize 20 k nonce
+        -- Poly1305 key
+        (rs, crypt) = XSalsa.combine st0 zs where Nonce32 zs = zeros32
+        -- Since rs is 32 bytes, this pattern should never fail...
+        Cryptonite.CryptoPassed hash = Poly1305.initialize rs
+    in State hash crypt
  where
     -- diffie helman
     shared = ecdh (Proxy :: Proxy Curve_X25519) sk recipient
     -- shared secret XSalsa key
     k = hsalsa20 shared zeros24
-    -- cipher state
-    st0 = XSalsa.initialize 20 k nonce
-    -- Poly1305 key
-    (rs, crypt) = XSalsa.combine st0 zs where Nonce32 zs = zeros32
-    -- Since rs is 32 bytes, this pattern should never fail...
-    Cryptonite.CryptoPassed hash = Poly1305.initialize rs
+
+unsafeFirstWord64 :: ByteArrayAccess ba => ba -> Word64
+unsafeFirstWord64 ba = unsafePerformIO $ BA.withByteArray ba peek
+
+instance Hashable PublicKey where
+    hashWithSalt salt pk = hashWithSalt salt (unsafeFirstWord64 pk)
+
+instance Hashable SecretKey where
+    hashWithSalt salt sk = hashWithSalt salt (unsafeFirstWord64 sk)
+
+instance Ord PublicKey where compare = unsafeCompare32Bytes
+instance Ord SecretKey where compare = unsafeCompare32Bytes
+
+unsafeCompare32Bytes :: (ByteArrayAccess ba, ByteArrayAccess bb)
+                        => ba -> bb -> Ordering
+unsafeCompare32Bytes ba bb =
+    unsafePerformIO $ BA.withByteArray ba
+                    $ \pa -> BA.withByteArray bb
+                    $ \pb -> unsafeCompare32Bytes' 3 pa pb
+
+unsafeCompare32Bytes' :: Int -> Ptr Word64 -> Ptr Word64 -> IO Ordering
+unsafeCompare32Bytes' !n !pa !pb = do
+    a <- peek pa
+    b <- peek pb
+    case compare a b of
+        EQ     -> if n == 0
+                    then return EQ
+                    else unsafeCompare32Bytes' (n - 1)
+                                               (pa `plusPtr` 8)
+                                               (pb `plusPtr` 8)
+        result -> return result
+
+
 
 lookupSharedSecret :: TransportCrypto -> SecretKey -> PublicKey -> Nonce24 -> IO State
-lookupSharedSecret crypto sk recipient nonce = do
-    -- TODO
-    return $ computeSharedSecret sk recipient nonce
+lookupSharedSecret TransportCrypto{secretsCache} sk recipient nonce = do
+    now <- getPOSIXTime
+    atomically $ do
+    mm <- readTVar $ sharedSecret secretsCache
+    case MM.lookup' recipient mm of
+        Nothing -> do
+            let miss = computeSharedSecret sk recipient
+            writeTVar (sharedSecret secretsCache)
+                      (MM.insertTake' 160 recipient (MM.singleton' sk miss (Down now)) (Down now) mm)
+            return $ miss nonce
+        Just (stamp,smm) -> do
+            let (r,v) = case MM.lookup' sk smm of
+                    Nothing | let miss = computeSharedSecret sk recipient
+                                      -> (miss, MM.insertTake' 3 sk miss (Down now) smm)
+                    Just (stamp2,hit) -> (hit , MM.insertTake' 3 sk hit  (Down now) smm)
+            writeTVar (sharedSecret secretsCache)
+                      (MM.insertTake' 160 recipient v (Down now) mm)
+            return $ r nonce
+
 
 hsalsa20 :: (ByteArrayAccess t, ByteArrayAccess t1) => t1 -> t -> BA.ScrubbedBytes
 hsalsa20 k n = BA.append a b
@@ -415,6 +473,16 @@ getAliasedAsymm = flip Asymm <$> get <*> getPublicKey <*> get
 putAliasedAsymm :: Serialize a => Asymm a -> Put
 putAliasedAsymm (Asymm key nonce dta) = put nonce >> putPublicKey key >> put dta
 
+data SecretsCache = SecretsCache
+    { sharedSecret :: TVar (MinMaxPSQ' PublicKey
+                                       (Down POSIXTime)
+                                       (MinMaxPSQ' SecretKey (Down POSIXTime) (Nonce24 -> State)))
+    }
+
+newSecretsCache :: IO SecretsCache
+newSecretsCache = atomically (SecretsCache <$> newTVar MM.empty)
+
+
 newtype SymmetricKey = SymmetricKey ByteString
 
 data TransportCrypto = TransportCrypto
@@ -428,6 +496,7 @@ data TransportCrypto = TransportCrypto
     , transportNewNonce :: STM Nonce24
     , userKeys :: TVar [(SecretKey,PublicKey)]
     , pendingCookies :: TVar [(SockAddr, (Int, PublicKey))]
+    , secretsCache :: SecretsCache
     }
 
 getPublicKey :: S.Get PublicKey