krypteia-arcana 0.1.0

//! Stateful streaming cipher context (init / update / finalize) with
//! padding, on top of the existing block-cipher and mode primitives.
//!
//! This module provides the [`Cipher`] object that wraps the block
//! ciphers (AES, DES, 3DES) and the unauthenticated modes (ECB, CBC,
//! CTR) into a single uniform `init / update / finalize` API. The
//! design intent is to mirror what OpenSSL's `EVP_CIPHER_CTX` and PSA
//! Crypto's `psa_cipher_operation_t` offer, while keeping all buffers
//! caller-provided so the implementation is suitable for embedded
//! targets without an allocator.
//!
//! AEAD modes (GCM, CCM, ChaCha20-Poly1305) are intentionally **not**
//! routed through this object; they remain function-oriented in
//! their own modules to avoid the trap of releasing unverified
//! plaintext during streaming decryption.
//!
//! # Cycle of life
//!
//! ```text
//!   new(algo, mode, padding)
//!         │
//!         ▼
//!   init(direction, key, iv) ◄──┐  (may be called again to rekey
//!         │                     │   or change direction)
//!         ▼                     │
//!   update(input, output) ──────┘
//!         │   (0..N times)
//!         ▼
//!   finalize(output)
//! ```
//!
//! # Buffer ownership
//!
//! Both `update` and `finalize` write into a caller-provided output
//! slice and return the number of bytes actually written. Callers can
//! query a safe upper bound for the output size with
//! [`Cipher::update_output_size`] / [`Cipher::finalize_output_size`]
//! before sizing the buffer. Convenience helpers
//! [`Cipher::update_to_vec`] / [`Cipher::finalize_to_vec`] are
//! provided for callers that prefer allocation.
//!
//! # Example: AES-256-CBC with PKCS#7 padding
//!
//! ```rust,ignore
//! use arcana::cipher::ctx::{Cipher, Algorithm, Mode, Padding, Direction};
//!
//! let key = [0u8; 32];
//! let iv  = [0u8; 16];
//!
//! let mut c = Cipher::new(Algorithm::Aes256, Mode::Cbc, Padding::Pkcs7).unwrap();
//! c.init(Direction::Encrypt, &key, &iv).unwrap();
//!
//! let mut out = vec![0u8; 64];
//! let mut written = 0;
//! written += c.update(b"hello, ", &mut out[written..]).unwrap();
//! written += c.update(b"world!",  &mut out[written..]).unwrap();
//! written += c.finalize(&mut out[written..]).unwrap();
//! out.truncate(written);
//! ```

use crate::BlockCipher;
use crate::cipher::aes::Aes;
use crate::cipher::des::{Des, TripleDes};

// ============================================================
// Public enums
// ============================================================

/// Symmetric block cipher selection for [`Cipher`].
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Algorithm {
    /// AES with a 128-bit key (FIPS 197). 16-byte block.
    Aes128,
    /// AES with a 192-bit key (FIPS 197). 16-byte block.
    Aes192,
    /// AES with a 256-bit key (FIPS 197). 16-byte block.
    Aes256,
    /// Single DES (FIPS 46-3). 8-byte block, 8-byte key (56 effective bits).
    /// **Legacy / cryptographically broken**: shipped for interop only.
    Des,
    /// Triple DES (FIPS 46-3, EDE3). 8-byte block, 24-byte key.
    /// **Legacy / deprecated**: shipped for interop only.
    TripleDes,
}

/// Mode of operation selection for [`Cipher`].
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Mode {
    /// Electronic Codebook (NIST SP 800-38A §6.1). Insecure for any
    /// general use; ships for interop / test-vector replay only.
    Ecb,
    /// Cipher Block Chaining (NIST SP 800-38A §6.2).
    Cbc,
    /// Counter mode (NIST SP 800-38A §6.5). Padding **must** be
    /// [`Padding::None`]; output length always equals input length.
    Ctr,
}

/// Padding scheme used at the end of an [`Cipher`] block-mode stream.
///
/// Padding only applies to ECB and CBC. CTR mode requires
/// [`Padding::None`] and the constructor will return
/// [`Error::InvalidPaddingForMode`] otherwise.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Padding {
    /// No padding. Caller must feed a multiple of the block size in
    /// total; otherwise `finalize` returns [`Error::UnpaddedInput`].
    None,
    /// PKCS#7 padding (RFC 5652 §6.3). Pads with `n` copies of the
    /// byte `n`, where `n` is the number of bytes added (1..=block).
    /// Always adds at least one byte. Equivalent to JCE's
    /// `PKCS5Padding` for 8-byte blocks.
    Pkcs7,
    /// ISO/IEC 9797-1 Padding Method 1: append `0x00` bytes until the
    /// block is full. Adds **zero** bytes if the input is already
    /// block-aligned, which makes the operation lossy on decryption
    /// (trailing zero bytes of the original message are
    /// indistinguishable from padding). Also known as **zero
    /// padding**.
    Iso9797M1,
    /// ISO/IEC 9797-1 Padding Method 2: append `0x80` followed by as
    /// many `0x00` bytes as needed to fill the block. Always adds at
    /// least one byte. Identical to ISO/IEC 7816-4 padding (also
    /// called "bit padding").
    Iso9797M2,
    /// ANSI X9.23 padding: append `0x00` bytes followed by a final
    /// byte holding the padding length. Always adds at least one
    /// byte. Legacy; ships for interop with X9.23-encoded data.
    AnsiX923,
}

/// Direction selection for [`Cipher::init`].
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Direction {
    /// Encrypt plaintext to ciphertext.
    Encrypt,
    /// Decrypt ciphertext to plaintext.
    Decrypt,
}

