1 //
2 // Copyright 2021 Signal Messenger, LLC.
3 // SPDX-License-Identifier: AGPL-3.0-only
4 //
5 
6 use crate::error::{Error, Result};
7 use aes::cipher::{FromBlockCipher, StreamCipher, StreamCipherSeek};
8 use aes::{Aes256, NewBlockCipher};
9 
10 /// A wrapper around [`aes::Aes256Ctr`] that uses a smaller nonce and supports an initial counter.
11 pub struct Aes256Ctr32(aes::Aes256Ctr);
12 
13 impl Aes256Ctr32 {
14     pub const NONCE_SIZE: usize = aes::BLOCK_SIZE - 4;
15 
new(aes256: Aes256, nonce: &[u8], init_ctr: u32) -> Result<Self>16     pub fn new(aes256: Aes256, nonce: &[u8], init_ctr: u32) -> Result<Self> {
17         if nonce.len() != Self::NONCE_SIZE {
18             return Err(Error::InvalidNonceSize);
19         }
20 
21         let mut nonce_block = [0u8; aes::BLOCK_SIZE];
22         nonce_block[0..Self::NONCE_SIZE].copy_from_slice(nonce);
23 
24         let mut ctr = aes::Aes256Ctr::from_block_cipher(aes256, &nonce_block.into());
25         ctr.seek((aes::BLOCK_SIZE as u64) * (init_ctr as u64));
26 
27         Ok(Self(ctr))
28     }
29 
from_key(key: &[u8], nonce: &[u8], init_ctr: u32) -> Result<Self>30     pub fn from_key(key: &[u8], nonce: &[u8], init_ctr: u32) -> Result<Self> {
31         Self::new(
32             Aes256::new_from_slice(key).map_err(|_| Error::InvalidKeySize)?,
33             nonce,
34             init_ctr,
35         )
36     }
37 
process(&mut self, buf: &mut [u8]) -> Result<()>38     pub fn process(&mut self, buf: &mut [u8]) -> Result<()> {
39         self.0.apply_keystream(buf);
40         Ok(())
41     }
42 }
43