rns_crypto/
hkdf.rs

1use alloc::vec::Vec;
2use core::fmt;
3
4use crate::hmac::{hmac_sha256, HmacSha256};
5
6#[derive(Debug, PartialEq)]
7pub enum HkdfError {
8    InvalidLength,
9    EmptyInput,
10}
11
12impl fmt::Display for HkdfError {
13    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
14        match self {
15            HkdfError::InvalidLength => write!(f, "Invalid output key length"),
16            HkdfError::EmptyInput => write!(f, "Cannot derive key from empty input material"),
17        }
18    }
19}
20
21/// Custom HKDF implementation matching RNS/Cryptography/HKDF.py.
22/// WARNING: This is NOT RFC 5869. The counter wraps modulo 256.
23pub fn hkdf(
24    length: usize,
25    derive_from: &[u8],
26    salt: Option<&[u8]>,
27    context: Option<&[u8]>,
28) -> Result<Vec<u8>, HkdfError> {
29    let hash_len: usize = 32;
30
31    if length < 1 {
32        return Err(HkdfError::InvalidLength);
33    }
34
35    if derive_from.is_empty() {
36        return Err(HkdfError::EmptyInput);
37    }
38
39    let salt = match salt {
40        Some(s) if !s.is_empty() => s.to_vec(),
41        _ => alloc::vec![0u8; hash_len],
42    };
43
44    let context = context.unwrap_or(b"");
45
46    // Extract
47    let prk = hmac_sha256(&salt, derive_from);
48
49    // Expand
50    let expansion_hmac = HmacSha256::new(&prk);
51    let mut block = [0u8; 32];
52    let mut has_block = false;
53    let mut derived = Vec::with_capacity(length);
54
55    let iterations = length.div_ceil(hash_len);
56    for i in 0..iterations {
57        let mut hmac = expansion_hmac.clone();
58        if has_block {
59            hmac.update(&block);
60        }
61        hmac.update(context);
62        hmac.update(&[((i + 1) % 256) as u8]);
63        block = hmac.finalize();
64        has_block = true;
65        derived.extend_from_slice(&block);
66    }
67
68    derived.truncate(length);
69    Ok(derived)
70}
71
72// Cargo discovers this module through the crate test harness, so the HKDF
73// vectors run under `cargo test --workspace` without a manual runner registry.
74#[cfg(test)]
75mod tests {
76    use super::*;
77
78    fn decode_hex(value: &str) -> Vec<u8> {
79        value
80            .as_bytes()
81            .chunks_exact(2)
82            .map(|pair| {
83                let digit = |byte: u8| match byte {
84                    b'0'..=b'9' => byte - b'0',
85                    b'a'..=b'f' => byte - b'a' + 10,
86                    _ => panic!("invalid test vector"),
87                };
88                (digit(pair[0]) << 4) | digit(pair[1])
89            })
90            .collect()
91    }
92
93    fn reference_hkdf(length: usize, ikm: &[u8], salt: &[u8], context: &[u8]) -> Vec<u8> {
94        let prk = hmac_sha256(salt, ikm);
95        let mut block = Vec::new();
96        let mut derived = Vec::with_capacity(length);
97        for index in 0..length.div_ceil(32) {
98            let mut input = block;
99            input.extend_from_slice(context);
100            input.push(((index + 1) % 256) as u8);
101            block = hmac_sha256(&prk, &input).to_vec();
102            derived.extend_from_slice(&block);
103        }
104        derived.truncate(length);
105        derived
106    }
107
108    #[test]
109    fn optimized_expansion_matches_reference_across_counter_wrap() {
110        let ikm = b"deterministic input material";
111        let salt = b"deterministic salt";
112        let context = b"counter-wrap";
113        for length in [8191, 8192, 8193] {
114            assert_eq!(
115                hkdf(length, ikm, Some(salt), Some(context)).unwrap(),
116                reference_hkdf(length, ikm, salt, context)
117            );
118        }
119    }
120
121    #[test]
122    fn rfc5869_sha256_vectors() {
123        let cases = [
124            (
125                vec![0x0b; 22],
126                (0x00..0x0d).collect::<Vec<_>>(),
127                (0xf0..0xfa).collect::<Vec<_>>(),
128                42,
129                "3cb25f25faacd57a90434f64d0362f2a2d2d0a90cf1a5a4c5db02d56ecc4c5bf34007208d5b887185865",
130            ),
131            (
132                (0x00..0x50).collect::<Vec<_>>(),
133                (0x60..0xb0).collect::<Vec<_>>(),
134                (0xb0..=0xff).collect::<Vec<_>>(),
135                82,
136                "b11e398dc80327a1c8e7f78c596a49344f012eda2d4efad8a050cc4c19afa97c59045a99cac7827271cb41c65e590e09da3275600c2f09b8367793a9aca3db71cc30c58179ec3e87c14c01d5c1f3434f1d87",
137            ),
138            (
139                vec![0x0b; 22],
140                Vec::new(),
141                Vec::new(),
142                42,
143                "8da4e775a563c18f715f802a063c5a31b8a11f5c5ee1879ec3454e5f3c738d2d9d201395faa4b61a96c8",
144            ),
145        ];
146
147        for (ikm, salt, context, length, expected) in cases {
148            assert_eq!(
149                hkdf(length, &ikm, Some(&salt), Some(&context)).unwrap(),
150                decode_hex(expected)
151            );
152        }
153    }
154
155    #[test]
156    fn test_hkdf_32bytes() {
157        let ikm = b"input key material";
158        let salt = b"salt value";
159        let result = hkdf(32, ikm, Some(salt), None).unwrap();
160        assert_eq!(result.len(), 32);
161    }
162
163    #[test]
164    fn test_hkdf_64bytes() {
165        let ikm = b"input key material";
166        let salt = b"salt value";
167        let result = hkdf(64, ikm, Some(salt), None).unwrap();
168        assert_eq!(result.len(), 64);
169    }
170
171    #[test]
172    fn test_hkdf_with_context() {
173        let ikm = b"input key material";
174        let salt = b"salt";
175        let ctx = b"context info";
176        let result = hkdf(32, ikm, Some(salt), Some(ctx)).unwrap();
177        assert_eq!(result.len(), 32);
178        // With context should differ from without
179        let result2 = hkdf(32, ikm, Some(salt), None).unwrap();
180        assert_ne!(result, result2);
181    }
182
183    #[test]
184    fn test_hkdf_none_salt() {
185        let ikm = b"input key material";
186        let result = hkdf(32, ikm, None, None).unwrap();
187        assert_eq!(result.len(), 32);
188    }
189
190    #[test]
191    fn test_hkdf_empty_salt() {
192        let ikm = b"input key material";
193        let result1 = hkdf(32, ikm, Some(b""), None).unwrap();
194        let result2 = hkdf(32, ikm, None, None).unwrap();
195        // Empty salt and None salt should produce same result
196        assert_eq!(result1, result2);
197    }
198
199    #[test]
200    fn test_hkdf_invalid_length() {
201        assert_eq!(hkdf(0, b"ikm", None, None), Err(HkdfError::InvalidLength));
202    }
203
204    #[test]
205    fn test_hkdf_empty_ikm() {
206        assert_eq!(hkdf(32, b"", None, None), Err(HkdfError::EmptyInput));
207    }
208}
rns_crypto/hkdf.rs

rns_crypto/
hkdf.rs