/// Errors returned by [`Cipher`].
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Error {
    /// The supplied key length does not match the algorithm.
    InvalidKeyLen,
    /// The supplied IV length does not match the mode's block size.
    InvalidIvLen,
    /// `Padding::None` is required for stream-shaped modes (CTR), and
    /// padded modes are required for block-shaped modes (ECB, CBC) if
    /// the caller wants to feed unaligned data.
    InvalidPaddingForMode,
    /// Caller-provided output buffer is too small for the bytes that
    /// the call would have produced. The `needed` field reports the
    /// total number of bytes the call would have written.
    OutputBufferTooSmall {
        /// The total number of output bytes the call would have written.
        needed: usize,
    },
    /// `update` / `finalize` was called before [`Cipher::init`].
    NotInitialized,
    /// `update` / `init` was called after [`Cipher::finalize`] without
    /// re-initializing.
    AlreadyFinalized,
    /// Encrypting or decrypting with [`Padding::None`] but the total
    /// input was not a multiple of the block size.
    UnpaddedInput,
    /// Decrypted padding bytes do not match the declared padding
    /// scheme.
    BadPadding,
}

// ============================================================
// Internal state
// ============================================================

const MAX_BLOCK: usize = 16;

/// Algorithm-specific key schedule, set up at [`Cipher::init`] time.
enum Key {
    None,
    Aes(Aes),
    Des(Des),
    TripleDes(TripleDes),
}

impl Key {
    fn encrypt_block(&self, blk: &mut [u8]) {
        match self {
            Key::None => unreachable!("encrypt_block on uninitialised Cipher"),
            Key::Aes(c) => c.encrypt_block(blk),
            Key::Des(c) => c.encrypt_block(blk),
            Key::TripleDes(c) => c.encrypt_block(blk),
        }
    }

    fn decrypt_block(&self, blk: &mut [u8]) {
        match self {
            Key::None => unreachable!("decrypt_block on uninitialised Cipher"),
            Key::Aes(c) => c.decrypt_block(blk),
            Key::Des(c) => c.decrypt_block(blk),
            Key::TripleDes(c) => c.decrypt_block(blk),
        }
    }
}

// ============================================================
// Cipher object
// ============================================================

/// Stateful symmetric cipher context.
///
/// See the [module-level documentation](self) for the cycle of life
/// and the buffer ownership model.
pub struct Cipher {
    algo: Algorithm,
    mode: Mode,
    padding: Padding,
    block_len: usize,

    key: Key,
    direction: Option<Direction>,

    /// For CBC: previous ciphertext block (or IV before the first
    /// block). For CTR: current counter block. Unused for ECB.
    iv: [u8; MAX_BLOCK],

    /// CTR keystream buffer and read position. `ks_pos == block_len`
    /// means the next byte requires a fresh keystream block.
    ctr_ks: [u8; MAX_BLOCK],
    ctr_ks_pos: usize,

    /// Input reliquat / look-ahead buffer (block-mode only).
    ///
    /// At any moment between `update` calls, `buf[..buf_len]` holds
    /// the most recent input bytes that have not yet been encrypted
    /// or decrypted. The buffer can be larger than one block on
    /// padded decryption, where one full ciphertext block is held
    /// back so that `finalize` can strip the padding.
    buf: [u8; 2 * MAX_BLOCK],
    buf_len: usize,

    /// True once `finalize` has run successfully.
    finalized: bool,
}

impl Cipher {
    /// Create a new cipher context with the given algorithm, mode and
    /// padding scheme. Validates that the combination is supported;
    /// the actual key and IV are loaded later by [`Self::init`].
    ///
    /// # Errors
    ///
    /// Returns [`Error::InvalidPaddingForMode`] when the requested
    /// padding is incompatible with the requested mode (e.g. any
    /// padding other than `None` with `Mode::Ctr`).
    pub fn new(algo: Algorithm, mode: Mode, padding: Padding) -> Result<Self, Error> {
        let block_len = match algo {
            Algorithm::Aes128 | Algorithm::Aes192 | Algorithm::Aes256 => 16,
            Algorithm::Des | Algorithm::TripleDes => 8,
        };

        // CTR is a stream mode: padding makes no sense.
        if matches!(mode, Mode::Ctr) && !matches!(padding, Padding::None) {
            return Err(Error::InvalidPaddingForMode);
        }

        Ok(Self {
            algo,
            mode,
            padding,
            block_len,
            key: Key::None,
            direction: None,
            iv: [0u8; MAX_BLOCK],
            ctr_ks: [0u8; MAX_BLOCK],
            ctr_ks_pos: MAX_BLOCK, // forces refill on first byte
            buf: [0u8; 2 * MAX_BLOCK],
            buf_len: 0,
            finalized: false,
        })
    }

