1 files changed, 14 insertions, 7 deletions

diff --git a/tox-crypto/src/Crypto/Tox.hs b/tox-crypto/src/Crypto/Tox.hs
index 602ead0a..aea684ec 100644
--- a/tox-crypto/src/Crypto/Tox.hs
+++ b/tox-crypto/src/Crypto/Tox.hs
@@ -41,6 +41,7 @@ module Crypto.Tox
     , decodePlain
     -- , computeSharedSecret
     , lookupSharedSecret
+    , lookupSharedSecretSTM
     , lookupNonceFunction
     , lookupNonceFunctionSTM
     , Payload(..)
@@ -369,11 +370,14 @@ lookupSharedSecret :: TransportCrypto -> SecretKey -> PublicKey -> Nonce24 -> IO
 lookupSharedSecret crypto sk recipient nonce
     = ($ nonce) <$> lookupNonceFunction crypto sk recipient
 
+lookupSharedSecretSTM :: TransportCrypto -> SecretKey -> PublicKey -> Nonce24 -> STM State
+lookupSharedSecretSTM crypto sk recipient nonce
+    = ($ nonce) <$> lookupNonceFunctionSTM crypto sk recipient
+
 {-# INLINE lookupNonceFunction #-}
 lookupNonceFunction :: TransportCrypto -> SecretKey -> PublicKey -> IO (Nonce24 -> State)
 lookupNonceFunction c@(TransportCrypto{secretsCache}) sk recipient = do
-    now <- getPOSIXTime
-    atomically $ lookupNonceFunctionSTM now c sk recipient
+    atomically $ lookupNonceFunctionSTM c sk recipient
 
 -- | Dummy, used to indicate computeSharedSecret failure.  This is a way to provide
 -- the old non-failing interface compatibly.
@@ -389,8 +393,10 @@ dummyState = State hash xsalsa
 --   of the transaction fails, we may end up forgoing a computation that could have been cached.
 --   Use with care. In most circumstances you probably want 'lookupNonceFunction'. It also commits
 --   us to using TVars to store the cache.
-lookupNonceFunctionSTM :: POSIXTime -> TransportCrypto -> SecretKey -> PublicKey -> STM (Nonce24 -> State)
-lookupNonceFunctionSTM now TransportCrypto{secretsCache} sk recipient = do
+lookupNonceFunctionSTM :: TransportCrypto -> SecretKey -> PublicKey -> STM (Nonce24 -> State)
+lookupNonceFunctionSTM TransportCrypto{secretsCache} sk recipient = do
+    now <- readTVar $ secretsTime secretsCache
+    writeTVar (secretsTime secretsCache) $! now + 1
     mm <- readTVar $ sharedSecret secretsCache
     case MM.lookup' recipient mm of
         Nothing -> case computeSharedSecret sk recipient of
@@ -579,12 +585,13 @@ 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)))
+                                       (Down Word64)
+                                       (MinMaxPSQ' SecretKey (Down Word64) (Nonce24 -> State)))
+    , secretsTime :: TVar Word64
     }
 
 newSecretsCache :: IO SecretsCache
-newSecretsCache = atomically (SecretsCache <$> newTVar MM.empty)
+newSecretsCache = atomically (SecretsCache <$> newTVar MM.empty <*> newTVar 0)
 
 
 newtype SymmetricKey = SymmetricKey ByteString