From b6c72f75230f7fbd4312eb9b63f6a40bf4cb3063 Mon Sep 17 00:00:00 2001
From: Han Xu <keepsimple@gmail.com>
Date: Sat, 30 Sep 2023 21:54:39 -0700
Subject: [PATCH] support changing auto_padding setting in CBC mode

---
 src/lib.rs   | 35 +++++++++++++++++++++++++++++++++--
 tests/aes.rs | 25 +++++++++++++++++++++++++
 2 files changed, 58 insertions(+), 2 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 8d4991a..60fb378 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -58,6 +58,9 @@ const RD_KEY_MAX_SIZE: usize = 4 * (AES_MAX_ROUNDS + 1);
 pub struct Cipher {
     encrypt_key: AesKey,
     decrypt_key: AesKey,
+
+    /// By default, this is true. It is only used in CBC mode now.
+    auto_padding: bool,
 }
 
 impl Cipher {
@@ -90,6 +93,7 @@ impl Cipher {
         Cipher {
             encrypt_key,
             decrypt_key,
+            auto_padding: true,
         }
     }
 
@@ -122,6 +126,7 @@ impl Cipher {
         Cipher {
             encrypt_key,
             decrypt_key,
+            auto_padding: true,
         }
     }
 
@@ -154,9 +159,22 @@ impl Cipher {
         Cipher {
             encrypt_key,
             decrypt_key,
+            auto_padding: true,
         }
     }
 
+    /// Changes the `auto_padding` setting of the cipher. Used in CBC mode.
+    /// Note that `auto_padding` is ignored in CFB mode.
+    pub fn set_auto_padding(&mut self, auto_padding: bool) {
+        self.auto_padding = auto_padding;
+    }
+
+    /// Returns the `auto_padding` setting in this cipher. If true,
+    /// the padding is done in PKCS7.
+    pub fn auto_padding(&self) -> bool {
+        self.auto_padding
+    }
+
     /// Encrypt in CBC mode.
     ///
     /// The input data is not modified. The output is a new Vec. Padding (PKCS7) is included.
@@ -164,6 +182,10 @@ impl Cipher {
     ///
     /// This method works for all key sizes.
     ///
+    /// Note: auto padding is enabled by default. If the cipher disabled auto padding, this
+    /// method will not add padding. If the input data length is not multiple of AES_BLOCK_SIZE,
+    /// the output will be an empty Vec (i.e. error).
+    ///
     /// # Examples
     ///
     /// ```
@@ -176,7 +198,14 @@ impl Cipher {
     /// let encrypted = cipher.cbc_encrypt(iv, plaintext);
     /// ```
     pub fn cbc_encrypt(&self, iv: &[u8], data: &[u8]) -> Vec<u8> {
-        let mut padded = pad(data);
+        let mut padded = if self.auto_padding {
+            pad(data)
+        } else if data.len() % AES_BLOCK_SIZE == 0 {
+            data.to_vec()
+        } else {
+            return Vec::new();
+        };
+
         let mut my_iv = iv;
         let mut rest = &mut padded[..];
         while rest.len() >= AES_BLOCK_SIZE {
@@ -228,7 +257,9 @@ impl Cipher {
             xor_with_iv(block, my_iv);
             my_iv = &data[i..i + AES_BLOCK_SIZE];
         }
-        unpad(&mut new);
+        if self.auto_padding {
+            unpad(&mut new);
+        }
         new
     }
 
diff --git a/tests/aes.rs b/tests/aes.rs
index 21700f8..5d7aa75 100644
--- a/tests/aes.rs
+++ b/tests/aes.rs
@@ -321,3 +321,28 @@ fn invalid_input_decrypt() {
     let decrypted = cipher.cbc_decrypt(iv, &empty_ciphertext[..]);
     assert!(decrypted.is_empty());
 }
+
+#[test]
+fn cbc_without_auto_padding() {
+    let plaintext: &[u8; 16] = b"\x59\xf9\x62\x18\xd8\xec\xca\xb2\x77\xed\x47\x7a\x33\xdc\xb7\xf3";
+    let ciphertext: &[u8; 16] = b"\x9f\x27\x07\xbc\x98\xbb\x57\x81\xd4\xe7\xb4\x61\xbf\xfe\x62\x70";
+    let key: &[u8; 16] = b"\x64\xc5\xfd\x55\xdd\x3a\xd9\x88\x32\x5b\xaa\xec\x52\x43\xdb\x98";
+    let iv: &[u8; 16] = b"\x00\x04\x00\x8c\x00\x16\x58\x00\x00\x00\x00\x00\x00\x00\x00\x00";
+
+    // create a cipher and disable its auto_padding.
+    let mut aes: Cipher = Cipher::new_128(key);
+    aes.set_auto_padding(false);
+
+    // verify decrypt.
+    let decrypted = aes.cbc_decrypt(iv, ciphertext);
+    assert_eq!(&decrypted, plaintext);
+
+    // verify encrypt.
+    let encrypted = aes.cbc_encrypt(iv, plaintext);
+    assert_eq!(&encrypted, ciphertext);
+
+    // encrypt returns empty if the plaintext is not valid block-sized.
+    let invalidblock: &[u8; 10] = b"\x59\xf9\x62\x18\xd8\xec\xca\xb2\x77\xed";
+    let encrypted = aes.cbc_encrypt(iv, invalidblock);
+    assert!(encrypted.is_empty());
+}