    /// Initialise (or re-initialise) the context with a key, IV and
    /// direction. Loading a new key wipes any pending input from a
    /// previous run, so the same `Cipher` object can be reused across
    /// many independent encryptions or decryptions.
    ///
    /// `iv` must be `block_len()` bytes long for CBC and CTR; for ECB
    /// the slice is ignored but a `&[]` is the conventional choice.
    ///
    /// # Errors
    ///
    /// - [`Error::InvalidKeyLen`] if `key.len()` does not match the
    ///   algorithm.
    /// - [`Error::InvalidIvLen`] if `iv.len()` does not match the
    ///   block size for a mode that requires an IV.
    pub fn init(&mut self, direction: Direction, key: &[u8], iv: &[u8]) -> Result<(), Error> {
        // Validate key length up front.
        let expected_key_len = match self.algo {
            Algorithm::Aes128 => 16,
            Algorithm::Aes192 => 24,
            Algorithm::Aes256 => 32,
            Algorithm::Des => 8,
            Algorithm::TripleDes => 24,
        };
        if key.len() != expected_key_len {
            return Err(Error::InvalidKeyLen);
        }

        // Validate IV length per mode.
        match self.mode {
            Mode::Ecb => { /* iv ignored */ }
            Mode::Cbc | Mode::Ctr => {
                if iv.len() != self.block_len {
                    return Err(Error::InvalidIvLen);
                }
            }
        }

        // Build the key schedule.
        self.key = match self.algo {
            Algorithm::Aes128 | Algorithm::Aes192 | Algorithm::Aes256 => Key::Aes(Aes::new(key)),
            Algorithm::Des => Key::Des(Des::new(key)),
            Algorithm::TripleDes => Key::TripleDes(TripleDes::new(key)),
        };

        // Reset mode state.
        self.iv = [0u8; MAX_BLOCK];
        if matches!(self.mode, Mode::Cbc | Mode::Ctr) {
            self.iv[..self.block_len].copy_from_slice(iv);
        }
        self.ctr_ks = [0u8; MAX_BLOCK];
        self.ctr_ks_pos = self.block_len; // empty keystream → refill on first byte
        self.buf = [0u8; 2 * MAX_BLOCK];
        self.buf_len = 0;
        self.direction = Some(direction);
        self.finalized = false;

        Ok(())
    }

    /// Block size of the underlying cipher in bytes (16 for AES, 8
    /// for DES / 3DES).
    pub const fn block_len(&self) -> usize {
        self.block_len
    }

    /// IV / nonce length expected by [`Self::init`] for the configured
    /// mode. Returns 0 for ECB.
    pub const fn iv_len(&self) -> usize {
        match self.mode {
            Mode::Ecb => 0,
            Mode::Cbc | Mode::Ctr => self.block_len,
        }
    }

    /// Safe upper bound on the number of bytes a single
    /// [`Self::update`] call would write into `output`, given an
    /// `input_len`-byte input. Use this to size caller-provided
    /// buffers.
    pub fn update_output_size(&self, input_len: usize) -> usize {
        match self.mode {
            Mode::Ctr => input_len,
            Mode::Ecb | Mode::Cbc => {
                // At most everything we have buffered plus the new
                // input, rounded up to a full block.
                let total = self.buf_len + input_len;
                if total < self.block_len {
                    0
                } else {
                    // Worst case: emit floor(total / block_len) blocks.
                    (total / self.block_len) * self.block_len
                }
            }
        }
    }

    /// Safe upper bound on the number of bytes [`Self::finalize`]
    /// would write into `output`. Always at most one block.
    pub const fn finalize_output_size(&self) -> usize {
        match self.mode {
            Mode::Ctr => 0,
            Mode::Ecb | Mode::Cbc => self.block_len,
        }
    }

    /// Feed `input` bytes through the cipher and write any complete
    /// output bytes into `output`. Returns the number of bytes
    /// actually written.
    ///
    /// Either slice may be empty. Bytes that don't yet form a full
    /// output block stay buffered inside the `Cipher` until enough
    /// data arrives or [`Self::finalize`] is called.
    ///
    /// # Errors
    ///
    /// - [`Error::NotInitialized`] if [`Self::init`] has not been called.
    /// - [`Error::AlreadyFinalized`] if [`Self::finalize`] has already run.
    /// - [`Error::OutputBufferTooSmall`] if `output` is shorter than
    ///   the number of bytes the call would write.
    pub fn update(&mut self, input: &[u8], output: &mut [u8]) -> Result<usize, Error> {
        let dir = self.direction.ok_or(Error::NotInitialized)?;
        if self.finalized {
            return Err(Error::AlreadyFinalized);
        }

        match self.mode {
            Mode::Ctr => self.update_ctr(input, output),
            Mode::Ecb | Mode::Cbc => self.update_block(dir, input, output),
        }
    }

    /// Finish the operation: pad and emit the trailing block (encrypt)
    /// or strip the padding from the last buffered ciphertext block
    /// (decrypt). After this call the context is in a finalized state
    /// and must be re-initialised with [`Self::init`] before further
    /// use.
    ///
    /// # Errors
    ///
    /// - [`Error::NotInitialized`] if [`Self::init`] has not been called.
    /// - [`Error::AlreadyFinalized`] if `finalize` has already run.
    /// - [`Error::OutputBufferTooSmall`] if `output` is too small.
    /// - [`Error::UnpaddedInput`] if `Padding::None` was selected and
    ///   the total input length is not a multiple of the block size.
    /// - [`Error::BadPadding`] (decryption only) if the trailing
    ///   plaintext does not match the declared padding scheme.
    pub fn finalize(&mut self, output: &mut [u8]) -> Result<usize, Error> {
        let dir = self.direction.ok_or(Error::NotInitialized)?;
        if self.finalized {
            return Err(Error::AlreadyFinalized);
        }

        let written = match self.mode {
            Mode::Ctr => {
                // Stream mode: nothing buffered, nothing to flush.
                0
            }
            Mode::Ecb | Mode::Cbc => match dir {
                Direction::Encrypt => self.finalize_block_encrypt(output)?,
                Direction::Decrypt => self.finalize_block_decrypt(output)?,
            },
        };

        self.finalized = true;
        Ok(written)
    }

    // --------------------------------------------------------
    // CTR mode update path
    // --------------------------------------------------------

