sponge-hash-aes256 1.2.0

// SPDX-License-Identifier: 0BSD
// Copyright (C) 2025 by LoRd_MuldeR <mulder2@gmx.de>

use crate::utilities::{BLOCK_SIZE, aes256_encrypt, xor_arrays};
use zeroize::Zeroize;

/// Default digest size, in bytes
///
/// The default digest size is currently defined as 32 bytes, i.e., 256 bits.
pub const DEFAULT_DIGEST_SIZE: usize = 2usize * BLOCK_SIZE;

/// Default number of permutation rounds to be performed
///
/// The default number of permutation rounds is currently defined as `3usize`.
pub const DEFAULT_PERMUTE_ROUNDS: usize = 3usize;

// ---------------------------------------------------------------------------
// Tracing
// ---------------------------------------------------------------------------

#[cfg(feature = "tracing")]
macro_rules! trace {
    ($self:tt, $arg:tt) => {
        log::trace!(
            "SpongeHash256@{:p}: {} --> {:02x?} {:02x?} {:02x?}",
            &$self,
            $arg,
            &$self.state0,
            &$self.state1,
            &$self.state2
        );
    };
}

#[cfg(not(feature = "tracing"))]
macro_rules! trace {
    ($self:tt, $arg:tt) => {};
}

// ---------------------------------------------------------------------------
// Non-zero argument constraint
// ---------------------------------------------------------------------------

/// Validates that the const generic parameter is non-zero
struct NoneZero<const N: usize>;

impl<const N: usize> NoneZero<N> {
    const OK: () = assert!(N > 0, "Const generic argument must be a non-zero value!");
}

// ---------------------------------------------------------------------------
// Streaming API
// ---------------------------------------------------------------------------

/// This struct encapsulates the state for a “streaming” (incremental) SpongeHash-AES256 computation.
///
/// The const generic parameter `R` specifies the number of permutation rounds to be performed, which must be a *positive* value. The default number of permutation rounds is given by [`DEFAULT_PERMUTE_ROUNDS`]. Using a greater value slows down the hash calculation, which helps to increase the security in some usage scenarios, e.g., password hashing.
///
/// ### Usage Example
///
/// The easiest way to use the **`SpongeHash256`** structure is as follows:
///
/// ```rust
/// use hex::encode_to_slice;
/// use sponge_hash_aes256::{DEFAULT_DIGEST_SIZE, SpongeHash256};
///
/// fn main() {
///     // Create new hash instance
///     let mut hash = SpongeHash256::default();
///
///     // Process message
///     hash.update(b"The quick brown fox jumps over the lazy dog");
///
///     // Retrieve the final digest
///     let digest = hash.digest::<DEFAULT_DIGEST_SIZE>();
///
///     // Encode to hex
///     let mut hex_buffer = [0u8; 2usize * DEFAULT_DIGEST_SIZE];
///     encode_to_slice(&digest, &mut hex_buffer).unwrap();
///
///     // Print the digest (hex format)
///     println!("0x{}", str::from_utf8(&hex_buffer).unwrap());
/// }
/// ```
/// &nbsp;
///
/// Optionally, additional “context” information may be provided via the `info` parameter:
///
/// ```rust
/// use sponge_hash_aes256::{DEFAULT_DIGEST_SIZE, SpongeHash256};
///
/// fn main() {
///     // Create new hash instance with “info”
///     let mut hash: SpongeHash256 = SpongeHash256::with_info("my_application");
///
///     /* ... */
/// }
/// ```
/// &nbsp;
///
/// <div class="warning">
///
/// The [`compute()`] and [`compute_to_slice()`] convenience functions may be used as an alternative to working with the `SpongeHash256` struct directly. This is especially useful, if *all* data to be hashed is available at once.
///
/// </div>
pub struct SpongeHash256<const R: usize = DEFAULT_PERMUTE_ROUNDS> {
    state0: [u8; BLOCK_SIZE],
    state1: [u8; BLOCK_SIZE],
    state2: [u8; BLOCK_SIZE],
    offset: usize,
}

impl<const R: usize> SpongeHash256<R> {
    const BIT_MASK_0: [u8; BLOCK_SIZE] = [0x5Cu8; BLOCK_SIZE];
    const BIT_MASK_1: [u8; BLOCK_SIZE] = [0x36u8; BLOCK_SIZE];

