Go back to lazy town

src/Tahoe/CHK/Encrypt.hs

 -- | Support the encryption requirements of CHK.
 module Tahoe.CHK.Encrypt (encrypt, encryptLazy, decrypt, decryptLazy) where
 
-import Crypto.Cipher.Types (BlockCipher (ctrCombine), nullIV)
+import Crypto.Cipher.Types (BlockCipher (ctrCombine), ivAdd, nullIV)
 import Data.ByteArray (ByteArray)
 import qualified Data.ByteString.Lazy as LBS
+import Data.List (unfoldr)
 
 {- | CTR-mode encrypt a byte string using some block cipher.
 
@@ -16,18 +17,27 @@ import qualified Data.ByteString.Lazy as LBS
 encrypt :: (BlockCipher cipher, ByteArray ba) => cipher -> ba -> ba
 encrypt key = ctrCombine key nullIV
 
-{- | Like encrypt but operate on lazy bytestrings.  TODO: Make this more
- efficient than converting to/from strict ByteString!
--}
+-- | Like encrypt but operate on lazy bytestrings.
 encryptLazy :: BlockCipher cipher => cipher -> LBS.ByteString -> LBS.ByteString
-encryptLazy cipher lbs = LBS.fromStrict (encrypt cipher (LBS.toStrict lbs))
+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
+
+    ivs = iterate (`ivAdd` (workingBlockSize `div` 16)) nullIV
+    blocks = LBS.toStrict <$> unfoldr takeChunk lbs
+
+    takeChunk "" = Nothing
+    takeChunk xs = Just . LBS.splitAt (fromIntegral workingBlockSize) $ xs
+
+-- LBS.fromStrict (encrypt cipher (LBS.toStrict lbs))
 
 -- | AES128-CTR decrypt a byte string in the manner used by CHK.
 decrypt :: (BlockCipher cipher, ByteArray ba) => cipher -> ba -> ba
 decrypt = encrypt
 
-{- | Like decrypt but operate on lazy bytestrings.  TODO: Make this more
- efficient than converting to/from strict ByteString!
--}
+-- | Like decrypt but operate on lazy bytestrings.
 decryptLazy :: BlockCipher cipher => cipher -> LBS.ByteString -> LBS.ByteString
 decryptLazy = encryptLazy