    fn update_ctr(&mut self, input: &[u8], output: &mut [u8]) -> Result<usize, Error> {
        if output.len() < input.len() {
            return Err(Error::OutputBufferTooSmall { needed: input.len() });
        }
        let bs = self.block_len;
        for i in 0..input.len() {
            if self.ctr_ks_pos == bs {
                // Generate next keystream block from the current counter,
                // then increment the counter (rightmost 64 bits, big-endian).
                self.ctr_ks[..bs].copy_from_slice(&self.iv[..bs]);
                self.key.encrypt_block(&mut self.ctr_ks[..bs]);
                ctr_increment(&mut self.iv[..bs]);
                self.ctr_ks_pos = 0;
            }
            output[i] = input[i] ^ self.ctr_ks[self.ctr_ks_pos];
            self.ctr_ks_pos += 1;
        }
        Ok(input.len())
    }

    // --------------------------------------------------------
    // Block-mode (ECB / CBC) update path
    // --------------------------------------------------------

    /// Number of bytes we must keep buffered until `finalize` so that
    /// the padding can be stripped on decryption. Zero in every other
    /// configuration.
    fn keep_back(&self, dir: Direction) -> usize {
        if matches!(dir, Direction::Decrypt) && !matches!(self.padding, Padding::None) {
            self.block_len
        } else {
            0
        }
    }

    fn update_block(&mut self, dir: Direction, input: &[u8], output: &mut [u8]) -> Result<usize, Error> {
        let bs = self.block_len;
        let kb = self.keep_back(dir);
        let mut written = 0usize;
        let mut in_pos = 0usize;
        let cap = 2 * bs;

        while in_pos < input.len() {
            // 1. Fill the internal buffer as much as it can hold.
            let space = cap - self.buf_len;
            if space > 0 {
                let take = space.min(input.len() - in_pos);
                self.buf[self.buf_len..self.buf_len + take].copy_from_slice(&input[in_pos..in_pos + take]);
                self.buf_len += take;
                in_pos += take;
            }

            // 2. Drain as many full blocks as we can while still
            //    leaving `kb` bytes behind for finalize.
            while self.buf_len >= bs && self.buf_len - bs >= kb {
                if output.len() - written < bs {
                    return Err(Error::OutputBufferTooSmall { needed: written + bs });
                }
                let mut blk = [0u8; MAX_BLOCK];
                blk[..bs].copy_from_slice(&self.buf[..bs]);
                self.process_one_block(dir, &mut blk[..bs]);
                output[written..written + bs].copy_from_slice(&blk[..bs]);
                written += bs;
                // Shift remaining bytes to the start of buf.
                self.buf.copy_within(bs..self.buf_len, 0);
                self.buf_len -= bs;
            }
        }

        Ok(written)
    }

    /// Process exactly one block in-place, applying the mode's
    /// chaining if any.
    fn process_one_block(&mut self, dir: Direction, blk: &mut [u8]) {
        let bs = self.block_len;
        match (self.mode, dir) {
            (Mode::Ecb, Direction::Encrypt) => {
                self.key.encrypt_block(blk);
            }
            (Mode::Ecb, Direction::Decrypt) => {
                self.key.decrypt_block(blk);
            }
            (Mode::Cbc, Direction::Encrypt) => {
                for i in 0..bs {
                    blk[i] ^= self.iv[i];
                }
                self.key.encrypt_block(blk);
                self.iv[..bs].copy_from_slice(&blk[..bs]);
            }
            (Mode::Cbc, Direction::Decrypt) => {
                let ct_copy = {
                    let mut tmp = [0u8; MAX_BLOCK];
                    tmp[..bs].copy_from_slice(&blk[..bs]);
                    tmp
                };
                self.key.decrypt_block(blk);
                for i in 0..bs {
                    blk[i] ^= self.iv[i];
                }
                self.iv[..bs].copy_from_slice(&ct_copy[..bs]);
            }
            (Mode::Ctr, _) => unreachable!("CTR uses update_ctr"),
        }
    }

    // --------------------------------------------------------
    // Finalize: encrypt path
    // --------------------------------------------------------

    fn finalize_block_encrypt(&mut self, output: &mut [u8]) -> Result<usize, Error> {
        let bs = self.block_len;

        match self.padding {
            Padding::None => {
                if self.buf_len != 0 {
                    return Err(Error::UnpaddedInput);
                }
                Ok(0)
            }
            Padding::Pkcs7 | Padding::Iso9797M2 | Padding::Iso9797M1 | Padding::AnsiX923 => {
                // Iso9797M1 = zero padding: if already aligned, no padding
                // is added at all.
                if matches!(self.padding, Padding::Iso9797M1) && self.buf_len == 0 {
                    return Ok(0);
                }

                if output.len() < bs {
                    return Err(Error::OutputBufferTooSmall { needed: bs });
                }
                let pad_len = bs - self.buf_len; // 1..=bs for the others
                apply_padding(self.padding, &mut self.buf[..bs], self.buf_len, pad_len);
                self.buf_len = bs;

                let mut blk = [0u8; MAX_BLOCK];
                blk[..bs].copy_from_slice(&self.buf[..bs]);
                self.process_one_block(Direction::Encrypt, &mut blk[..bs]);
                output[..bs].copy_from_slice(&blk[..bs]);
                self.buf_len = 0;
                Ok(bs)
            }
        }
    }

    // --------------------------------------------------------
    // Finalize: decrypt path
    // --------------------------------------------------------

