Expand API; improve test structure

Author: justsmth

aws-lc-rs/src/cipher/stream.rs

@@ -58,15 +58,27 @@ impl StreamingEncryptingKey {
     }
 
     pub fn cbc_pkcs7(key: UnboundCipherKey) -> Result<Self, Unspecified> {
-        let mode = OperatingMode::CBC;
-        let context = key.algorithm().new_encryption_context(mode)?;
-        StreamingEncryptingKey::new(key, mode, context)
+        let context = key.algorithm().new_encryption_context(OperatingMode::CBC)?;
+        Self::less_safe_cbc_pkcs7(key, context)
+    }
+
+    pub fn less_safe_cbc_pkcs7(
+        key: UnboundCipherKey,
+        context: EncryptionContext,
+    ) -> Result<Self, Unspecified> {
+        StreamingEncryptingKey::new(key, OperatingMode::CBC, context)
     }
 
     pub fn ctr(key: UnboundCipherKey) -> Result<Self, Unspecified> {
-        let mode = OperatingMode::CTR;
-        let context = key.algorithm().new_encryption_context(mode)?;
-        StreamingEncryptingKey::new(key, mode, context)
+        let context = key.algorithm().new_encryption_context(OperatingMode::CTR)?;
+        Self::less_safe_ctr(key, context)
+    }
+
+    pub fn less_safe_ctr(
+        key: UnboundCipherKey,
+        context: EncryptionContext,
+    ) -> Result<Self, Unspecified> {
+        StreamingEncryptingKey::new(key, OperatingMode::CTR, context)
     }
 
     /// Returns the cipher operating mode.
