1 files changed, 18 insertions, 3 deletions

diff --git a/src/Crypto/Tox.hs b/src/Crypto/Tox.hs
index 18cdb5d2..624da233 100644
--- a/src/Crypto/Tox.hs
+++ b/src/Crypto/Tox.hs
@@ -34,11 +34,13 @@ module Crypto.Tox
     , decodePlain
     -- , computeSharedSecret
     , lookupSharedSecret
+    , lookupNonceFunction
     , encrypt
     , decrypt
     , Nonce8(..)
     , Nonce24(..)
     , incrementNonce24
+    , nonce24ToWord16
     , addtoNonce24
     , Nonce32(..)
     , getRemainingEncrypted
@@ -300,10 +302,19 @@ lookupSharedSecret :: TransportCrypto -> SecretKey -> PublicKey -> Nonce24 -> IO
 lookupSharedSecret crypto sk recipient nonce
     = ($ nonce) <$> lookupNonceFunction crypto sk recipient
 
+{-# INLINE lookupNonceFunction #-}
 lookupNonceFunction :: TransportCrypto -> SecretKey -> PublicKey -> IO (Nonce24 -> State)
-lookupNonceFunction TransportCrypto{secretsCache} sk recipient = do
+lookupNonceFunction c@(TransportCrypto{secretsCache}) sk recipient = do
     now <- getPOSIXTime
-    atomically $ do
+    atomically $ lookupNonceFunctionSTM now c sk recipient
+
+{-# INLINE lookupNonceFunctionSTM #-}
+-- | This version of 'lookupNonceFunction' is STM instead of IO, this means if some later part of
+--   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
     mm <- readTVar $ sharedSecret secretsCache
     case MM.lookup' recipient mm of
         Nothing -> do
@@ -332,7 +343,10 @@ hsalsa20 k n = BA.append a b
 
 
 newtype Nonce24 = Nonce24 ByteString
- deriving (Eq, Ord, ByteArrayAccess,Data)
+ deriving (Eq, Ord, ByteArrayAccess, Data)
+
+nonce24ToWord16 :: Nonce24 -> Word16
+nonce24ToWord16 (Nonce24 n24) = fromIntegral (B.index n24 23) + 256 * fromIntegral (B.index n24 22)
 
 addtoNonce24 :: Nonce24 -> Word -> IO Nonce24
 addtoNonce24 (Nonce24 n24) n = Nonce24 <$> BA.copy n24 init
@@ -376,6 +390,7 @@ addtoNonce24 (Nonce24 n24) n = Nonce24 <$> BA.copy n24 init
             pokeElemOff ptr 5 $ tBE32 (W# sum_)
         init _ = error "incrementNonce24: I only support 64 and 32 bits"
 
+{-# INLINE incrementNonce24 #-}
 incrementNonce24 :: Nonce24 -> IO Nonce24
 incrementNonce24 nonce24 = addtoNonce24 nonce24 1