rns-crypto 0.2.2

use alloc::vec::Vec;
use core::fmt;

use crate::aes128::Aes128;
use crate::aes256::Aes256;
use aes::cipher::{array::Array, BlockModeDecrypt};

use crate::hmac::HmacSha256;
use crate::pkcs7;
use crate::sha256::Sha256;
use crate::Rng;

pub const TOKEN_OVERHEAD: usize = 48; // 16 IV + 32 HMAC

#[derive(Debug, PartialEq)]
pub enum TokenError {
    InvalidKeyLength,
    InvalidToken,
    HmacMismatch,
    DecryptionFailed,
}

impl fmt::Display for TokenError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            TokenError::InvalidKeyLength => write!(f, "Token key must be 32 or 64 bytes"),
            TokenError::InvalidToken => write!(f, "Token too short"),
            TokenError::HmacMismatch => write!(f, "Token HMAC was invalid"),
            TokenError::DecryptionFailed => write!(f, "Could not decrypt token"),
        }
    }
}

#[derive(Clone)]
enum AesMode {
    Aes128(Aes128),
    Aes256(Aes256),
}

/// CBC chaining state carried across the pieces of one ciphertext. It lives
/// on the stack for one call, so the variants' size difference is harmless.
#[allow(clippy::large_enum_variant)]
enum CbcDecryptor {
    Aes128(cbc::Decryptor<aes::Aes128>),
    Aes256(cbc::Decryptor<aes::Aes256>),
}

impl CbcDecryptor {
    /// Decrypt whole blocks in place, continuing the chain.
    fn decrypt_blocks(&mut self, buf: &mut [u8]) {
        let (blocks, rest) = Array::slice_as_chunks_mut(buf);
        debug_assert!(rest.is_empty());
        match self {
            CbcDecryptor::Aes128(dec) => dec.decrypt_blocks(blocks),
            CbcDecryptor::Aes256(dec) => dec.decrypt_blocks(blocks),
        }
    }
}

pub struct Token {
    /// HMAC keyed once with the signing key: each tag clones it instead of
    /// hashing the key's inner and outer pads again.
    mac: HmacSha256,
    mode: AesMode,
}

impl Token {
    pub fn new(key: &[u8]) -> Result<Self, TokenError> {
        match key.len() {
            32 => {
                let signing_key = &key[..16];
                let encryption_key: [u8; 16] = key[16..32].try_into().unwrap();
                Ok(Token {
                    mac: HmacSha256::new(signing_key),
                    mode: AesMode::Aes128(Aes128::new(&encryption_key)),
                })
            }
            64 => {
                let signing_key = &key[..32];
                let encryption_key: [u8; 32] = key[32..64].try_into().unwrap();
                Ok(Token {
                    mac: HmacSha256::new(signing_key),
                    mode: AesMode::Aes256(Aes256::new(&encryption_key)),
                })
            }
            _ => Err(TokenError::InvalidKeyLength),
        }
    }

    pub fn encrypt(&self, plaintext: &[u8], rng: &mut dyn Rng) -> Vec<u8> {
        let mut iv = [0u8; 16];
        rng.fill_bytes(&mut iv);
        self.encrypt_with_iv(plaintext, &iv)
    }

    /// Length of the token [`Token::encrypt`] produces for a plaintext of
    /// `plaintext_len` bytes.
    pub fn token_len(plaintext_len: usize) -> usize {
        16 + plaintext_len + (pkcs7::BLOCK_SIZE - plaintext_len % pkcs7::BLOCK_SIZE) + 32
    }

    /// [`Token::encrypt`] into `out`, which must be exactly
    /// [`Token::token_len`]`(plaintext.len())` bytes long.
    pub fn encrypt_into(&self, plaintext: &[u8], rng: &mut dyn Rng, out: &mut [u8]) {
        let mut iv = [0u8; 16];
        rng.fill_bytes(&mut iv);
        self.encrypt_with_iv_into(plaintext, &iv, out)
    }