    fn finalize_block_decrypt(&mut self, output: &mut [u8]) -> Result<usize, Error> {
        let bs = self.block_len;

        match self.padding {
            Padding::None => {
                if self.buf_len != 0 {
                    return Err(Error::UnpaddedInput);
                }
                Ok(0)
            }
            _ => {
                // We must have exactly one full block buffered (the
                // reserved last ciphertext block). If buf_len == 0 the
                // caller fed nothing at all.
                if self.buf_len != bs {
                    return Err(Error::UnpaddedInput);
                }
                let mut blk = [0u8; MAX_BLOCK];
                blk[..bs].copy_from_slice(&self.buf[..bs]);
                self.process_one_block(Direction::Decrypt, &mut blk[..bs]);

                let unpadded = strip_padding(self.padding, &blk[..bs])?;
                if output.len() < unpadded {
                    return Err(Error::OutputBufferTooSmall { needed: unpadded });
                }
                output[..unpadded].copy_from_slice(&blk[..unpadded]);
                self.buf_len = 0;
                Ok(unpadded)
            }
        }
    }

    // --------------------------------------------------------
    // Allocating helpers
    // --------------------------------------------------------

    /// Convenience wrapper around [`Self::update`] that returns the
    /// freshly produced bytes as a `Vec<u8>`. Allocates.
    pub fn update_to_vec(&mut self, input: &[u8]) -> Result<Vec<u8>, Error> {
        let upper = self.update_output_size(input.len());
        let mut out = vec![0u8; upper];
        let n = self.update(input, &mut out)?;
        out.truncate(n);
        Ok(out)
    }

    /// Convenience wrapper around [`Self::finalize`] that returns the
    /// trailing bytes as a `Vec<u8>`. Allocates.
    pub fn finalize_to_vec(&mut self) -> Result<Vec<u8>, Error> {
        let upper = self.finalize_output_size();
        let mut out = vec![0u8; upper];
        let n = self.finalize(&mut out)?;
        out.truncate(n);
        Ok(out)
    }
}

// ============================================================
// Padding helpers
// ============================================================

/// Append `pad_len` padding bytes starting at offset `data_len` in
/// `block`. The block must already be sized to one full block; the
/// caller is responsible for `data_len + pad_len == block.len()`.
fn apply_padding(padding: Padding, block: &mut [u8], data_len: usize, pad_len: usize) {
    let bs = block.len();
    debug_assert_eq!(data_len + pad_len, bs);
    match padding {
        Padding::None => {}
        Padding::Pkcs7 => {
            for b in &mut block[data_len..bs] {
                *b = pad_len as u8;
            }
        }
        Padding::Iso9797M1 => {
            for b in &mut block[data_len..bs] {
                *b = 0x00;
            }
        }
        Padding::Iso9797M2 => {
            block[data_len] = 0x80;
            for b in &mut block[data_len + 1..bs] {
                *b = 0x00;
            }
        }
        Padding::AnsiX923 => {
            for b in &mut block[data_len..bs - 1] {
                *b = 0x00;
            }
            block[bs - 1] = pad_len as u8;
        }
    }
}

/// Validate the padding bytes of a freshly decrypted last block and
/// return the unpadded length.
fn strip_padding(padding: Padding, block: &[u8]) -> Result<usize, Error> {
    let bs = block.len();
    match padding {
        Padding::None => Ok(bs),
        Padding::Pkcs7 => {
            let pad_len = block[bs - 1] as usize;
            if pad_len == 0 || pad_len > bs {
                return Err(Error::BadPadding);
            }
            for &b in &block[bs - pad_len..bs] {
                if b as usize != pad_len {
                    return Err(Error::BadPadding);
                }
            }
            Ok(bs - pad_len)
        }
        Padding::Iso9797M1 => {
            // Zero padding is ambiguous: we cannot tell trailing zero
            // bytes of the message from padding. We adopt the
            // convention "strip all trailing zeros", which matches
            // most ISO 9797-1 M1 deployments.
            let mut len = bs;
            while len > 0 && block[len - 1] == 0x00 {
                len -= 1;
            }
            Ok(len)
        }
        Padding::Iso9797M2 => {
            // Find the rightmost 0x80 preceded only by 0x00 bytes.
            let mut i = bs;
            while i > 0 {
                i -= 1;
                if block[i] == 0x80 {
                    return Ok(i);
                }
                if block[i] != 0x00 {
                    return Err(Error::BadPadding);
                }
            }
            Err(Error::BadPadding)
        }
        Padding::AnsiX923 => {
            let pad_len = block[bs - 1] as usize;
            if pad_len == 0 || pad_len > bs {
                return Err(Error::BadPadding);
            }
            for &b in &block[bs - pad_len..bs - 1] {
                if b != 0x00 {
                    return Err(Error::BadPadding);
                }
            }
            Ok(bs - pad_len)
        }
    }
}

// ============================================================
// CTR counter helper
// ============================================================

/// Increment the rightmost 64 bits of an n-byte counter (big-endian).
/// Used by CTR mode update path.
fn ctr_increment(counter: &mut [u8]) {
    let n = counter.len();
    let lo = n.saturating_sub(8);
    for i in (lo..n).rev() {
        let (v, carry) = counter[i].overflowing_add(1);
        counter[i] = v;
        if !carry {
            return;
        }
    }
}

// ============================================================
// Tests
// ============================================================

#[cfg(test)]
mod tests {
    use super::*;
    use crate::BlockCipher;
    use crate::cipher::aes::Aes128;
    use crate::cipher::modes::{cbc_decrypt, cbc_encrypt, ctr_encrypt, ecb_encrypt};

    fn hex(s: &str) -> Vec<u8> {
        (0..s.len())
            .step_by(2)
            .map(|i| u8::from_str_radix(&s[i..i + 2], 16).unwrap())
            .collect()
    }

    // ---------- ECB ----------

