faster lazy secret_box

src/Simplex/FileTransfer/Client/Main.hs

@@ -231,18 +231,18 @@ cliSendFile SendOptions {filePath, outputDir, numRecipients, xftpServers, retryC
       let fileHdr = smpEncode FileHeader {fileName, fileExtra = Nothing}
           fileSize' = fromIntegral (B.length fileHdr) + fileSize
           chunkSizes = prepareChunkSizes $ fileSize' + fileSizeLen + cbAuthTagLen
-          paddedSize = fromIntegral $ sum chunkSizes
-      encrypt fileHdr key nonce fileSize' paddedSize encPath
+          encSize = fromIntegral $ sum chunkSizes
+      encrypt fileHdr key nonce fileSize' encSize encPath
       digest <- liftIO $ LC.sha512Hash <$> LB.readFile encPath
       let chunkSpecs = prepareChunkSpecs encPath chunkSizes
-          fd = FileDescription {size = FileSize paddedSize, digest = FileDigest digest, key, nonce, chunkSize = FileSize defaultChunkSize, chunks = []}
+          fd = FileDescription {size = FileSize encSize, digest = FileDigest digest, key, nonce, chunkSize = FileSize defaultChunkSize, chunks = []}
       pure (encPath, fd, chunkSpecs)
       where
         encrypt :: ByteString -> C.SbKey -> C.CbNonce -> Int64 -> Int64 -> FilePath -> ExceptT CLIError IO ()
-        encrypt fileHdr key nonce fileSize' paddedSize encFile = do
+        encrypt fileHdr key nonce fileSize' encSize encFile = do
           f <- liftIO $ LB.readFile filePath
           let f' = LB.fromStrict fileHdr <> f
-          c <- liftEither $ first (CLIError . show) $ LC.sbEncrypt key nonce f' fileSize' $ paddedSize - cbAuthTagLen
+          c <- liftEither $ first (CLIError . show) $ LC.sbEncrypt key nonce f' fileSize' $ encSize - cbAuthTagLen
           liftIO $ LB.writeFile encFile c
     -- let padSize = paddedSize - fileSize - fromIntegral (B.length fileHdr)
     -- when (padSize > 0) . LB.hPut h $ LB.replicate padSize '#'

src/Simplex/Messaging/Crypto/Lazy.hs

@@ -1,5 +1,8 @@
+{-# LANGUAGE BangPatterns #-}
 {-# LANGUAGE LambdaCase #-}
 {-# LANGUAGE PatternSynonyms #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TupleSections #-}
 
 module Simplex.Messaging.Crypto.Lazy
   ( sha512Hash,
@@ -18,10 +21,13 @@ import Crypto.Hash.Algorithms (SHA512)
 import qualified Crypto.MAC.Poly1305 as Poly1305
 import Data.ByteArray (ByteArrayAccess)
 import qualified Data.ByteArray as BA
+import qualified Data.ByteString as S
 import Data.ByteString.Char8 (ByteString)
 import qualified Data.ByteString.Char8 as B
 import qualified Data.ByteString.Lazy.Char8 as LB
+import qualified Data.ByteString.Lazy.Internal as LB
 import Data.Int (Int64)
+import Data.List.NonEmpty (NonEmpty (..))
 import Foreign (sizeOf)
 import Simplex.Messaging.Crypto (CbNonce, CryptoError (..), SbKey, pattern CbNonce, pattern SbKey)
 import Simplex.Messaging.Encoding
@@ -68,51 +74,64 @@ unPad padded
 
 -- | NaCl @secret_box@ lazy encrypt with a symmetric 256-bit key and 192-bit nonce.
 sbEncrypt :: SbKey -> CbNonce -> LazyByteString -> Int64 -> Int64 -> Either CryptoError LazyByteString
-sbEncrypt (SbKey key) = sbEncrypt_ key
-
-sbEncrypt_ :: ByteArrayAccess key => key -> CbNonce -> LazyByteString -> Int64 -> Int64 -> Either CryptoError LazyByteString
-sbEncrypt_ secret (CbNonce nonce) msg len paddedLen = cryptoBox secret nonce =<< pad msg len paddedLen
+sbEncrypt (SbKey key) (CbNonce nonce) msg len paddedLen =
+  prependTag <$> (secretBox sbEncryptChunk key nonce =<< pad msg len paddedLen)
+  where
+    prependTag (tag :| cs) = LB.Chunk tag $ LB.fromChunks cs
 
 -- | NaCl @secret_box@ decrypt with a symmetric 256-bit key and 192-bit nonce.
 sbDecrypt :: SbKey -> CbNonce -> LazyByteString -> Either CryptoError LazyByteString
-sbDecrypt (SbKey key) = sbDecrypt_ key
-
--- | NaCl @crypto_box@ decrypt with a shared DH secret and 192-bit nonce.
-sbDecrypt_ :: ByteArrayAccess key => key -> CbNonce -> LazyByteString -> Either CryptoError LazyByteString
-sbDecrypt_ secret (CbNonce nonce) packet
+sbDecrypt (SbKey key) (CbNonce nonce) packet
   | LB.length tag' < 16 = Left CBDecryptError
-  | otherwise = case poly1305auth rs c of
-    Right tag
-      | BA.constEq (LB.toStrict tag') tag -> unPad msg
+  | otherwise = case secretBox sbDecryptChunk key nonce c of
+    Right (tag :| cs)
+      | BA.constEq (LB.toStrict tag') tag -> unPad $ LB.fromChunks cs
       | otherwise -> Left CBDecryptError
     Left e -> Left e
   where
     (tag', c) = LB.splitAt 16 packet
-    (rs, msg) = xSalsa20 secret nonce c
 
-cryptoBox :: ByteArrayAccess key => key -> ByteString -> LazyByteString -> Either CryptoError LazyByteString
-cryptoBox secret nonce s = (<> c) . LB.fromStrict . BA.convert <$> tag
+secretBox :: ByteArrayAccess key => (SbState -> ByteString -> (ByteString, SbState)) -> key -> ByteString -> LazyByteString -> Either CryptoError (NonEmpty ByteString)
+secretBox sbProcess secret nonce msg = run <$> sbInit secret nonce
   where
-    (rs, c) = xSalsa20 secret nonce s
-    tag = poly1305auth rs c
+    process state = foldlChunks update ([], state) msg
+    update (cs, st) chunk = let (c, st') = sbProcess st chunk in (c : cs, st')
+    run state = let (cs, state') = process state in BA.convert (sbAuth state') :| reverse cs
 
-poly1305auth :: ByteString -> LazyByteString -> Either CryptoError Poly1305.Auth
-poly1305auth rs c = authTag <$> cryptoPassed (Poly1305.initialize rs)
-  where
-    authTag state = Poly1305.finalize $ Poly1305.updates state $ LB.toChunks c
-    cryptoPassed = \case
-      CE.CryptoPassed a -> Right a
-      CE.CryptoFailed e -> Left $ CryptoPoly1305Error e
+type SbState = (XSalsa.State, Poly1305.State)
 
-xSalsa20 :: ByteArrayAccess key => key -> ByteString -> LazyByteString -> (ByteString, LazyByteString)
-xSalsa20 secret nonce msg = (rs, msg')
+sbInit :: ByteArrayAccess key => key -> ByteString -> Either CryptoError SbState
+sbInit secret nonce = (state2,) <$> cryptoPassed (Poly1305.initialize rs)
   where
     zero = B.replicate 16 $ toEnum 0
     (iv0, iv1) = B.splitAt 8 nonce
     state0 = XSalsa.initialize 20 secret (zero `B.append` iv0)
     state1 = XSalsa.derive state0 iv1
-    (rs, state2) = XSalsa.generate state1 32
-    (msg', _) = foldl update (LB.empty, state2) $ LB.toChunks msg
-    update (acc, state) chunk =
-      let (c, state') = XSalsa.combine state chunk
-       in (acc `LB.append` LB.fromStrict c, state')
+    (rs :: ByteString, state2) = XSalsa.generate state1 32
+
+sbEncryptChunk :: SbState -> ByteString -> (ByteString, SbState)
+sbEncryptChunk (st, authSt) chunk =
+  let (c, st') = XSalsa.combine st chunk
+      authSt' = Poly1305.update authSt c
+   in (c, (st', authSt'))
+
+sbDecryptChunk :: SbState -> ByteString -> (ByteString, SbState)
+sbDecryptChunk (st, authSt) chunk =
+  let (s, st') = XSalsa.combine st chunk
+      authSt' = Poly1305.update authSt chunk
+   in (s, (st', authSt'))
+
+sbAuth :: SbState -> Poly1305.Auth
+sbAuth = Poly1305.finalize . snd
+
+cryptoPassed :: CE.CryptoFailable b -> Either CryptoError b
+cryptoPassed = \case
+  CE.CryptoPassed a -> Right a
+  CE.CryptoFailed e -> Left $ CryptoPoly1305Error e
+
+foldlChunks :: (a -> S.ByteString -> a) -> a -> LazyByteString -> a
+foldlChunks f = go
+  where
+    go !a LB.Empty = a
+    go !a (LB.Chunk c cs) = go (f a c) cs
+{-# INLINE foldlChunks #-}