krypteia-quantica 0.1.0

//! 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 |

/// 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 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
    }

    // ---------- 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);
    }
}