    #[test]
    fn ecb_aes128_no_padding_round_trip() {
        let key = hex("2b7e151628aed2a6abf7158809cf4f3c");
        let pt = hex("6bc1bee22e409f96e93d7e117393172a"); // 16 bytes

        let mut enc = Cipher::new(Algorithm::Aes128, Mode::Ecb, Padding::None).unwrap();
        enc.init(Direction::Encrypt, &key, &[]).unwrap();
        let ct = {
            let mut out = vec![0u8; 32];
            let n = enc.update(&pt, &mut out).unwrap();
            let m = enc.finalize(&mut out[n..]).unwrap();
            out.truncate(n + m);
            out
        };

        let mut dec = Cipher::new(Algorithm::Aes128, Mode::Ecb, Padding::None).unwrap();
        dec.init(Direction::Decrypt, &key, &[]).unwrap();
        let mut got = vec![0u8; 32];
        let n = dec.update(&ct, &mut got).unwrap();
        let m = dec.finalize(&mut got[n..]).unwrap();
        got.truncate(n + m);
        assert_eq!(got, pt);
    }

    #[test]
    fn ecb_aes128_pkcs7_unaligned() {
        let key = [0x42u8; 16];
        let pt: Vec<u8> = (0u8..30).collect(); // 30 bytes → 2 blocks of CT

        // Reference: pad with PKCS#7 by hand and call ecb_encrypt.
        let mut padded = pt.clone();
        let pad = 16 - (padded.len() % 16);
        padded.extend(std::iter::repeat(pad as u8).take(pad));
        let cipher_ref = Aes128::new(&key);
        ecb_encrypt(&cipher_ref, &mut padded);
        let ct_ref = padded;

        // Round-trip via Cipher.
        let mut enc = Cipher::new(Algorithm::Aes128, Mode::Ecb, Padding::Pkcs7).unwrap();
        enc.init(Direction::Encrypt, &key, &[]).unwrap();
        let ct = {
            let mut out = vec![0u8; 64];
            let n = enc.update(&pt, &mut out).unwrap();
            let m = enc.finalize(&mut out[n..]).unwrap();
            out.truncate(n + m);
            out
        };
        assert_eq!(ct, ct_ref, "Cipher PKCS#7 ECB encrypt mismatch");

        let mut dec = Cipher::new(Algorithm::Aes128, Mode::Ecb, Padding::Pkcs7).unwrap();
        dec.init(Direction::Decrypt, &key, &[]).unwrap();
        let pt_back = {
            let mut out = vec![0u8; 64];
            let n = dec.update(&ct, &mut out).unwrap();
            let m = dec.finalize(&mut out[n..]).unwrap();
            out.truncate(n + m);
            out
        };
        assert_eq!(pt_back, pt);
    }

    // ---------- CBC ----------

    #[test]
    fn cbc_aes128_pkcs7_matches_modes_helper() {
        let key = hex("2b7e151628aed2a6abf7158809cf4f3c");
        let iv = hex("000102030405060708090a0b0c0d0e0f");
        let pt: Vec<u8> = (0u8..37).collect(); // unaligned

        // Reference path through cipher::modes::cbc_encrypt.
        let mut padded = pt.clone();
        let pad = 16 - (padded.len() % 16);
        padded.extend(std::iter::repeat(pad as u8).take(pad));
        let cipher_ref = Aes128::new(&key);
        cbc_encrypt(&cipher_ref, &iv, &mut padded);
        let ct_ref = padded;

        let mut enc = Cipher::new(Algorithm::Aes128, Mode::Cbc, Padding::Pkcs7).unwrap();
        enc.init(Direction::Encrypt, &key, &iv).unwrap();
        let mut ct = vec![0u8; 64];
        // Feed in two pieces to exercise the streaming path.
        let mut w = enc.update(&pt[..10], &mut ct).unwrap();
        w += enc.update(&pt[10..], &mut ct[w..]).unwrap();
        w += enc.finalize(&mut ct[w..]).unwrap();
        ct.truncate(w);
        assert_eq!(ct, ct_ref, "Cipher CBC PKCS#7 encrypt mismatch");

        // Decrypt via cbc_decrypt and via Cipher.
        let mut ref_pt = ct.clone();
        cbc_decrypt(&cipher_ref, &iv, &mut ref_pt);
        // Strip PKCS#7 from the reference.
        let pad = *ref_pt.last().unwrap() as usize;
        ref_pt.truncate(ref_pt.len() - pad);
        assert_eq!(ref_pt, pt);

        let mut dec = Cipher::new(Algorithm::Aes128, Mode::Cbc, Padding::Pkcs7).unwrap();
        dec.init(Direction::Decrypt, &key, &iv).unwrap();
        let mut got = vec![0u8; 64];
        let mut n = dec.update(&ct[..7], &mut got).unwrap();
        n += dec.update(&ct[7..], &mut got[n..]).unwrap();
        n += dec.finalize(&mut got[n..]).unwrap();
        got.truncate(n);
        assert_eq!(got, pt);
    }

    #[test]
    fn cbc_aes128_none_padding_aligned() {
        let key = hex("2b7e151628aed2a6abf7158809cf4f3c");
        let iv = hex("000102030405060708090a0b0c0d0e0f");
        let pt = hex("6bc1bee22e409f96e93d7e117393172aae2d8a571e03ac9c9eb76fac45af8e51");

        let mut enc = Cipher::new(Algorithm::Aes128, Mode::Cbc, Padding::None).unwrap();
        enc.init(Direction::Encrypt, &key, &iv).unwrap();
        let mut ct = vec![0u8; 64];
        let mut n = enc.update(&pt, &mut ct).unwrap();
        n += enc.finalize(&mut ct[n..]).unwrap();
        ct.truncate(n);

        // Reference.
        let cipher_ref = Aes128::new(&key);
        let mut ct_ref = pt.clone();
        cbc_encrypt(&cipher_ref, &iv, &mut ct_ref);
        assert_eq!(ct, ct_ref);
    }

