1use crypto_core::{CryptoError, HkdfFailureKind, HkdfHash};
6use hkdf::Hkdf;
7use sha2::Sha256;
8#[cfg(feature = "sha3")]
9use sha3::{Digest, Sha3_256};
10#[cfg(feature = "sha3")]
11use zeroize::{Zeroize, Zeroizing};
12
13use crate::material::{HkdfInfo, HkdfInputKeyMaterial, HkdfOutput, HkdfSalt};
14use crate::policy::{DomainKeyPurpose, DomainTag, HkdfSuite};
15
16#[cfg(feature = "sha3")]
17const SHA3_256_DIGEST_LEN: usize = 32;
18#[cfg(feature = "sha3")]
19const SHA3_256_HMAC_BLOCK_LEN: usize = 136;
20#[cfg(feature = "sha3")]
21const HKDF_MAX_BLOCKS: usize = 255;
22
23pub struct DeriveRequest<'a> {
25 pub suite: HkdfSuite,
27 pub ikm: &'a HkdfInputKeyMaterial,
29 pub salt: Option<&'a HkdfSalt>,
31 pub info: &'a HkdfInfo,
33}
34
35#[cfg(feature = "sha3")]
36fn hmac_sha3_256(key: &[u8], message: &[u8]) -> [u8; SHA3_256_DIGEST_LEN] {
37 let mut key_block = Zeroizing::new([0u8; SHA3_256_HMAC_BLOCK_LEN]);
38
39 if key.len() > SHA3_256_HMAC_BLOCK_LEN {
40 let hashed_key = Sha3_256::digest(key);
41 key_block[..SHA3_256_DIGEST_LEN].copy_from_slice(&hashed_key);
42 } else {
43 key_block[..key.len()].copy_from_slice(key);
44 }
45
46 let mut inner_pad = Zeroizing::new([0x36u8; SHA3_256_HMAC_BLOCK_LEN]);
47 let mut outer_pad = Zeroizing::new([0x5cu8; SHA3_256_HMAC_BLOCK_LEN]);
48 for index in 0..SHA3_256_HMAC_BLOCK_LEN {
49 inner_pad[index] ^= key_block[index];
50 outer_pad[index] ^= key_block[index];
51 }
52
53 let mut inner = Sha3_256::new();
54 inner.update(&inner_pad[..]);
55 inner.update(message);
56 let inner_digest = inner.finalize();
57
58 let mut outer = Sha3_256::new();
59 outer.update(&outer_pad[..]);
60 outer.update(inner_digest.as_slice());
61 let outer_digest = outer.finalize();
62
63 let mut tag = [0u8; SHA3_256_DIGEST_LEN];
64 tag.copy_from_slice(&outer_digest);
65 tag
66}
67
68#[cfg(feature = "sha3")]
69fn derive_sha3_256<const N: usize>(request: &DeriveRequest<'_>) -> Result<[u8; N], CryptoError> {
70 let block_count = N
71 .checked_add(SHA3_256_DIGEST_LEN - 1)
72 .map(|value| value / SHA3_256_DIGEST_LEN)
73 .ok_or(CryptoError::Hkdf {
74 hash: HkdfHash::Sha3_256,
75 kind: HkdfFailureKind::InvalidOutputLength,
76 })?;
77 if block_count > HKDF_MAX_BLOCKS {
78 return Err(CryptoError::Hkdf {
79 hash: HkdfHash::Sha3_256,
80 kind: HkdfFailureKind::InvalidOutputLength,
81 });
82 }
83
84 let zero_salt = [0u8; SHA3_256_DIGEST_LEN];
85 let salt = match request.salt {
86 Some(salt) => salt.as_bytes(),
87 None => &zero_salt,
88 };
89 let mut prk = Zeroizing::new(hmac_sha3_256(salt, request.ikm.as_bytes()));
90 let mut output = [0u8; N];
91 let mut previous_block = Zeroizing::new(Vec::<u8>::new());
92 let mut written = 0usize;
93
94 for block_index in 1..=block_count {
95 let capacity = previous_block
96 .len()
97 .checked_add(request.info.as_bytes().len())
98 .and_then(|value| value.checked_add(1))
99 .ok_or(CryptoError::Hkdf {
100 hash: HkdfHash::Sha3_256,
101 kind: HkdfFailureKind::LengthOverflow,
102 })?;
103 let mut block_input = Zeroizing::new(Vec::<u8>::with_capacity(capacity));
104 block_input.extend_from_slice(&previous_block);
105 block_input.extend_from_slice(request.info.as_bytes());
106 let counter = u8::try_from(block_index).map_err(|_| CryptoError::Hkdf {
107 hash: HkdfHash::Sha3_256,
108 kind: HkdfFailureKind::InvalidOutputLength,
109 })?;
110 block_input.push(counter);
111
112 let block = hmac_sha3_256(&*prk, &block_input);
113 previous_block.clear();
114 previous_block.extend_from_slice(&block);
115
116 let remaining = N.checked_sub(written).ok_or(CryptoError::Hkdf {
117 hash: HkdfHash::Sha3_256,
118 kind: HkdfFailureKind::LengthOverflow,
119 })?;
120 let take = remaining.min(SHA3_256_DIGEST_LEN);
121 let end = written.checked_add(take).ok_or(CryptoError::Hkdf {
122 hash: HkdfHash::Sha3_256,
123 kind: HkdfFailureKind::LengthOverflow,
124 })?;
125 output[written..end].copy_from_slice(&previous_block[..take]);
126 written = end;
127 }
128
129 prk.zeroize();
130 Ok(output)
131}
132
133pub fn derive<const N: usize>(request: &DeriveRequest<'_>) -> Result<HkdfOutput<N>, CryptoError> {
160 if request.ikm.as_bytes().is_empty() {
161 return Err(CryptoError::Hkdf {
162 hash: request.suite.hash(),
163 kind: HkdfFailureKind::InvalidIkmLength,
164 });
165 }
166
167 if N == 0 {
168 return Err(CryptoError::Hkdf {
169 hash: request.suite.hash(),
170 kind: HkdfFailureKind::InvalidOutputLength,
171 });
172 }
173
174 let mut output = [0u8; N];
175
176 match request.suite {
177 HkdfSuite::Sha2_256 => {
178 let hkdf =
179 Hkdf::<Sha256>::new(request.salt.map(HkdfSalt::as_bytes), request.ikm.as_bytes());
180 hkdf.expand(request.info.as_bytes(), &mut output)
181 .map_err(|_| CryptoError::Hkdf {
182 hash: HkdfHash::Sha2_256,
183 kind: HkdfFailureKind::ExpandFailed,
184 })?;
185 }
186 HkdfSuite::Sha3_256 => {
187 #[cfg(feature = "sha3")]
188 {
189 output = derive_sha3_256::<N>(request)?;
190 }
191
192 #[cfg(not(feature = "sha3"))]
193 {
194 return Err(CryptoError::Unsupported);
195 }
196 }
197 }
198
199 Ok(HkdfOutput::from_array(output))
200}
201
202pub fn derive_domain_key_32(
209 ikm: &HkdfInputKeyMaterial,
210 salt: Option<&HkdfSalt>,
211 purpose: DomainKeyPurpose,
212 domain_tag: &DomainTag,
213) -> Result<HkdfOutput<32>, CryptoError> {
214 const PREFIX: &[u8] = b"reallyme/crypto/hkdf/v1/";
215 const SEPARATOR: &[u8] = b"/";
216
217 let capacity = PREFIX
218 .len()
219 .checked_add(purpose.as_bytes().len())
220 .and_then(|value| value.checked_add(SEPARATOR.len()))
221 .and_then(|value| value.checked_add(domain_tag.as_bytes().len()))
222 .ok_or(CryptoError::Hkdf {
223 hash: HkdfHash::Sha3_256,
224 kind: HkdfFailureKind::LengthOverflow,
225 })?;
226
227 let mut info = Vec::<u8>::with_capacity(capacity);
228 info.extend_from_slice(PREFIX);
229 info.extend_from_slice(purpose.as_bytes());
230 info.extend_from_slice(SEPARATOR);
231 info.extend_from_slice(domain_tag.as_bytes());
232
233 let hkdf_info = HkdfInfo::from_vec(info);
234 derive::<32>(&DeriveRequest {
235 suite: HkdfSuite::Sha3_256,
236 ikm,
237 salt,
238 info: &hkdf_info,
239 })
240}