From dacaee60c418595c913e447faef61bfcc723faa6 Mon Sep 17 00:00:00 2001
From: Alexander Bondarenko <486682+dpwiz@users.noreply.github.com>
Date: Fri, 12 Apr 2024 16:09:57 +0300
Subject: [PATCH] add doubleStremaing case

---
 benchmarks/Bench/Crypto/Lazy.hs | 42 +++++++++++++++++++++++++++++----
 1 file changed, 37 insertions(+), 5 deletions(-)

diff --git a/benchmarks/Bench/Crypto/Lazy.hs b/benchmarks/Bench/Crypto/Lazy.hs
index 6226b9fd7..e5c71478d 100644
--- a/benchmarks/Bench/Crypto/Lazy.hs
+++ b/benchmarks/Bench/Crypto/Lazy.hs
@@ -1,4 +1,5 @@
 {-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE NamedFieldPuns #-}
 
 module Bench.Crypto.Lazy where
 
@@ -6,7 +7,7 @@ module Bench.Crypto.Lazy where
 import Test.Tasty.Bench
 
 import Control.Concurrent.STM (atomically)
-import Control.Monad.Except (runExceptT)
+import Control.Monad.Except (runExceptT, throwError)
 import qualified Data.ByteString.Lazy.Char8 as LB
 import qualified Simplex.Messaging.Crypto as C
 import Simplex.Messaging.Crypto.File (CryptoFile (..))
@@ -15,6 +16,10 @@ import System.Directory (removeFile)
 import Test.Tasty (TestTree, withResource)
 import System.IO (IOMode(..))
 import Control.Monad.IO.Class (liftIO)
+import Control.Monad
+import UnliftIO.Directory (getFileSize)
+import qualified Data.ByteString as B
+import Control.Monad.Trans.Except (ExceptT)
 
 benchCryptoLazy :: [Benchmark]
 benchCryptoLazy =
@@ -22,7 +27,8 @@ benchCryptoLazy =
       "File"
       [ withSomeFile $ bench "cf-readFile" . nfAppIO (>>= benchReadFile),
         withSomeFile $ bcompare "cf-readFile" . bench "cf-streamFromFile" . nfAppIO (>>= benchStreamFromFile),
-        withSomeFile $ bcompare "cf-readFile" . bench "cf-passthrough" . nfAppIO (>>= benchPassthrough)
+        withSomeFile $ bcompare "cf-readFile" . bench "cf-passthrough" . nfAppIO (>>= benchPassthrough),
+        withSomeFile $ bcompare "cf-readFile" . bench "cf-double-streaming" . nfAppIO (>>= benchDoubleStreaming)
       ]
   ]
 
@@ -39,6 +45,27 @@ benchStreamFromFile (cfIn, cfOut) = fmap (either (error . show) id) . runExceptT
 benchPassthrough :: (CryptoFile, CryptoFile) -> IO ()
 benchPassthrough (CryptoFile pathIn _, CryptoFile pathOut _) = LB.readFile pathIn >>= LB.writeFile pathOut
 
+benchDoubleStreaming :: (CryptoFile, CryptoFile) -> IO ()
+benchDoubleStreaming (fromCF@CryptoFile {filePath, cryptoArgs = fromArgs}, toCF@CryptoFile {cryptoArgs = cfArgs}) = do
+  fromSizeFull <- getFileSize filePath
+  let fromSize = fromSizeFull - maybe 0 (const $ toInteger C.authTagSize) fromArgs
+  fmap (either (error . show) id) . runExceptT $
+    CF.withFile fromCF ReadMode $ \fromH ->
+      CF.withFile toCF WriteMode $ \toH -> do
+        decryptChunks fromH fromSize (liftIO . CF.hPut toH . LB.fromStrict)
+        forM_ fromArgs $ \_ -> CF.hGetTag fromH
+        forM_ cfArgs $ \_ -> liftIO $ CF.hPutTag toH
+  where
+    decryptChunks :: CF.CryptoFileHandle -> Integer -> (B.ByteString -> ExceptT CF.FTCryptoError IO ()) -> ExceptT CF.FTCryptoError IO ()
+    decryptChunks r size f = do
+      let chSize = min size 0xFFFF
+          chSize' = fromIntegral chSize
+          size' = size - chSize
+      ch <- liftIO $ CF.hGet r chSize'
+      when (B.length ch /= chSize') $ throwError $ CF.FTCEFileIOError "encrypting file: unexpected EOF"
+      f ch
+      when (size' > 0) $ decryptChunks r size' f
+
 withSomeFile :: (IO (CryptoFile, CryptoFile) -> TestTree) -> TestTree
 withSomeFile = withResource createCF deleteCF
   where
@@ -49,12 +76,17 @@ withSomeFile = withResource createCF deleteCF
       -- let cfIn  = CryptoFile pathIn Nothing
       -- LB.writeFile pathIn $ LB.replicate (256 * 1024 * 1024) '#'
       cfIn <- atomically $ CryptoFile pathIn . Just <$> CF.randomArgs g
+      -- Right () <- runExceptT $ CF.withFile cfIn WriteMode $ \cbh -> liftIO $ do
+      --   replicateM_ 256 $ CF.hPut cbh dummyChunk
+      --   CF.hPutTag cbh
       Right () <- runExceptT $ CF.writeFile cfIn $ LB.replicate (256 * 1024 * 1024) '#'
       -- gen out args
       cfOut <- atomically $ CryptoFile "./some-file.out" . Just <$> CF.randomArgs g
       -- let cfOut  = CryptoFile "./some-file.out" Nothing
       pure (cfIn, cfOut)
     deleteCF (CryptoFile pathIn _, CryptoFile pathOut _) = do
-      -- removeFile pathIn
-      -- removeFile pathOut
-      pure ()
+      removeFile pathIn
+      removeFile pathOut
+
+dummyChunk :: LB.ByteString
+dummyChunk = LB.replicate (1024 * 1024) '#'