    /// [`Token::encrypt_with_iv`] into `out`, which must be exactly
    /// [`Token::token_len`]`(plaintext.len())` bytes long.
    pub fn encrypt_with_iv_into(&self, plaintext: &[u8], iv: &[u8; 16], out: &mut [u8]) {
        assert_eq!(
            out.len(),
            Self::token_len(plaintext.len()),
            "token buffer length"
        );
        let (body, tag) = out.split_at_mut(out.len() - 32);
        let padding = body.len() - 16 - plaintext.len();
        body[..16].copy_from_slice(iv);
        body[16..16 + plaintext.len()].copy_from_slice(plaintext);
        body[16 + plaintext.len()..].fill(padding as u8);
        match &self.mode {
            AesMode::Aes128(aes) => aes.encrypt_cbc_in_place(&mut body[16..], iv),
            AesMode::Aes256(aes) => aes.encrypt_cbc_in_place(&mut body[16..], iv),
        }
        .expect("PKCS7 padding always produces block-aligned plaintext");
        tag.copy_from_slice(&self.tag(body));
    }

    pub fn encrypt_with_iv(&self, plaintext: &[u8], iv: &[u8; 16]) -> Vec<u8> {
        // Build the whole token in one buffer: the plaintext is copied once,
        // padded and encrypted in place, then signed and the tag appended.
        let padding = pkcs7::BLOCK_SIZE - plaintext.len() % pkcs7::BLOCK_SIZE;
        let mut token = Vec::with_capacity(16 + plaintext.len() + padding + 32);
        token.extend_from_slice(iv);
        token.extend_from_slice(plaintext);
        token.resize(token.len() + padding, padding as u8);
        match &self.mode {
            AesMode::Aes128(aes) => aes.encrypt_cbc_in_place(&mut token[16..], iv),
            AesMode::Aes256(aes) => aes.encrypt_cbc_in_place(&mut token[16..], iv),
        }
        .expect("PKCS7 padding always produces block-aligned plaintext");

        let mac = self.tag(&token);
        token.extend_from_slice(&mac);
        token
    }

    /// [`Token::encrypt`] with the tag left to the caller: the token's last
    /// 32 bytes are zero, and the returned MAC, fed every byte before them,
    /// finalizes to the tag. The caller can hash those bytes together with
    /// other data before writing it.
    pub fn encrypt_untagged(&self, plaintext: &[u8], rng: &mut dyn Rng) -> (Vec<u8>, HmacSha256) {
        let mut iv = [0u8; 16];
        rng.fill_bytes(&mut iv);
        let mut token = alloc::vec![0u8; Self::token_len(plaintext.len())];
        let signed = token.len() - 32;
        let body = &mut token[..signed];
        let padding = body.len() - 16 - plaintext.len();
        body[..16].copy_from_slice(&iv);
        body[16..16 + plaintext.len()].copy_from_slice(plaintext);
        body[16 + plaintext.len()..].fill(padding as u8);
        match &self.mode {
            AesMode::Aes128(aes) => aes.encrypt_cbc_in_place(&mut body[16..], &iv),
            AesMode::Aes256(aes) => aes.encrypt_cbc_in_place(&mut body[16..], &iv),
        }
        .expect("PKCS7 padding always produces block-aligned plaintext");
        (token, self.mac.clone())
    }

    /// HMAC-SHA256 of `data` under the signing key.
    fn tag(&self, data: &[u8]) -> [u8; 32] {
        let mut mac = self.mac.clone();
        mac.update(data);
        mac.finalize()
    }

    pub fn verify_hmac(&self, token: &[u8]) -> Result<bool, TokenError> {
        if token.len() <= 32 {
            return Err(TokenError::InvalidToken);
        }
        let received_hmac = &token[token.len() - 32..];
        let mut mac = self.mac.clone();
        mac.update(&token[..token.len() - 32]);
        Ok(mac.verify(received_hmac))
    }