    #[test]
    fn cbc_aes128_none_padding_unaligned_errors() {
        let key = [0u8; 16];
        let iv = [0u8; 16];
        let mut enc = Cipher::new(Algorithm::Aes128, Mode::Cbc, Padding::None).unwrap();
        enc.init(Direction::Encrypt, &key, &iv).unwrap();
        let mut out = vec![0u8; 64];
        enc.update(&[0u8; 13], &mut out).unwrap();
        assert_eq!(enc.finalize(&mut out), Err(Error::UnpaddedInput));
    }

    // ---------- CTR ----------

    #[test]
    fn ctr_aes128_streaming_matches_one_shot() {
        let key = hex("2b7e151628aed2a6abf7158809cf4f3c");
        let nonce = hex("f0f1f2f3f4f5f6f7f8f9fafbfcfdfeff"); // 16-byte counter block
        let pt = hex("6bc1bee22e409f96e93d7e117393172aae2d8a571e03ac9c9eb76fac45af8e51\
             30c81c46a35ce411e5fbc1191a0a52eff69f2445df4f9b17ad2b417be66c3710");

        // Reference: cipher::modes::ctr_encrypt with shorter nonce; we
        // need to mimic the CTR layout. For interoperability we use
        // the full 16-byte counter block path: build the cipher_ref
        // and step the counter manually.
        let cipher_ref = Aes128::new(&key);
        let mut ct_ref = pt.clone();
        // Use ctr_encrypt with nonce of length < block to match its API:
        // here we use nonce of 8 bytes, counter starts at 1 in the rest.
        // For Cipher we will use the same equivalent counter block:
        //   nonce[0..8] || 0x0000000000000001
        let nonce8 = &nonce[..8];
        ctr_encrypt(&cipher_ref, nonce8, &mut ct_ref);

        // Build the matching 16-byte initial counter block for Cipher.
        let mut iv = [0u8; 16];
        iv[..8].copy_from_slice(nonce8);
        iv[15] = 1;

        let mut enc = Cipher::new(Algorithm::Aes128, Mode::Ctr, Padding::None).unwrap();
        enc.init(Direction::Encrypt, &key, &iv).unwrap();
        let mut ct = vec![0u8; pt.len()];
        // Feed in awkward chunks.
        let mut w = 0;
        for chunk in pt.chunks(7) {
            w += enc.update(chunk, &mut ct[w..]).unwrap();
        }
        w += enc.finalize(&mut ct[w..]).unwrap();
        assert_eq!(w, pt.len());
        assert_eq!(ct, ct_ref, "Cipher CTR streaming mismatch");

        // Decrypt path (CTR is symmetric).
        let mut dec = Cipher::new(Algorithm::Aes128, Mode::Ctr, Padding::None).unwrap();
        dec.init(Direction::Decrypt, &key, &iv).unwrap();
        let mut pt_back = vec![0u8; ct.len()];
        let mut w = 0;
        for chunk in ct.chunks(11) {
            w += dec.update(chunk, &mut pt_back[w..]).unwrap();
        }
        w += dec.finalize(&mut pt_back[w..]).unwrap();
        assert_eq!(w, ct.len());
        assert_eq!(pt_back, pt);
    }

    // ---------- Padding variants ----------

    fn pad_round_trip(padding: Padding, msg_len: usize) {
        let key = [0x33u8; 16];
        let iv = [0x77u8; 16];
        let pt: Vec<u8> = (0..msg_len).map(|i| (i as u8).wrapping_mul(7)).collect();

        let mut enc = Cipher::new(Algorithm::Aes128, Mode::Cbc, padding).unwrap();
        enc.init(Direction::Encrypt, &key, &iv).unwrap();
        let mut ct = vec![0u8; pt.len() + 16];
        let mut n = enc.update(&pt, &mut ct).unwrap();
        n += enc.finalize(&mut ct[n..]).unwrap();
        ct.truncate(n);

        let mut dec = Cipher::new(Algorithm::Aes128, Mode::Cbc, padding).unwrap();
        dec.init(Direction::Decrypt, &key, &iv).unwrap();
        let mut got = vec![0u8; ct.len() + 16];
        let mut n = dec.update(&ct, &mut got).unwrap();
        n += dec.finalize(&mut got[n..]).unwrap();
        got.truncate(n);
        assert_eq!(got, pt, "round-trip {:?} len={}", padding, msg_len);
    }

    #[test]
    fn pkcs7_round_trips() {
        for len in &[0, 1, 15, 16, 17, 31, 32, 33, 100] {
            pad_round_trip(Padding::Pkcs7, *len);
        }
    }

    #[test]
    fn iso9797_m2_round_trips() {
        for len in &[0, 1, 15, 16, 17, 31, 32, 33, 100] {
            pad_round_trip(Padding::Iso9797M2, *len);
        }
    }

    #[test]
    fn ansix923_round_trips() {
        for len in &[1, 15, 16, 17, 31, 32, 33, 100] {
            pad_round_trip(Padding::AnsiX923, *len);
        }
    }

