krypteia-quantica 0.2.0

// SPDX-License-Identifier: Apache-2.0
// Copyright 2026 Cédric Mesnil <cslashm@pm.me>

//! ADRS (Address) structure for SLH-DSA (FIPS 205, Section 4.2).
//!
//! The address structure is a 32-byte value organized as 8 big-endian `u32` words.
//! It is used to domain-separate every hash call in SLH-DSA so that distinct positions
//! in the tree hierarchy produce independent hash outputs.
//!
//! # Layout
//!
//! | Bytes     | Field             | Description                       |
//! |-----------|-------------------|-----------------------------------|
//! | `[0..4]`  | Layer address     | Hypertree layer (0 = bottom)      |
//! | `[4..16]` | Tree address      | Index of the XMSS tree (3 words)  |
//! | `[16..20]`| Type              | Address type (see constants below)|
//! | `[20..32]`| Type-specific     | Meaning depends on the type field |
//!
//! # Stability (hazmat tier)
//!
//! This module is part of the `hazmat` expert building-block layer: it is
//! public only with the `hazmat` Cargo feature, for KAT/test-vector tooling
//! and research. Read each function's invariants before use. **No stability,
//! misuse-resistance, or side-channel promise** applies to this tier — its
//! API may change in any release. The typed facade is the guaranteed API.

/// Address type for WOTS+ hash chain evaluation.
pub const WOTS_HASH: u32 = 0;
/// Address type for WOTS+ public key compression.
pub const WOTS_PK: u32 = 1;
/// Address type for XMSS/hypertree internal Merkle tree nodes.
pub const TREE: u32 = 2;
/// Address type for FORS tree nodes.
pub const FORS_TREE: u32 = 3;
/// Address type for FORS root compression (combining k roots into one public key).
pub const FORS_ROOTS: u32 = 4;
/// Address type for WOTS+ secret key generation via PRF.
pub const WOTS_PRF: u32 = 5;
/// Address type for FORS secret key generation via PRF.
pub const FORS_PRF: u32 = 6;

/// A 32-byte address structure used to domain-separate hash function calls in SLH-DSA.
///
/// Every hash invocation in SLH-DSA includes an `Adrs` value that encodes the position
/// within the signing hierarchy (layer, tree, leaf, chain step, etc.). This ensures that
/// hash outputs at different positions are cryptographically independent.
///
/// The type-specific fields (bytes 20..32) have different interpretations depending on
/// the address type:
///
/// - **WOTS_HASH / WOTS_PK / WOTS_PRF**: key pair address, chain address, hash address
/// - **TREE**: padding (0), tree height, tree index
/// - **FORS_TREE / FORS_ROOTS / FORS_PRF**: key pair address, tree height, tree index
#[derive(Clone, Debug)]
pub struct Adrs {
    data: [u8; 32],
}

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

impl Adrs {
    /// Create a new address with all fields initialized to zero.
    pub fn new() -> Self {
        Self { data: [0u8; 32] }
    }

    /// Get a reference to the raw 32-byte address value.
    ///
    /// This is passed directly to hash functions as part of their input.
    pub fn as_bytes(&self) -> &[u8; 32] {
        &self.data
    }

    /// FIPS 205 §11.2 compressed address `ADRS^c` (22 bytes), used by the
    /// SHA-2 instantiation in place of the full 32-byte address.
    ///
    /// Keeps the low byte of the layer address, the 8-byte tree address,
    /// the low byte of the type, and the 12 type-specific bytes (words
    /// 5/6/7) — the three high bytes of the layer and type words and the
    /// four high bytes of the tree address (always zero for the supported
    /// `u64` tree index range) are dropped. Offsets are compile-time
    /// constants over public address data: no secret-dependent access.
    pub fn as_bytes_compressed(&self) -> [u8; 22] {
        let mut c = [0u8; 22];
        c[0] = self.data[3]; // layer address, low byte (word 0)
        c[1..9].copy_from_slice(&self.data[8..16]); // tree address (words 2-3, the u64)
        c[9] = self.data[19]; // type, low byte (word 4)
        c[10..22].copy_from_slice(&self.data[20..32]); // type-specific (words 5-7)
        c
    }

    // ---------- Word-level helpers ----------

    fn get_word(&self, offset: usize) -> u32 {
        u32::from_be_bytes([
            self.data[offset],
            self.data[offset + 1],
            self.data[offset + 2],
            self.data[offset + 3],
        ])
    }

    fn set_word(&mut self, offset: usize, val: u32) {
        let bytes = val.to_be_bytes();
        self.data[offset..offset + 4].copy_from_slice(&bytes);
    }

    // ---------- Layer address (bytes 0..4) ----------

    /// Get the hypertree layer address (0 = bottom layer).
    pub fn get_layer_address(&self) -> u32 {
        self.get_word(0)
    }

    /// Set the hypertree layer address.
    ///
    /// Layer 0 is the bottom of the hypertree (closest to the FORS trees);
    /// layer `d - 1` is the top.
    pub fn set_layer_address(&mut self, val: u32) {
        self.set_word(0, val);
    }

    // ---------- Tree address (bytes 4..16, 12 bytes = 3 words) ----------

    /// Set the tree address from a `u64` index.
    ///
    /// The tree address occupies bytes 4..16 (words 1, 2, 3). The `u64` value is stored
    /// in the lower 8 bytes (words 2 and 3); the upper 4 bytes (word 1) are zeroed.
    /// This identifies which XMSS tree within the current layer is being addressed.
    pub fn set_tree_address(&mut self, val: u64) {
        // Word 1 (bytes 4..8): upper 32 bits beyond u64 range, set to 0.
        self.set_word(4, 0);
        // Word 2 (bytes 8..12): high 32 bits of val.
        self.set_word(8, (val >> 32) as u32);
        // Word 3 (bytes 12..16): low 32 bits of val.
        self.set_word(12, val as u32);
    }