    /// [`Token::decrypt`] into `out`, replacing its contents and reusing
    /// its capacity. On error `out` holds no plaintext.
    pub fn decrypt_into(&self, token: &[u8], out: &mut Vec<u8>) -> Result<(), TokenError> {
        out.clear();
        if token.len() <= TOKEN_OVERHEAD {
            return Err(TokenError::InvalidToken);
        }
        if !self
            .verify_hmac(token)
            .map_err(|_| TokenError::InvalidToken)?
        {
            return Err(TokenError::HmacMismatch);
        }
        let iv: [u8; 16] = token[..16].try_into().unwrap();
        let ciphertext = &token[16..token.len() - 32];
        if ciphertext.is_empty() || !ciphertext.len().is_multiple_of(16) {
            return Err(TokenError::DecryptionFailed);
        }
        out.extend_from_slice(ciphertext);
        let decrypted = match &self.mode {
            AesMode::Aes128(aes) => aes.decrypt_cbc_in_place(out, &iv),
            AesMode::Aes256(aes) => aes.decrypt_cbc_in_place(out, &iv),
        };
        let unpadded_len = decrypted
            .ok()
            .and_then(|()| pkcs7::unpad(out, 16).ok().map(<[u8]>::len));
        match unpadded_len {
            Some(len) => {
                out.truncate(len);
                Ok(())
            }
            None => {
                out.clear();
                Err(TokenError::DecryptionFailed)
            }
        }
    }

    pub fn decrypt(&self, token: &[u8]) -> Result<Vec<u8>, TokenError> {
        if token.len() <= TOKEN_OVERHEAD {
            return Err(TokenError::InvalidToken);
        }

        if !self
            .verify_hmac(token)
            .map_err(|_| TokenError::InvalidToken)?
        {
            return Err(TokenError::HmacMismatch);
        }

        let iv: [u8; 16] = token[..16].try_into().unwrap();
        let ciphertext = &token[16..token.len() - 32];

        // CBC ciphertext must be a non-empty whole number of 16-byte blocks.
        if ciphertext.is_empty() || !ciphertext.len().is_multiple_of(16) {
            return Err(TokenError::DecryptionFailed);
        }

        let mut decrypted = match &self.mode {
            AesMode::Aes128(aes) => aes.decrypt_cbc(ciphertext, &iv),
            AesMode::Aes256(aes) => aes.decrypt_cbc(ciphertext, &iv),
        }
        .map_err(|_| TokenError::DecryptionFailed)?;

        let unpadded_len = pkcs7::unpad(&decrypted, 16)
            .map_err(|_| TokenError::DecryptionFailed)?
            .len();
        decrypted.truncate(unpadded_len);
        Ok(decrypted)
    }

    /// Decrypt a token that arrives as consecutive chunks, such as the parts
    /// of a resource, without joining them first. Each chunk is authenticated
    /// and decrypted while it is still in cache; nothing is returned unless
    /// the tag matches. The result is [`Token::decrypt`] of the joined bytes
    /// with its first `skip` bytes (at most one block) dropped, and the errors
    /// are the same, except that a plaintext shorter than `skip` is
    /// [`TokenError::DecryptionFailed`].
    pub fn decrypt_chunks(&self, chunks: &[&[u8]], skip: usize) -> Result<Vec<u8>, TokenError> {
        let total: usize = chunks.iter().map(|chunk| chunk.len()).sum();
        let mut stream = self.decrypt_stream(skip);
        stream.reserve(total.saturating_sub(TOKEN_OVERHEAD));
        for chunk in chunks {
            stream.update(chunk);
        }
        stream.finish()
    }

    /// Start a [`Token::decrypt_chunks`] whose chunks are fed one at a time,
    /// as they become available.
    pub fn decrypt_stream(&self, skip: usize) -> TokenDecryptStream {
        assert!(skip <= 16, "decrypt_chunks skips at most one block");
        TokenDecryptStream {
            mac: self.mac.clone(),
            mode: self.mode.clone(),
            decryptor: None,
            skip,
            signed: 0,
            iv: [0; 16],
            tail: [0; 32],
            tail_len: 0,
            block: [0; 16],
            block_len: 0,
            first_block: true,
            last_byte: 0,
            plaintext: Vec::new(),
            plain_hash: None,
        }
    }
}