    #[test]
    fn iso9797_m1_round_trips_when_no_trailing_zero() {
        // Zero padding can't survive trailing zero bytes; exercise it
        // only with messages that don't end in 0x00.
        let padding = Padding::Iso9797M1;
        let key = [0u8; 16];
        let iv = [0u8; 16];

        for len in &[1usize, 5, 15, 17, 30] {
            let pt: Vec<u8> = (0..*len).map(|i| (i as u8) | 1).collect();
            let mut enc = Cipher::new(Algorithm::Aes128, Mode::Cbc, padding).unwrap();
            enc.init(Direction::Encrypt, &key, &iv).unwrap();
            let mut ct = vec![0u8; pt.len() + 16];
            let mut n = enc.update(&pt, &mut ct).unwrap();
            n += enc.finalize(&mut ct[n..]).unwrap();
            ct.truncate(n);

            let mut dec = Cipher::new(Algorithm::Aes128, Mode::Cbc, padding).unwrap();
            dec.init(Direction::Decrypt, &key, &iv).unwrap();
            let mut got = vec![0u8; ct.len() + 16];
            let mut n = dec.update(&ct, &mut got).unwrap();
            n += dec.finalize(&mut got[n..]).unwrap();
            got.truncate(n);
            assert_eq!(got, pt);
        }
    }

    // ---------- Error paths ----------

    #[test]
    fn invalid_key_len_rejected() {
        let mut c = Cipher::new(Algorithm::Aes128, Mode::Cbc, Padding::Pkcs7).unwrap();
        assert_eq!(
            c.init(Direction::Encrypt, &[0u8; 17], &[0u8; 16]),
            Err(Error::InvalidKeyLen)
        );
    }

    #[test]
    fn invalid_iv_len_rejected() {
        let mut c = Cipher::new(Algorithm::Aes128, Mode::Cbc, Padding::Pkcs7).unwrap();
        assert_eq!(
            c.init(Direction::Encrypt, &[0u8; 16], &[0u8; 8]),
            Err(Error::InvalidIvLen)
        );
    }

    #[test]
    fn ctr_with_padding_rejected() {
        assert_eq!(
            Cipher::new(Algorithm::Aes128, Mode::Ctr, Padding::Pkcs7).err(),
            Some(Error::InvalidPaddingForMode)
        );
    }

    #[test]
    fn output_too_small_reported() {
        let key = [0u8; 16];
        let iv = [0u8; 16];
        let mut enc = Cipher::new(Algorithm::Aes128, Mode::Cbc, Padding::Pkcs7).unwrap();
        enc.init(Direction::Encrypt, &key, &iv).unwrap();
        let mut tiny = [0u8; 4];
        let err = enc.update(&[0u8; 32], &mut tiny);
        assert!(matches!(err, Err(Error::OutputBufferTooSmall { .. })));
    }

    #[test]
    fn bad_padding_detected() {
        // Encrypt some data, flip the last ciphertext byte, expect BadPadding.
        let key = [1u8; 16];
        let iv = [2u8; 16];
        let pt = b"hello, world! Some content here.";

        let mut enc = Cipher::new(Algorithm::Aes128, Mode::Cbc, Padding::Pkcs7).unwrap();
        enc.init(Direction::Encrypt, &key, &iv).unwrap();
        let mut ct = vec![0u8; 64];
        let mut n = enc.update(pt, &mut ct).unwrap();
        n += enc.finalize(&mut ct[n..]).unwrap();
        ct.truncate(n);
        let last = ct.len() - 1;
        ct[last] ^= 0xFF;

        let mut dec = Cipher::new(Algorithm::Aes128, Mode::Cbc, Padding::Pkcs7).unwrap();
        dec.init(Direction::Decrypt, &key, &iv).unwrap();
        let mut out = vec![0u8; 64];
        let n = dec.update(&ct, &mut out).unwrap();
        let r = dec.finalize(&mut out[n..]);
        assert_eq!(r, Err(Error::BadPadding));
    }

    // ---------- DES / 3DES ----------

    #[test]
    fn tripledes_cbc_pkcs7_round_trip() {
        let key = [0x11u8; 24];
        let iv = [0x22u8; 8];
        let pt: Vec<u8> = (0u8..23).collect();

        let mut enc = Cipher::new(Algorithm::TripleDes, Mode::Cbc, Padding::Pkcs7).unwrap();
        enc.init(Direction::Encrypt, &key, &iv).unwrap();
        let mut ct = vec![0u8; 64];
        let mut n = enc.update(&pt, &mut ct).unwrap();
        n += enc.finalize(&mut ct[n..]).unwrap();
        ct.truncate(n);
        assert_eq!(ct.len() % 8, 0);

        let mut dec = Cipher::new(Algorithm::TripleDes, Mode::Cbc, Padding::Pkcs7).unwrap();
        dec.init(Direction::Decrypt, &key, &iv).unwrap();
        let mut got = vec![0u8; 64];
        let mut n = dec.update(&ct, &mut got).unwrap();
        n += dec.finalize(&mut got[n..]).unwrap();
        got.truncate(n);
        assert_eq!(got, pt);
    }

    #[test]
    fn reuse_after_init() {
        let key = [9u8; 16];
        let iv = [8u8; 16];
        let mut c = Cipher::new(Algorithm::Aes128, Mode::Cbc, Padding::Pkcs7).unwrap();

        for &msg in &[b"first message".as_slice(), b"second", b"a third one!"] {
            c.init(Direction::Encrypt, &key, &iv).unwrap();
            let ct = {
                let mut out = vec![0u8; 64];
                let mut n = c.update(msg, &mut out).unwrap();
                n += c.finalize(&mut out[n..]).unwrap();
                out.truncate(n);
                out
            };
            c.init(Direction::Decrypt, &key, &iv).unwrap();
            let pt = {
                let mut out = vec![0u8; 64];
                let mut n = c.update(&ct, &mut out).unwrap();
                n += c.finalize(&mut out[n..]).unwrap();
                out.truncate(n);
                out
            };
            assert_eq!(pt, msg);
        }
    }
}