    /// Creates a new SpongeHash-AES256 instance and initializes the hash computation.
    ///
    /// **Note:** This function implies an *empty* [`info`](Self::with_info()) string.
    pub fn new() -> Self {
        Self::with_info(Default::default())
    }

    /// Creates a new SpongeHash-AES256 instance and initializes the hash computation with the given `info` string.
    ///
    /// **Note:** The length of the `info` string **must not** exceed a length of 255 characters!
    pub fn with_info(info: &str) -> Self {
        let () = NoneZero::<R>::OK;

        let mut instance = Self {
            state0: [0u8; BLOCK_SIZE],
            state1: [0u8; BLOCK_SIZE],
            state2: [0u8; BLOCK_SIZE],
            offset: 0usize,
        };

        instance.initialize(info.as_bytes());
        instance
    }

    /// Initializes the internal state with the given `info` string
    fn initialize(&mut self, info_data: &[u8]) {
        trace!(self, "initlz::enter");

        match info_data.len().try_into() {
            Ok(length) => {
                self.update(u8::to_be_bytes(length));
                self.update(info_data);
            }
            Err(_) => panic!("Info length exceeds the allowable maximum!"),
        };

        trace!(self, "initlz::leave");
    }

    /// Processes the next message chunk, as given by the slice referenced by `chunk`.
    ///
    /// The internal state of the hash computation is updated by this function.
    pub fn update<T: AsRef<[u8]>>(&mut self, chunk: T) {
        trace!(self, "update::enter");

        for byte in chunk.as_ref() {
            self.state0[self.offset] ^= byte;
            self.offset += 1usize;

            if self.offset >= BLOCK_SIZE {
                self.permute();
                self.offset = 0usize;
            }
        }

        trace!(self, "update::leave");
    }

    /// Concludes the hash computation and returns the final digest.
    ///
    /// The hash value (digest) of the concatenation of all processed message chunks is returned as an new array of size `N`.
    ///
    /// **Note:** The digest output size `N`, in bytes, must be a *positive* value! &#x1F6A8;
    pub fn digest<const N: usize>(self) -> [u8; N] {
        let mut digest = [0u8; N];
        self.digest_to_slice(&mut digest);
        digest
    }

    /// Concludes the hash computation and returns the final digest.
    ///
    /// The hash value (digest) of the concatenation of all processed message chunks is written into the slice referenced by `digest`.
    ///
    /// **Note:** The digest output size `N`, in bytes, must be a *positive* value! &#x1F6A8;
    pub fn digest_to_slice<const N: usize>(mut self, digest_out: &mut [u8; N]) {
        let () = NoneZero::<N>::OK;
        debug_assert!(self.offset < BLOCK_SIZE);

        trace!(self, "digest::enter");

        self.state0[self.offset] ^= 0x80u8;
        let mut pos = 0usize;

        while pos < N {
            let copy_len = BLOCK_SIZE.min(N - pos);
            self.permute();
            digest_out[pos..(pos + copy_len)].copy_from_slice(&self.state0[..copy_len]);
            pos += copy_len;
        }

        trace!(self, "digest::leave");
    }

    /// Pseudorandom permutation, based on the AES-256 block cipher
    fn permute(&mut self) {
        trace!(self, "permfn::enter");

        let mut temp0 = [0u8; BLOCK_SIZE];
        let mut temp1 = [0u8; BLOCK_SIZE];
        let mut temp2 = [0u8; BLOCK_SIZE];

        for _ in 0..R {
            aes256_encrypt(&mut temp0, &self.state0, &self.state1, &self.state2);
            aes256_encrypt(&mut temp1, &self.state1, &self.state2, &self.state0);
            aes256_encrypt(&mut temp2, &self.state2, &self.state0, &self.state1);

            xor_arrays(&mut self.state0, &temp0);
            xor_arrays(&mut self.state1, &temp1);
            xor_arrays(&mut self.state2, &temp2);

            xor_arrays(&mut self.state1, &Self::BIT_MASK_0);
            xor_arrays(&mut self.state2, &Self::BIT_MASK_1);
        }

        temp0.zeroize();
        temp1.zeroize();
        temp2.zeroize();

        trace!(self, "permfn::leave");
    }
}

impl Default for SpongeHash256 {
    fn default() -> Self {
        Self::new()
    }
}