/// An incremental [`Token::decrypt_chunks`]: the token's bytes are fed in
/// order with [`TokenDecryptStream::update`], and [`TokenDecryptStream::finish`]
/// checks the tag and returns the plaintext. The last 32 bytes seen are held
/// back as the candidate tag; everything before them is authenticated and
/// decrypted as it arrives. Decrypted bytes are not authentic until `finish`
/// accepts the tag.
pub struct TokenDecryptStream {
    mac: HmacSha256,
    mode: AesMode,
    decryptor: Option<CbcDecryptor>,
    skip: usize,
    /// Bytes released to the MAC: the IV and the ciphertext so far.
    signed: usize,
    iv: [u8; 16],
    /// The last bytes seen, up to 32: the tag if the token ends here.
    tail: [u8; 32],
    tail_len: usize,
    /// A ciphertext block that straddles updates, and always the first
    /// block, which carries the skipped bytes.
    block: [u8; 16],
    block_len: usize,
    first_block: bool,
    last_byte: u8,
    plaintext: Vec<u8>,
    /// SHA-256 of the settled plaintext, and how many bytes it covers.
    plain_hash: Option<(Sha256, usize)>,
}

impl TokenDecryptStream {
    /// Also hash the plaintext, together with the tag's MAC, as it settles;
    /// [`Self::finish_hashed`] returns the hash. Call before any update.
    pub fn hash_plaintext(&mut self) {
        debug_assert!(self.signed == 0 && self.tail_len == 0);
        self.plain_hash = Some((Sha256::new(), 0));
    }

    /// Reserve room for `additional` more plaintext bytes.
    pub fn reserve(&mut self, additional: usize) {
        self.plaintext.reserve(additional);
    }

    /// The plaintext decrypted so far, without the skipped bytes. Its last
    /// block may still hold the padding, and none of it is authentic until
    /// [`Self::finish`] succeeds.
    pub fn plaintext(&self) -> &[u8] {
        &self.plaintext
    }

    /// Feed the token's next bytes.
    pub fn update(&mut self, data: &[u8]) {
        let total = self.tail_len + data.len();
        if total <= 32 {
            self.tail[self.tail_len..total].copy_from_slice(data);
            self.tail_len = total;
            return;
        }
        // Everything but the last 32 bytes is signed data.
        let release = total - 32;
        let from_tail = release.min(self.tail_len);
        if from_tail > 0 {
            let tail = self.tail;
            self.absorb_signed(&tail[..from_tail]);
        }
        let from_data = release - from_tail;
        self.absorb_signed(&data[..from_data]);
        let kept = self.tail_len - from_tail;
        self.tail.copy_within(from_tail..self.tail_len, 0);
        self.tail[kept..].copy_from_slice(&data[from_data..]);
        self.tail_len = 32;
    }

    fn absorb_signed(&mut self, data: &[u8]) {
        if data.is_empty() {
            return;
        }
        self.decrypt_signed(data);
        match self.plain_hash.as_mut() {
            Some((hash, hashed)) => {
                // The last block may still lose its padding.
                let settled = self.plaintext.len().saturating_sub(16).max(*hashed);
                self.mac
                    .update_paired(data, hash, &self.plaintext[*hashed..settled]);
                *hashed = settled;
            }
            None => self.mac.update(data),
        }
    }

    fn decrypt_signed(&mut self, mut data: &[u8]) {
        if self.signed < 16 {
            let n = (16 - self.signed).min(data.len());
            self.iv[self.signed..self.signed + n].copy_from_slice(&data[..n]);
            self.signed += n;
            data = &data[n..];
            if data.is_empty() {
                return;
            }
        }
        self.signed += data.len();
        let (mode, iv) = (&self.mode, &self.iv);
        let decryptor = self.decryptor.get_or_insert_with(|| match mode {
            AesMode::Aes128(aes) => CbcDecryptor::Aes128(aes.cbc_decryptor(iv)),
            AesMode::Aes256(aes) => CbcDecryptor::Aes256(aes.cbc_decryptor(iv)),
        });
        while !data.is_empty() && (self.block_len > 0 || self.first_block) {
            let n = (16 - self.block_len).min(data.len());
            self.block[self.block_len..self.block_len + n].copy_from_slice(&data[..n]);
            self.block_len += n;
            data = &data[n..];
            if self.block_len == 16 {
                decryptor.decrypt_blocks(&mut self.block);
                let from = if self.first_block { self.skip } else { 0 };
                self.plaintext.extend_from_slice(&self.block[from..]);
                self.last_byte = self.block[15];
                self.first_block = false;
                self.block_len = 0;
            }
        }
        if !data.is_empty() {
            let whole = data.len() - data.len() % 16;
            if whole > 0 {
                let at = self.plaintext.len();
                self.plaintext.extend_from_slice(&data[..whole]);
                decryptor.decrypt_blocks(&mut self.plaintext[at..]);
                self.last_byte = self.plaintext[self.plaintext.len() - 1];
            }
            let rest = &data[whole..];
            self.block[..rest.len()].copy_from_slice(rest);
            self.block_len = rest.len();
        }
    }