@@ -81,7 +93,7 @@ impl StreamingEncryptingKey {
     }
 
     pub fn update<'a>(&self, input: &[u8], output: &'a mut [u8]) -> Result<usize, Unspecified> {
-        if output.len() < (input.len() + self.algorithm.block_len - 1) {
+        if output.len() < (input.len() + self.algorithm.block_len) {
             return Err(Unspecified);
         }
 
@@ -215,11 +227,86 @@ impl StreamingDecryptingKey {
 #[cfg(test)]
 mod tests {
     use crate::cipher::stream::{StreamingDecryptingKey, StreamingEncryptingKey};
-    use crate::cipher::{Algorithm, UnboundCipherKey, AES_256, AES_256_KEY_LEN};
+    use crate::cipher::{Algorithm, DecryptionContext, UnboundCipherKey, AES_256, AES_256_KEY_LEN};
     use crate::rand::{SecureRandom, SystemRandom};
     use paste::*;
 
-    macro_rules! helper_stream_step_test {
+    fn step_encrypt(
+        encrypting_key: StreamingEncryptingKey,
+        plaintext: &[u8],
+        step: usize,
+    ) -> (Box<[u8]>, DecryptionContext) {
+        let alg = encrypting_key.algorithm();
+        let mode = encrypting_key.mode();
+        let n = plaintext.len();
+        let mut ciphertext = vec![0u8; n + alg.block_len()];
+
+        let mut in_idx: usize = 0;
+        let mut out_idx: usize = 0;
+        loop {
+            let mut in_end = in_idx + step;
+            if in_end > n {
+                in_end = n;
+            }
+            let out_end = out_idx + (in_end - in_idx) + alg.block_len();
+            let outlen = encrypting_key
+                .update(
+                    &plaintext[in_idx..in_end],
+                    &mut ciphertext[out_idx..out_end],
+                )
+                .unwrap();
+            in_idx += step;
+            out_idx += outlen;
+            if in_idx >= n {
+                break;
+            }
+        }
+        let out_end = out_idx + alg.block_len();
+        let (decrypt_iv, outlen) = encrypting_key
+            .finish(&mut ciphertext[out_idx..out_end])
+            .unwrap();
+
+        ciphertext.truncate(out_idx + outlen);
+        (ciphertext.into_boxed_slice(), decrypt_iv)
+    }
+
+    fn step_decrypt(
+        decrypting_key: StreamingDecryptingKey,
+        ciphertext: &[u8],
+        step: usize,
+    ) -> Box<[u8]> {
+        let alg = decrypting_key.algorithm();
+        let n = ciphertext.len();
+        let mut output = vec![0u8; n + alg.block_len()];
+
+        let mut in_idx: usize = 0;
+        let mut out_idx: usize = 0;
+        loop {
+            let mut in_end = in_idx + step;
+            if in_end > n {
+                in_end = n;
+            }
+            let out_end = out_idx + (in_end - in_idx) + alg.block_len();
+            let outlen = decrypting_key
+                .update(&ciphertext[in_idx..in_end], &mut output[out_idx..out_end])
+                .unwrap();
+            in_idx += step;
+            out_idx += outlen;
+            if in_idx >= n {
+                break;
+            }
+        }
+        let out_end = out_idx + alg.block_len();
+        let outlen = decrypting_key
+            .finish(&mut output[out_idx..out_end])
+            .unwrap();
+
+        output.truncate(out_idx + outlen);
+
+        output.into_boxed_slice()
+    }
+
+    macro_rules! helper_stream_step_encrypt_test {
         ($mode:ident) => {
             paste! {
                 fn [<helper_test_ $mode _stream_encrypt_step_n_bytes>](
@@ -228,82 +315,58 @@ mod tests {
                     n: usize,
                     step: usize,
                 ) {
-                    let mut input: Vec<u8> = Vec::with_capacity(n);
-                    for i in 0..n {
-                        let byte: u8 = (i % 256).try_into().unwrap();
-                        input.push(byte);
-                    }
+                    let mut input = vec![0u8; n];
+                    let random = SystemRandom::new();
+                    random.fill(&mut input).unwrap();
 
                     let cipher_key = UnboundCipherKey::new(alg, key).unwrap();
                     let encrypting_key = StreamingEncryptingKey::$mode(cipher_key).unwrap();
 
-                    let mut ciphertext = vec![0u8; n + alg.block_len()];
-
-                    let mut in_idx: usize = 0;
-                    let mut out_idx: usize = 0;
-                    loop {
-                        let mut in_end = in_idx + step;
-                        if in_end > n {
-                            in_end = n;
-                        }
-                        let out_end = out_idx + (in_end - in_idx) + alg.block_len() - 1;
-                        let outlen = encrypting_key
-                            .update(&input[in_idx..in_end], &mut ciphertext[out_idx..out_end])
-                            .unwrap();
-                        in_idx += step;
-                        out_idx += outlen;
-                        if in_idx >= n {
-                            break;
-                        }
-                    }
-                    let out_end = out_idx + alg.block_len();
-                    let (decrypt_iv, outlen) = encrypting_key
-                        .finish(&mut ciphertext[out_idx..out_end])
-                        .unwrap();
-                    eprintln!("{} ciphertext length: n: {}, step: {}, cipherlength: {}", stringify!($mode), n, step, out_idx + outlen);
-                    ciphertext.truncate(out_idx + outlen);
+                    let (ciphertext, decrypt_iv) = step_encrypt(encrypting_key, &input, step);
 
                     let cipher_key2 = UnboundCipherKey::new(alg, key).unwrap();
                     let decrypting_key = StreamingDecryptingKey::$mode(cipher_key2, decrypt_iv).unwrap();
 
-                    let mut plaintext = vec![0u8; n + 2*alg.block_len()];
-                    let outlen = decrypting_key
-                        .update(ciphertext.as_slice(), plaintext.as_mut_slice())
-                        .unwrap();
-                    let outlen = outlen + decrypting_key.finish(&mut plaintext[outlen..n]).unwrap();
+                    let plaintext = step_decrypt(decrypting_key, &ciphertext, step);
 
-                    assert_eq!(input.as_slice(), &plaintext[0..outlen]);
+                    assert_eq!(input.as_slice(), &*plaintext);
                 }
             }
 
         };
     }
 
-    helper_stream_step_test!(cbc_pkcs7);
-    helper_stream_step_test!(ctr);
+    helper_stream_step_encrypt_test!(cbc_pkcs7);
+    helper_stream_step_encrypt_test!(ctr);
 
     #[test]
-    fn test_cbc() {
+    fn test_step_cbc() {
         let random = SystemRandom::new();
         let mut key = [0u8; AES_256_KEY_LEN];
         random.fill(&mut key).unwrap();
 
         for i in 13..=21 {
-            for j in 125..=131 {
+            for j in 124..=131 {
                 let _ = helper_test_cbc_pkcs7_stream_encrypt_step_n_bytes(&key, &AES_256, j, i);
             }
+            for j in 124..=131 {
+                let _ = helper_test_cbc_pkcs7_stream_encrypt_step_n_bytes(&key, &AES_256, j, j - i);
+            }
         }
     }
 
     #[test]
-    fn test_ctr() {
+    fn test_step_ctr() {
         let random = SystemRandom::new();
         let mut key = [0u8; AES_256_KEY_LEN];
         random.fill(&mut key).unwrap();
         for i in 13..=21 {
-            for j in 125..=131 {
+            for j in 124..=131 {
                 let _ = helper_test_ctr_stream_encrypt_step_n_bytes(&key, &AES_256, j, i);
             }
+            for j in 124..=131 {
+                let _ = helper_test_ctr_stream_encrypt_step_n_bytes(&key, &AES_256, j, j - i);
+            }
         }
     }
 }