From b4482ac3f5daa18ee890b101c95b02f4eefaf698 Mon Sep 17 00:00:00 2001
From: Jean-Paul Calderone <exarkun@twistedmatrix.com>
Date: Fri, 29 Sep 2023 13:30:20 -0400
Subject: [PATCH] Un-hard-code the block size and related minor tweaks

---
 src/Tahoe/CHK/Encrypt.hs | 21 ++++++++++++++-------
 1 file changed, 14 insertions(+), 7 deletions(-)

diff --git a/src/Tahoe/CHK/Encrypt.hs b/src/Tahoe/CHK/Encrypt.hs
index bb83742..ce76a7d 100644
--- a/src/Tahoe/CHK/Encrypt.hs
+++ b/src/Tahoe/CHK/Encrypt.hs
@@ -1,7 +1,10 @@
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeApplications #-}
+
 -- | Support the encryption requirements of CHK.
 module Tahoe.CHK.Encrypt (encrypt, encryptLazy, decrypt, decryptLazy) where
 
-import Crypto.Cipher.Types (BlockCipher (ctrCombine), ivAdd, nullIV)
+import Crypto.Cipher.Types (BlockCipher (blockSize, ctrCombine), ivAdd, nullIV)
 import Data.ByteArray (ByteArray)
 import qualified Data.ByteString.Lazy as LBS
 import Data.List (unfoldr)
@@ -18,19 +21,23 @@ encrypt :: (BlockCipher cipher, ByteArray ba) => cipher -> ba -> ba
 encrypt key = ctrCombine key nullIV
 
 -- | Like encrypt but operate on lazy bytestrings.
-encryptLazy :: BlockCipher cipher => cipher -> LBS.ByteString -> LBS.ByteString
+encryptLazy :: forall cipher. BlockCipher cipher => cipher -> LBS.ByteString -> LBS.ByteString
 encryptLazy cipher lbs = LBS.concat . (LBS.fromStrict <$>) $ zipWith (ctrCombine cipher) ivs blocks
   where
     -- The underlying encryption function works on strict bytes.  Here's the
-    -- number of bytes to feed to it (that is, to make strict) at a time.
-    workingBlockSize :: Int
-    workingBlockSize = 1024 * 64
+    -- number of *blocks* to feed to it (that is, to make strict) at a time.
+    -- This value here is a magic number that is meant to represent a good
+    -- compromise between performance and number of bytes forced at one time.
+    workingBlocks = 1024 * 16
+
+    -- The size of a block is determined by the cipher.
+    workingBytes = workingBlocks * blockSize @cipher undefined
 
-    ivs = iterate (`ivAdd` (workingBlockSize `div` 16)) nullIV
+    ivs = iterate (`ivAdd` workingBlocks) nullIV
     blocks = LBS.toStrict <$> unfoldr takeChunk lbs
 
     takeChunk "" = Nothing
-    takeChunk xs = Just . LBS.splitAt (fromIntegral workingBlockSize) $ xs
+    takeChunk xs = Just . LBS.splitAt (fromIntegral workingBytes) $ xs
 
 -- | AES128-CTR decrypt a byte string in the manner used by CHK.
 decrypt :: (BlockCipher cipher, ByteArray ba) => cipher -> ba -> ba
-- 
GitLab