    /// [`Self::finish`], with the SHA-256 state over the whole plaintext
    /// when [`Self::hash_plaintext`] was called.
    pub fn finish_hashed(mut self) -> Result<(Vec<u8>, Option<Sha256>), TokenError> {
        let plain_hash = self.plain_hash.take();
        let plaintext = self.finish()?;
        Ok(match plain_hash {
            Some((mut hash, hashed)) => {
                hash.update(&plaintext[hashed..]);
                (plaintext, Some(hash))
            }
            None => (plaintext, None),
        })
    }

    /// Check the tag and return the plaintext, as [`Token::decrypt_chunks`]
    /// of every byte fed.
    pub fn finish(self) -> Result<Vec<u8>, TokenError> {
        if self.signed + self.tail_len <= TOKEN_OVERHEAD {
            return Err(TokenError::InvalidToken);
        }
        if !self.mac.verify(&self.tail) {
            return Err(TokenError::HmacMismatch);
        }
        // A misaligned ciphertext is rejected after the tag check, as in
        // `decrypt`.
        if self.first_block || self.block_len != 0 {
            return Err(TokenError::DecryptionFailed);
        }
        // The same padding check as `pkcs7::unpad`.
        let mut plaintext = self.plaintext;
        let padding = usize::from(self.last_byte);
        if padding > 16 || padding > plaintext.len() {
            return Err(TokenError::DecryptionFailed);
        }
        plaintext.truncate(plaintext.len() - padding);
        Ok(plaintext)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::hmac::hmac_sha256;

    #[test]
    fn decrypting_into_a_reused_buffer_matches_decrypt() {
        for key in [[7u8; 32].as_slice(), [9u8; 64].as_slice()] {
            let token = Token::new(key).unwrap();
            let other = Token::new(&[1u8; 32]).unwrap();
            let mut out = vec![0x55; 3];
            for len in [0usize, 1, 15, 16, 17, 100, 384] {
                let plaintext: Vec<u8> = (0..len).map(|i| i as u8).collect();
                let good = token.encrypt_with_iv(&plaintext, &[len as u8; 16]);
                let mut tampered = good.clone();
                tampered[20] ^= 1;
                let foreign = other.encrypt_with_iv(&plaintext, &[0; 16]);
                for input in [&good[..], &tampered[..], &foreign[..], &good[..40], &[][..]] {
                    let expected = token.decrypt(input);
                    let got = token.decrypt_into(input, &mut out);
                    assert_eq!(
                        format!("{:?}", got.map(|()| out.clone())),
                        format!("{expected:?}")
                    );
                    if expected.is_err() {
                        assert!(out.is_empty());
                    }
                }
            }
        }
    }

    #[test]
    fn encrypting_in_place_matches_encrypting_into_a_new_buffer() {
        for key in [[7u8; 32].as_slice(), [9u8; 64].as_slice()] {
            let token = Token::new(key).unwrap();
            for len in [0usize, 1, 15, 16, 17, 31, 32, 33, 100, 383, 384] {
                let plaintext: Vec<u8> = (0..len).map(|i| i as u8).collect();
                let iv = [len as u8; 16];
                let mut out = vec![0xaa; Token::token_len(len)];
                token.encrypt_with_iv_into(&plaintext, &iv, &mut out);
                assert_eq!(out, token.encrypt_with_iv(&plaintext, &iv));
                assert_eq!(token.decrypt(&out).unwrap(), plaintext);
            }
        }
    }
    use crate::FixedRng;

    #[test]
    fn test_token_new_32byte_key() {
        let key = [0u8; 32];
        assert!(Token::new(&key).is_ok());
    }

    #[test]
    fn test_token_new_64byte_key() {
        let key = [0u8; 64];
        assert!(Token::new(&key).is_ok());
    }

    #[test]
    fn test_token_new_48byte_key() {
        let key = [0u8; 48];
        assert!(matches!(
            Token::new(&key),
            Err(TokenError::InvalidKeyLength)
        ));
    }

    #[test]
    fn test_token_roundtrip_32() {
        let key = [0x42u8; 32];
        let token = Token::new(&key).unwrap();
        let mut rng = FixedRng::new(&[0xAA; 16]);
        let plaintext = b"Hello, Reticulum!";
        let encrypted = token.encrypt(plaintext, &mut rng);
        let decrypted = token.decrypt(&encrypted).unwrap();
        assert_eq!(decrypted, plaintext);
    }

    #[test]
    fn test_token_roundtrip_64() {
        let key = [0x42u8; 64];
        let token = Token::new(&key).unwrap();
        let mut rng = FixedRng::new(&[0xBB; 16]);
        let plaintext = b"Hello, Reticulum!";
        let encrypted = token.encrypt(plaintext, &mut rng);
        let decrypted = token.decrypt(&encrypted).unwrap();
        assert_eq!(decrypted, plaintext);
    }

    #[test]
    fn test_token_hmac_reject_tampered() {
        let key = [0x42u8; 64];
        let token = Token::new(&key).unwrap();
        let mut rng = FixedRng::new(&[0xCC; 16]);
        let encrypted = token.encrypt(b"test", &mut rng);
        let mut tampered = encrypted.clone();
        tampered[20] ^= 0xFF; // flip a bit in ciphertext
        assert!(token.decrypt(&tampered).is_err());
    }

    #[test]
    fn token_hmac_rejects_mismatches_across_the_entire_tag() {
        let key = [0x42u8; 64];
        let token = Token::new(&key).unwrap();
        let mut rng = FixedRng::new(&[0xCC; 16]);
        let encrypted = token.encrypt(b"timing regression", &mut rng);

        assert_eq!(token.verify_hmac(&encrypted), Ok(true));
        for tag_index in 0..32 {
            let mut corrupted = encrypted.clone();
            let index = corrupted.len() - 32 + tag_index;
            corrupted[index] ^= 0x01;
            assert_eq!(
                token.verify_hmac(&corrupted),
                Ok(false),
                "tag mismatch at byte {tag_index} must fail"
            );
        }
    }

    #[test]
    fn test_token_decrypt_truncated() {
        let key = [0x42u8; 64];
        let token = Token::new(&key).unwrap();
        assert!(matches!(
            token.decrypt(&[0u8; 10]),
            Err(TokenError::InvalidToken)
        ));
    }

    fn sign_token_parts(key: &[u8], iv_and_ciphertext: &[u8]) -> Vec<u8> {
        let signing_key = &key[..key.len() / 2];
        let mut token = iv_and_ciphertext.to_vec();
        token.extend_from_slice(&hmac_sha256(signing_key, iv_and_ciphertext));
        token
    }

    #[test]
    fn test_token_decrypt_rejects_misaligned_ciphertext_with_valid_hmac() {
        for key in [&[0x42u8; 32][..], &[0x42u8; 64][..]] {
            let token = Token::new(key).unwrap();
            for ciphertext_len in [1usize, 15, 17, 31, 33] {
                let mut parts = alloc::vec![0xAAu8; 16];
                parts.extend(core::iter::repeat_n(0x55u8, ciphertext_len));
                let signed = sign_token_parts(key, &parts);
                assert_eq!(token.verify_hmac(&signed), Ok(true));
                assert_eq!(
                    token.decrypt(&signed),
                    Err(TokenError::DecryptionFailed),
                    "ciphertext length {ciphertext_len} must be rejected"
                );
            }
        }
    }

    #[test]
    fn test_token_well_formed_roundtrip_after_alignment_check() {
        for key in [&[0x24u8; 32][..], &[0x24u8; 64][..]] {
            let token = Token::new(key).unwrap();
            for plaintext_len in [0usize, 1, 15, 16, 17, 32, 100] {
                let plaintext = alloc::vec![0x5Au8; plaintext_len];
                let encrypted = token.encrypt_with_iv(&plaintext, &[0x11; 16]);
                assert_eq!((encrypted.len() - TOKEN_OVERHEAD) % 16, 0);
                assert_eq!(token.decrypt(&encrypted).unwrap(), plaintext);
            }
        }
    }

    #[test]
    fn test_token_overhead() {
        assert_eq!(TOKEN_OVERHEAD, 48);
    }

    /// Split `bytes` at the given offsets (clamped, in order).
    fn split_at_offsets<'a>(bytes: &'a [u8], offsets: &[usize]) -> Vec<&'a [u8]> {
        let mut chunks = Vec::new();
        let mut from = 0;
        for &offset in offsets {
            let to = offset.clamp(from, bytes.len());
            chunks.push(&bytes[from..to]);
            from = to;
        }
        chunks.push(&bytes[from..]);
        chunks
    }

    fn chunked_as_joined(token: &Token, bytes: &[u8], chunks: &[&[u8]], skip: usize) {
        let joined = token.decrypt(bytes);
        let expected = match joined {
            Ok(plaintext) if plaintext.len() >= skip => Ok(plaintext[skip..].to_vec()),
            Ok(_) => Err(TokenError::DecryptionFailed),
            Err(error) => Err(error),
        };
        assert_eq!(
            token.decrypt_chunks(chunks, skip),
            expected,
            "{chunks:?} skip {skip}"
        );
    }

    #[test]
    fn chunked_decryption_matches_joined_decryption() {
        let mut rng = FixedRng::new(&[0x5D; 64]);
        for key in [[0x21u8; 32].to_vec(), [0x42u8; 64].to_vec()] {
            let token = Token::new(&key).unwrap();
            for len in [0usize, 1, 3, 4, 5, 12, 15, 16, 17, 33, 100, 1000] {
                let plaintext: Vec<u8> = (0..len).map(|i| (i * 7) as u8).collect();
                let bytes = token.encrypt(&plaintext, &mut rng);
                let n = bytes.len();
                let splits: Vec<Vec<usize>> = alloc::vec![
                    alloc::vec![],
                    alloc::vec![1],
                    alloc::vec![8, 16, 17],
                    alloc::vec![16, 20, 32, 47],
                    alloc::vec![n - 33, n - 32, n - 31],
                    alloc::vec![n - 40, n - 10],
                    (1..n).collect(),
                    (0..n).step_by(13).collect(),
                    alloc::vec![0, 0, 5, 5, n],
                ];
                for offsets in &splits {
                    let chunks = split_at_offsets(&bytes, offsets);
                    for skip in [0, 1, 4, 16] {
                        chunked_as_joined(&token, &bytes, &chunks, skip);
                    }
                }
                // Tampering anywhere is caught.
                for at in 0..n {
                    let mut tampered = bytes.clone();
                    tampered[at] ^= 0x01;
                    let chunks = split_at_offsets(&tampered, &[at / 2, at, at + 1]);
                    assert_eq!(
                        token.decrypt_chunks(&chunks, 4),
                        Err(TokenError::HmacMismatch)
                    );
                }
            }
            // Misaligned ciphertext under a valid tag, and short tokens.
            for ciphertext_len in [1usize, 15, 17, 31, 33] {
                let mut parts = alloc::vec![0xAAu8; 16];
                parts.extend(core::iter::repeat_n(0x55u8, ciphertext_len));
                let signed = sign_token_parts(&key, &parts);
                let chunks = split_at_offsets(&signed, &[7, 20, 40]);
                chunked_as_joined(&token, &signed, &chunks, 4);
            }
            // Signed garbage decrypts to arbitrary padding bytes.
            for fill in 0u8..64 {
                for ciphertext_len in [16usize, 32] {
                    let parts: Vec<u8> = (0..16 + ciphertext_len)
                        .map(|i| (i as u8).wrapping_mul(fill).wrapping_add(fill))
                        .collect();
                    let signed = sign_token_parts(&key, &parts);
                    for offsets in [alloc::vec![], alloc::vec![3, 19, 30]] {
                        let chunks = split_at_offsets(&signed, &offsets);
                        for skip in [0, 4, 16] {
                            chunked_as_joined(&token, &signed, &chunks, skip);
                        }
                    }
                }
            }
            for len in [0usize, 1, 48] {
                let short = alloc::vec![0u8; len];
                assert_eq!(
                    token.decrypt_chunks(&split_at_offsets(&short, &[1]), 0),
                    Err(TokenError::InvalidToken)
                );
            }
        }
    }

    #[test]
    fn untagged_encryption_with_its_mac_is_the_token() {
        for key in [[0x21u8; 32].to_vec(), [0x42u8; 64].to_vec()] {
            let token = Token::new(&key).unwrap();
            for len in [0usize, 1, 15, 16, 100, 5000] {
                let plaintext: Vec<u8> = (0..len).map(|i| (i * 5) as u8).collect();
                let expected = token.encrypt(&plaintext, &mut FixedRng::new(&[0x77; 16]));
                let (mut bytes, mut mac) =
                    token.encrypt_untagged(&plaintext, &mut FixedRng::new(&[0x77; 16]));
                let signed = bytes.len() - 32;
                assert!(bytes[signed..].iter().all(|&b| b == 0));
                mac.update(&bytes[..signed]);
                bytes[signed..].copy_from_slice(&mac.finalize());
                assert_eq!(bytes, expected);
            }
        }
    }

    #[test]
    fn streamed_plaintext_is_final_except_its_last_block() {
        let mut rng = FixedRng::new(&[0x3C; 64]);
        let token = Token::new(&[0x42u8; 64]).unwrap();
        for len in [4usize, 15, 16, 17, 100, 1000, 5000] {
            let plaintext: Vec<u8> = (0..len).map(|i| (i * 11) as u8).collect();
            let bytes = token.encrypt(&plaintext, &mut rng);
            for step in [1usize, 7, 16, 64, 500] {
                let mut stream = token.decrypt_stream(4);
                let mut partials = Vec::new();
                for chunk in bytes.chunks(step) {
                    stream.update(chunk);
                    partials.push(stream.plaintext().to_vec());
                }
                let result = stream.finish().unwrap();
                assert_eq!(result, plaintext[4..]);
                let mut hashed = token.decrypt_stream(4);
                hashed.hash_plaintext();
                for chunk in bytes.chunks(step) {
                    hashed.update(chunk);
                }
                let (again, hash) = hashed.finish_hashed().unwrap();
                assert_eq!(again, result);
                assert_eq!(hash.unwrap().digest(), crate::sha256::sha256(&result));
                for partial in partials {
                    let settled = partial.len().saturating_sub(16);
                    assert_eq!(partial[..settled], result[..settled]);
                }
            }
        }
    }

    #[test]
    fn single_buffer_encryption_matches_pad_encrypt_sign() {
        for key in [[0x21u8; 32].to_vec(), [0x42u8; 64].to_vec()] {
            let token = Token::new(&key).unwrap();
            let iv = [0x17u8; 16];
            for len in [0usize, 1, 15, 16, 17, 31, 32, 1000] {
                let plaintext: Vec<u8> = (0..len).map(|i| i as u8).collect();
                let padded = pkcs7::pad(&plaintext, 16);
                let ciphertext = match &token.mode {
                    AesMode::Aes128(aes) => aes.encrypt_cbc(&padded, &iv),
                    AesMode::Aes256(aes) => aes.encrypt_cbc(&padded, &iv),
                }
                .unwrap();
                let mut expected = iv.to_vec();
                expected.extend_from_slice(&ciphertext);
                // Tags from the raw signing key, not the token's keyed state.
                let mac = hmac_sha256(&key[..key.len() / 2], &expected);
                expected.extend_from_slice(&mac);
                let encrypted = token.encrypt_with_iv(&plaintext, &iv);
                assert_eq!(encrypted, expected, "key {} len {len}", key.len());
                assert_eq!(encrypted.capacity(), encrypted.len());
                assert_eq!(token.decrypt(&encrypted).unwrap(), plaintext);
            }
        }
    }
}