1 files changed, 21 insertions, 2 deletions

diff --git a/source/crypto/Aes.mc b/source/crypto/Aes.mc
index 242448b..9e4f320 100644
--- a/source/crypto/Aes.mc
+++ b/source/crypto/Aes.mc
@@ -42,19 +42,34 @@ module Aes {
 
     // Encrypt with PKCS#7 padding: 1 to 16 bytes, each holding the pad length.
     function cbcEncrypt(key as ByteArray, iv as ByteArray, plaintext as ByteArray) as ByteArray {
+        return cbcEncryptKeep(key, iv, plaintext, null);
+    }
+
+    // As cbcEncrypt(); with `leftovers`, the padded copy is added there for
+    // the caller to wipe (in steps, for kilobytes) instead of wiped here.
+    function cbcEncryptKeep(key as ByteArray, iv as ByteArray, plaintext as ByteArray, leftovers as Array<ByteArray>?) as ByteArray {
         var padLength = BLOCK_BYTES - (plaintext.size() % BLOCK_BYTES);
         var padded = Bytes.copy(plaintext);
         for (var i = 0; i < padLength; i++) {
             padded.add(padLength);
         }
         var out = cbcRaw(key, iv, padded, true);
-        Bytes.wipe(padded);
+        if (leftovers != null) {
+            leftovers.add(padded);
+        } else {
+            Bytes.wipe(padded);
+        }
         return out;
     }
 
     // Decrypt and remove PKCS#7 padding; null when the length or the padding is
     // wrong.
     function cbcDecrypt(key as ByteArray, iv as ByteArray, ciphertext as ByteArray) as ByteArray? {
+        return cbcDecryptKeep(key, iv, ciphertext, null);
+    }
+
+    // As cbcDecrypt(); with `leftovers`, the padded plaintext is added there.
+    function cbcDecryptKeep(key as ByteArray, iv as ByteArray, ciphertext as ByteArray, leftovers as Array<ByteArray>?) as ByteArray? {
         var size = ciphertext.size();
         if (size == 0 || size % BLOCK_BYTES != 0) {
             return null;
@@ -68,7 +83,11 @@ module Aes {
             }
         }
         var plaintext = valid ? padded.slice(0, size - padLength) : null;
-        Bytes.wipe(padded);
+        if (leftovers != null) {
+            leftovers.add(padded);
+        } else {
+            Bytes.wipe(padded);
+        }
         return plaintext;
     }
 }