Skip to main content

hanzo_pqc/
kdf.rs

1//! Key Derivation Functions (KDF) 
2//! SP 800-56C compliant HKDF and SP 800-108 compliant KDF
3
4use hkdf::Hkdf;
5use sha2::{Sha256, Sha384, Sha512};
6use sha3::{Sha3_256, Sha3_384, Sha3_512};
7use serde::{Deserialize, Serialize};
8use crate::{PqcError, Result};
9
10/// KDF algorithms (SP 800-56C and SP 800-108 compliant)
11#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
12pub enum KdfAlgorithm {
13    /// HKDF with SHA-256 (128-bit security)
14    HkdfSha256,
15    /// HKDF with SHA-384 (192-bit security) - RECOMMENDED for ML-KEM-768
16    HkdfSha384,
17    /// HKDF with SHA-512 (256-bit security)
18    HkdfSha512,
19    /// HKDF with SHA3-256 (128-bit security)
20    HkdfSha3_256,
21    /// HKDF with SHA3-384 (192-bit security)
22    HkdfSha3_384,
23    /// HKDF with SHA3-512 (256-bit security)
24    HkdfSha3_512,
25    /// BLAKE3 KDF (256-bit security)
26    Blake3Kdf,
27}
28
29impl Default for KdfAlgorithm {
30    fn default() -> Self {
31        Self::HkdfSha384 // Matches ML-KEM-768 security level
32    }
33}
34
35/// KDF trait for key derivation operations
36pub trait Kdf {
37    /// Extract a pseudorandom key from input keying material
38    fn extract(&self, salt: Option<&[u8]>, ikm: &[u8]) -> Vec<u8>;
39    
40    /// Expand a pseudorandom key to desired length
41    fn expand(&self, prk: &[u8], info: &[u8], okm_len: usize) -> Result<Vec<u8>>;
42    
43    /// Combined extract-and-expand operation
44    fn derive(&self, salt: Option<&[u8]>, ikm: &[u8], info: &[u8], okm_len: usize) -> Result<Vec<u8>>;
45}
46
47/// Generic HKDF implementation
48pub struct HkdfKdf {
49    algorithm: KdfAlgorithm,
50}
51
52impl HkdfKdf {
53    pub fn new(algorithm: KdfAlgorithm) -> Self {
54        Self { algorithm }
55    }
56}
57
58impl Kdf for HkdfKdf {
59    fn extract(&self, salt: Option<&[u8]>, ikm: &[u8]) -> Vec<u8> {
60        match self.algorithm {
61            KdfAlgorithm::HkdfSha256 => {
62                let (prk, _) = Hkdf::<Sha256>::extract(salt, ikm);
63                prk.to_vec()
64            }
65            KdfAlgorithm::HkdfSha384 => {
66                let (prk, _) = Hkdf::<Sha384>::extract(salt, ikm);
67                prk.to_vec()
68            }
69            KdfAlgorithm::HkdfSha512 => {
70                let (prk, _) = Hkdf::<Sha512>::extract(salt, ikm);
71                prk.to_vec()
72            }
73            KdfAlgorithm::HkdfSha3_256 => {
74                let (prk, _) = Hkdf::<Sha3_256>::extract(salt, ikm);
75                prk.to_vec()
76            }
77            KdfAlgorithm::HkdfSha3_384 => {
78                let (prk, _) = Hkdf::<Sha3_384>::extract(salt, ikm);
79                prk.to_vec()
80            }
81            KdfAlgorithm::HkdfSha3_512 => {
82                let (prk, _) = Hkdf::<Sha3_512>::extract(salt, ikm);
83                prk.to_vec()
84            }
85            KdfAlgorithm::Blake3Kdf => {
86                // BLAKE3 has its own KDF mode
87                let key = blake3::derive_key(
88                    salt.map(|s| std::str::from_utf8(s).unwrap_or("hanzo-pqc")).unwrap_or("hanzo-pqc"),
89                    ikm,
90                );
91                key.to_vec()
92            }
93        }
94    }
95    
96    fn expand(&self, prk: &[u8], info: &[u8], okm_len: usize) -> Result<Vec<u8>> {
97        let mut okm = vec![0u8; okm_len];
98        
99        match self.algorithm {
100            KdfAlgorithm::HkdfSha256 => {
101                let hk = Hkdf::<Sha256>::from_prk(prk)
102                    .map_err(|_| PqcError::KdfError("Invalid PRK length for SHA256".into()))?;
103                hk.expand(info, &mut okm)
104                    .map_err(|_| PqcError::KdfError("HKDF expand failed".into()))?;
105            }
106            KdfAlgorithm::HkdfSha384 => {
107                let hk = Hkdf::<Sha384>::from_prk(prk)
108                    .map_err(|_| PqcError::KdfError("Invalid PRK length for SHA384".into()))?;
109                hk.expand(info, &mut okm)
110                    .map_err(|_| PqcError::KdfError("HKDF expand failed".into()))?;
111            }
112            KdfAlgorithm::HkdfSha512 => {
113                let hk = Hkdf::<Sha512>::from_prk(prk)
114                    .map_err(|_| PqcError::KdfError("Invalid PRK length for SHA512".into()))?;
115                hk.expand(info, &mut okm)
116                    .map_err(|_| PqcError::KdfError("HKDF expand failed".into()))?;
117            }
118            KdfAlgorithm::HkdfSha3_256 => {
119                let hk = Hkdf::<Sha3_256>::from_prk(prk)
120                    .map_err(|_| PqcError::KdfError("Invalid PRK length for SHA3-256".into()))?;
121                hk.expand(info, &mut okm)
122                    .map_err(|_| PqcError::KdfError("HKDF expand failed".into()))?;
123            }
124            KdfAlgorithm::HkdfSha3_384 => {
125                let hk = Hkdf::<Sha3_384>::from_prk(prk)
126                    .map_err(|_| PqcError::KdfError("Invalid PRK length for SHA3-384".into()))?;
127                hk.expand(info, &mut okm)
128                    .map_err(|_| PqcError::KdfError("HKDF expand failed".into()))?;
129            }
130            KdfAlgorithm::HkdfSha3_512 => {
131                let hk = Hkdf::<Sha3_512>::from_prk(prk)
132                    .map_err(|_| PqcError::KdfError("Invalid PRK length for SHA3-512".into()))?;
133                hk.expand(info, &mut okm)
134                    .map_err(|_| PqcError::KdfError("HKDF expand failed".into()))?;
135            }
136            KdfAlgorithm::Blake3Kdf => {
137                // BLAKE3 XOF mode for expansion
138                let mut hasher = blake3::Hasher::new_keyed(
139                    &<[u8; 32]>::try_from(&prk[..32])
140                        .map_err(|_| PqcError::KdfError("BLAKE3 requires 32-byte key".into()))?
141                );
142                hasher.update(info);
143                let mut output = hasher.finalize_xof();
144                output.fill(&mut okm);
145            }
146        }
147        
148        Ok(okm)
149    }
150    
151    fn derive(&self, salt: Option<&[u8]>, ikm: &[u8], info: &[u8], okm_len: usize) -> Result<Vec<u8>> {
152        let prk = self.extract(salt, ikm);
153        self.expand(&prk, info, okm_len)
154    }
155}
156
157/// Combine multiple shared secrets (for hybrid mode)
158/// Per SP 800-56C Rev 2, Section 5.9.3
159pub fn combine_shared_secrets(
160    kdf: &impl Kdf,
161    secrets: &[&[u8]],
162    context: &[u8],
163    output_len: usize,
164) -> Result<Vec<u8>> {
165    // Concatenate all secrets with length prefixes
166    let mut combined = Vec::new();
167    for secret in secrets {
168        combined.extend_from_slice(&(secret.len() as u32).to_be_bytes());
169        combined.extend_from_slice(secret);
170    }
171    
172    // Derive final key material with context
173    kdf.derive(None, &combined, context, output_len)
174}
175
176/// Domain separation for different protocol contexts
177pub fn domain_separate(
178    kdf: &impl Kdf,
179    key_material: &[u8],
180    domain: &str,
181    output_len: usize,
182) -> Result<Vec<u8>> {
183    let info = format!("hanzo-pqc-v1|{domain}");
184    kdf.expand(key_material, info.as_bytes(), output_len)
185}
186
187#[cfg(test)]
188mod tests {
189    use super::*;
190    
191    #[test]
192    fn test_hkdf_sha384() {
193        let kdf = HkdfKdf::new(KdfAlgorithm::HkdfSha384);
194        
195        let ikm = b"input keying material";
196        let salt = b"salt";
197        let info = b"hanzo-test-v1";
198        
199        let okm = kdf.derive(Some(salt), ikm, info, 64).unwrap();
200        assert_eq!(okm.len(), 64);
201        
202        // Verify deterministic
203        let okm2 = kdf.derive(Some(salt), ikm, info, 64).unwrap();
204        assert_eq!(okm, okm2);
205    }
206    
207    #[test]
208    fn test_combine_secrets() {
209        let kdf = HkdfKdf::new(KdfAlgorithm::HkdfSha384);
210        
211        let secret1 = vec![1u8; 32]; // ML-KEM shared secret
212        let secret2 = vec![2u8; 32]; // X25519 shared secret
213        
214        let combined = combine_shared_secrets(
215            &kdf,
216            &[&secret1, &secret2],
217            b"hanzo-hybrid-v1",
218            48,
219        ).unwrap();
220        
221        assert_eq!(combined.len(), 48);
222    }
223}