    /// Get the tree address as a `u64` (lower 8 bytes of the 12-byte field).
    pub fn get_tree_address(&self) -> u64 {
        let hi = self.get_word(8) as u64;
        let lo = self.get_word(12) as u64;
        (hi << 32) | lo
    }

    // ---------- Type (bytes 16..20) ----------

    /// Set the address type and zero all type-specific fields (bytes 20..32).
    ///
    /// This must be called when switching address types to ensure leftover values
    /// from a previous type do not leak into the new context. Use one of the
    /// address type constants: [`WOTS_HASH`], [`WOTS_PK`], [`TREE`], [`FORS_TREE`],
    /// [`FORS_ROOTS`], [`WOTS_PRF`], or [`FORS_PRF`].
    pub fn set_type_and_clear(&mut self, addr_type: u32) {
        self.set_word(16, addr_type);
        // Zero out type-specific fields (bytes 20..32).
        for i in 20..32 {
            self.data[i] = 0;
        }
    }

    /// Get the current address type.
    pub fn get_type(&self) -> u32 {
        self.get_word(16)
    }

    // ---------- Type-specific fields (bytes 20..32) ----------
    // For WOTS_HASH / WOTS_PK / WOTS_PRF:
    //   word 5 (bytes 20..24): key pair address
    //   word 6 (bytes 24..28): chain address
    //   word 7 (bytes 28..32): hash address
    // For TREE:
    //   word 5 (bytes 20..24): padding (0)
    //   word 6 (bytes 24..28): tree height
    //   word 7 (bytes 28..32): tree index
    // For FORS_TREE / FORS_ROOTS / FORS_PRF:
    //   word 5 (bytes 20..24): key pair address
    //   word 6 (bytes 24..28): tree height
    //   word 7 (bytes 28..32): tree index

    /// Get the key pair address (word 5, bytes 20..24).
    ///
    /// Used by WOTS+ and FORS address types to identify which leaf key pair is being
    /// operated on within the current XMSS tree.
    pub fn get_key_pair_address(&self) -> u32 {
        self.get_word(20)
    }

    /// Set the key pair address (word 5, bytes 20..24).
    pub fn set_key_pair_address(&mut self, val: u32) {
        self.set_word(20, val);
    }

    /// Get the WOTS+ chain address (word 6, bytes 24..28).
    ///
    /// Identifies which of the `len` hash chains within a WOTS+ instance is being computed.
    pub fn get_chain_address(&self) -> u32 {
        self.get_word(24)
    }

    /// Set the WOTS+ chain address (word 6, bytes 24..28).
    pub fn set_chain_address(&mut self, val: u32) {
        self.set_word(24, val);
    }

    /// Get the WOTS+ hash address (word 7, bytes 28..32).
    ///
    /// Identifies the step within a WOTS+ hash chain (0 to `w - 2`).
    pub fn get_hash_address(&self) -> u32 {
        self.get_word(28)
    }

    /// Set the WOTS+ hash address (word 7, bytes 28..32).
    pub fn set_hash_address(&mut self, val: u32) {
        self.set_word(28, val);
    }

    /// Get the Merkle tree height (word 6, bytes 24..28).
    ///
    /// Used by TREE and FORS_TREE address types. Height 0 corresponds to leaves;
    /// increasing heights move toward the root.
    pub fn get_tree_height(&self) -> u32 {
        self.get_word(24)
    }

    /// Set the Merkle tree height (word 6, bytes 24..28).
    pub fn set_tree_height(&mut self, val: u32) {
        self.set_word(24, val);
    }

    /// Get the Merkle tree node index (word 7, bytes 28..32).
    ///
    /// Identifies the node's horizontal position within its tree level.
    pub fn get_tree_index(&self) -> u32 {
        self.get_word(28)
    }

    /// Set the Merkle tree node index (word 7, bytes 28..32).
    pub fn set_tree_index(&mut self, val: u32) {
        self.set_word(28, val);
    }
}

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

    // Lock the FIPS 205 §11.2 ADRS^c slice offsets: layer low byte,
    // 8-byte tree address, type low byte, 12 type-specific bytes.
    #[test]
    fn compressed_address_layout() {
        let mut a = Adrs::new();
        a.set_layer_address(0x11); // word 0 -> bytes [00 00 00 11]; low byte 0x11 at data[3]
        a.set_tree_address(0x2233_4455_6677_8899); // words 2-3 (data[8..16]); word1 (data[4..8]) stays 0
        a.set_type_and_clear(0x06); // word 4 -> bytes[16..20]=[00 00 00 06]; low byte 0x06 at data[19]
        a.set_key_pair_address(0xaabb_ccdd); // word 5 -> data[20..24]
        a.set_tree_height(0x01020304); // word 6 -> data[24..28]
        a.set_tree_index(0x05060708); // word 7 -> data[28..32]

        let c = a.as_bytes_compressed();
        assert_eq!(c.len(), 22);
        assert_eq!(c[0], 0x11); // layer low byte
        assert_eq!(&c[1..9], &0x2233_4455_6677_8899u64.to_be_bytes()); // tree address
        assert_eq!(c[9], 0x06); // type low byte
        assert_eq!(&c[10..14], &0xaabb_ccddu32.to_be_bytes()); // key-pair address
        assert_eq!(&c[14..18], &0x0102_0304u32.to_be_bytes()); // tree height
        assert_eq!(&c[18..22], &0x0506_0708u32.to_be_bytes()); // tree index
    }
}