tightbeam-rs 0.8.0

//! Multi-input key derivation functions for hybrid key agreement.
//!
//! Provides composable KDF primitives for protocols that combine multiple
//! shared secrets (e.g., ECDH + KEM in PQXDH).

#[cfg(not(feature = "std"))]
extern crate alloc;

#[cfg(not(feature = "std"))]
use alloc::vec::Vec;

use crate::crypto::kdf::KdfFunction;
use crate::crypto::profiles::CryptoProvider;
use crate::transport::handshake::error::HandshakeError;
use crate::zeroize::Zeroizing;

/// Multi-input HKDF combining multiple secrets with length prefixing.
///
/// Concatenates all input secrets with 4-byte big-endian length prefixes,
/// then derives a `key_size`-byte key using the provider's KDF.
///
/// # Length Prefixing
/// Each input is prefixed with its length as `u32` in big-endian format.
/// This prevents ambiguity in concatenation (e.g., "ab" + "cd" vs "abc" + "d").
///
/// # Parameters
/// - `inputs`: Array of input secret slices to combine
/// - `salt`: Salt bytes for KDF (minimum 16 bytes recommended)
/// - `info`: Application-specific context string
/// - `key_size`: Desired output key length in bytes (e.g. KEK size from the
///   negotiated key-wrap algorithm)
///
/// # Returns
/// Zeroizing vector containing the derived key
///
/// # Example
/// ```rust,ignore
/// let dh_secret = [0x42u8; 32];
/// let kem_secret = [0x99u8; 32];
/// let salt = [0xAAu8; 32];
///
/// let combined = multi_input_kdf::<DefaultCryptoProvider>(
///     &[&dh_secret, &kem_secret],
///     &salt,
///     b"MyApp-PQXDH-v1",
///     32,
/// )?;
/// ```
pub fn multi_input_kdf<P: CryptoProvider>(
	inputs: &[&[u8]],
	salt: &[u8],
	info: &[u8],
	key_size: usize,
) -> Result<Zeroizing<Vec<u8>>, HandshakeError> {
	// Concatenate all inputs with length prefixes. The buffer holds secret
	// material, so it is zeroized on drop.
	let mut combined = Zeroizing::new(Vec::new());
	for input in inputs {
		let len = u32::try_from(input.len()).map_err(|_| HandshakeError::IntegerOutOfRange)?;
		combined.extend_from_slice(&len.to_be_bytes());
		combined.extend_from_slice(input);
	}

	Ok(P::Kdf::derive_dynamic_key(&combined, info, Some(salt), key_size)?)
}

/// Chain multiple KDF operations where each output becomes the next salt.
///
/// Implements key derivation chaining as used in protocols like PQXDH:
/// ```text
/// KDF₁(input₁, salt₀, info₁) → output₁
/// KDF₂(input₂, output₁, info₂) → output₂
/// ...
/// ```
///
/// Each stage uses the previous stage's output as salt, creating a
/// dependency chain that incorporates all prior inputs.
///
/// # Parameters
/// - `stages`: Array of (input, info) pairs for each KDF stage
/// - `initial_salt`: Initial salt for first KDF stage
///
/// # Returns
/// Final derived key (32 bytes)
///
/// # Example
/// ```rust,ignore
/// let dh1 = [0x11u8; 32];
/// let dh2 = [0x22u8; 32];
/// let dh3 = [0x33u8; 32];
/// let kem_ss = [0x44u8; 32];
///
/// let final_key = kdf_chain::<DefaultCryptoProvider>(
///     &[
///         (&dh1, b"DH1"),
///         (&dh2, b"DH2"),
///         (&dh3, b"DH3"),
///         (&kem_ss, b"KEM"),
///     ],
///     b"InitialSalt"
/// )?;
/// ```
pub fn kdf_chain<P: CryptoProvider>(stages: &[(&[u8], &[u8])], initial_salt: &[u8]) -> Result<Vec<u8>, HandshakeError> {
	let mut current_salt = initial_salt.to_vec();
	let mut current_key = Vec::new();
	for (input, info) in stages {
		let derived = P::Kdf::derive_dynamic_key(input, info, Some(&current_salt), 32)?;
		current_key = derived.to_vec();
		current_salt = derived.to_vec(); // Next salt is previous output
	}

	Ok(current_key)
}

#[cfg(test)]
mod tests {
	use super::*;
	use crate::crypto::profiles::DefaultCryptoProvider;

	#[test]
	fn test_multi_input_kdf() -> Result<(), Box<dyn core::error::Error>> {
		let input1 = [0x42u8; 32];
		let input2 = [0x99u8; 32];
		let salt = [0xAAu8; 32];

		let result = multi_input_kdf::<DefaultCryptoProvider>(&[&input1, &input2], &salt, b"test-info", 32);
		assert!(result.is_ok());

		let key = result?;
		assert_eq!(key.len(), 32);

		Ok(())
	}

	#[test]
	fn test_multi_input_kdf_different_lengths() {
		let input1 = [0x42u8; 16];
		let input2 = [0x99u8; 48];
		let salt = [0xAAu8; 32];

		let result = multi_input_kdf::<DefaultCryptoProvider>(&[&input1, &input2], &salt, b"test-info", 32);
		assert!(result.is_ok());
	}

	#[test]
	fn test_kdf_chain() -> Result<(), Box<dyn core::error::Error>> {
		let input1 = [0x11u8; 32];
		let input2 = [0x22u8; 32];
		let initial_salt = [0xFFu8; 32];

		let result = kdf_chain::<DefaultCryptoProvider>(&[(&input1, b"stage1"), (&input2, b"stage2")], &initial_salt);
		assert!(result.is_ok());

		let key = result?;
		assert_eq!(key.len(), 32);

		Ok(())
	}

	#[test]
	fn test_kdf_chain_single_stage() {
		let input = [0x42u8; 32];
		let salt = [0xAAu8; 32];

		let result = kdf_chain::<DefaultCryptoProvider>(&[(&input, b"single")], &salt);
		assert!(result.is_ok());
	}
}