Skip to content

Commit

Permalink
Expand API; improve test structure
Browse files Browse the repository at this point in the history
  • Loading branch information
@justsmth
justsmth committed Apr 30, 2024
1 parent 505ef99 commit 0159744
Showing 1 changed file with 114 additions and 51 deletions.
165 changes: 114 additions & 51 deletions aws-lc-rs/src/cipher/stream.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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);
}

Expand Down Expand Up @@ -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>](
Expand All @@ -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);
}
}
}
}

0 comments on commit 0159744

Please sign in to comment.