impl<const R: usize> Drop for SpongeHash256<R> {
    fn drop(&mut self) {
        self.state0.zeroize();
        self.state1.zeroize();
        self.state2.zeroize();
    }
}

// ---------------------------------------------------------------------------
// One-Shot API
// ---------------------------------------------------------------------------

/// Convenience function for “one-shot” SpongeHash-AES256 computation
///
/// The hash value (digest) of the given `message` is returned as an new array of size `N`.
///
/// Optionally, an additional `info` string may be specified.
///
/// **Note:** The digest output size `N`, in bytes, must be a *positive* value! &#x1F6A8;
///
/// ### Usage Example
///
/// The **`compute()`** function can be used as follows:
///
/// ```rust
/// use hex::encode_to_slice;
/// use sponge_hash_aes256::{DEFAULT_DIGEST_SIZE, compute};
///
/// fn main() {
///     // Compute digest using the “one-shot” function
///     let digest: [u8; DEFAULT_DIGEST_SIZE] =
///         compute(None, b"The quick brown fox jumps over the lazy dog");
///
///     // Encode to hex
///     let mut hex_buffer = [0u8; 2usize * DEFAULT_DIGEST_SIZE];
///     encode_to_slice(&digest, &mut hex_buffer).unwrap();
///
///     // Print the digest (hex format)
///     println!("0x{}", str::from_utf8(&hex_buffer).unwrap());
/// }
/// ```
/// &nbsp;
///
/// Optionally, additional “context” information may be provided via the `info` parameter:
///
/// ```rust
/// use sponge_hash_aes256::{DEFAULT_DIGEST_SIZE, compute};
///
/// fn main() {
///     // Compute digest using the “one-shot” function with “info”
///     let digest: [u8; DEFAULT_DIGEST_SIZE] =
///         compute(Some("my_application"), b"The quick brown fox jumps over the lazy dog");
///
///     /* ... */
/// }
/// ```
/// &nbsp;
///
/// <div class="warning">
///
/// Applications that need to process *large* messages are recommended to use the [streaming API](SpongeHash256), which does **not** require *all* message data to be held in memory at once and which allows for an *incremental* hash computation.
///
/// </div>
pub fn compute<const N: usize, T: AsRef<[u8]>>(info: Option<&str>, message: T) -> [u8; N] {
    assert!(!info.is_some_and(str::is_empty), "Info must not be empty!");
    let mut state: SpongeHash256 = SpongeHash256::with_info(info.unwrap_or_default());
    state.update(message);
    state.digest()
}

/// Convenience function for “one-shot” SpongeHash-AES256 computation
///
/// The hash value (digest) of the given `message` is written into the slice of size `N` referenced by `digest`.
///
/// Optionally, an additional `info` string may be specified.
///
/// **Note:** The digest output size `N`, in bytes, must be a *positive* value! &#x1F6A8;
///
/// ### Usage Example
///
/// The **`compute_to_slice()`** function can be used as follows:
///
/// ```rust
/// use hex::encode_to_slice;
/// use sponge_hash_aes256::{DEFAULT_DIGEST_SIZE, compute_to_slice};
///
/// fn main() {
///     // Compute digest using the “one-shot” function
///     let mut digest = [0u8; DEFAULT_DIGEST_SIZE];
///     compute_to_slice(&mut digest, None, b"The quick brown fox jumps over the lazy dog");
///
///     // Encode to hex
///     let mut hex_buffer = [0u8; 2usize * DEFAULT_DIGEST_SIZE];
///     encode_to_slice(&digest, &mut hex_buffer).unwrap();
///
///     // Print the digest (hex format)
///     println!("0x{}", str::from_utf8(&hex_buffer).unwrap());
/// }
/// ```
/// &nbsp;
///
/// <div class="warning">
///
/// Applications that need to process *large* messages are recommended to use the [streaming API](SpongeHash256), which does **not** require *all* message data to be held in memory at once and which allows for an *incremental* hash computation.
///
/// </div>
pub fn compute_to_slice<const N: usize, T: AsRef<[u8]>>(digest_out: &mut [u8; N], info: Option<&str>, message: T) {
    assert!(!info.is_some_and(str::is_empty), "Info must not be empty!");
    let mut state: SpongeHash256 = SpongeHash256::with_info(info.unwrap_or_default());
    state.update(message);
    state.digest_to_slice(digest_out);
}