1pub use hkdf::Hkdf;
38
39use crate::constants::{
40 ECDH_SHARED_SECRET_SIZE, EC_PUBKEY_COMPRESSED_SIZE, EC_PUBKEY_UNCOMPRESSED_SIZE, MAX_HKDF_OUTPUT_SIZE, MIN_KEY_SIZE,
41};
42use crate::crypto::hash::{Digest, Sha3_256};
43use crate::crypto::secret::{SecretSlice, ToInsecure};
44use crate::zeroize::Zeroizing;
45use crate::ZeroizingArray;
46
47pub type Result<T> = ::core::result::Result<T, KdfError>;
48
49pub trait KdfFunction {
54 fn derive_key<const N: usize>(ikm: &[u8], info: &[u8], salt: Option<&[u8]>) -> Result<ZeroizingArray<N>>;
56
57 fn derive_dynamic_key(ikm: &[u8], info: &[u8], salt: Option<&[u8]>, key_size: usize) -> Result<Zeroizing<Vec<u8>>>;
74
75 fn derive_dual_keys<const N: usize>(
98 ikm: &[u8],
99 info: &[u8],
100 salt: Option<&[u8]>,
101 ) -> Result<(ZeroizingArray<N>, ZeroizingArray<N>)> {
102 let mut enc_info = Vec::with_capacity(info.len() + 11);
106 enc_info.extend_from_slice(info);
107 enc_info.extend_from_slice(b"-encryption");
108
109 let mut mac_info = Vec::with_capacity(info.len() + 4);
110 mac_info.extend_from_slice(info);
111 mac_info.extend_from_slice(b"-mac");
112
113 let k_enc = Self::derive_key::<N>(ikm, &enc_info, salt)?;
114 let k_mac = Self::derive_key::<N>(ikm, &mac_info, salt)?;
115
116 Ok((k_enc, k_mac))
117 }
118}
119pub struct HkdfSha3_256;
121
122crate::define_oid_wrapper!(
123 HkdfSha3_256Oid,
126 "2.16.840.1.101.3.4.2.8"
127);
128
129impl KdfFunction for HkdfSha3_256 {
130 fn derive_key<const N: usize>(ikm: &[u8], info: &[u8], salt: Option<&[u8]>) -> Result<ZeroizingArray<N>> {
131 let hk = Hkdf::<Sha3_256>::new(salt, ikm);
132 let mut output = Zeroizing::new([0u8; N]);
133 hk.expand(info, &mut output[..]).map_err(KdfError::DerivationFailed)?;
134 Ok(output)
135 }
136
137 fn derive_dynamic_key(ikm: &[u8], info: &[u8], salt: Option<&[u8]>, key_size: usize) -> Result<Zeroizing<Vec<u8>>> {
138 if !(MIN_KEY_SIZE..=MAX_HKDF_OUTPUT_SIZE).contains(&key_size) {
139 return Err(KdfError::DerivationFailed(hkdf::InvalidLength));
140 }
141
142 let hk = Hkdf::<Sha3_256>::new(salt, ikm);
143 let mut okm = vec![0u8; key_size];
144 hk.expand(info, &mut okm).map_err(KdfError::DerivationFailed)?;
145
146 Ok(Zeroizing::new(okm))
147 }
148
149 fn derive_dual_keys<const N: usize>(
152 ikm: &[u8],
153 info: &[u8],
154 salt: Option<&[u8]>,
155 ) -> Result<(ZeroizingArray<N>, ZeroizingArray<N>)> {
156 if N * 2 > MAX_HKDF_OUTPUT_SIZE {
158 return Err(KdfError::DerivationFailed(hkdf::InvalidLength));
159 }
160 let hk = Hkdf::<Sha3_256>::new(salt, ikm);
166 let mut combined = Zeroizing::new([0u8; MAX_HKDF_OUTPUT_SIZE]);
167 hk.expand(info, &mut combined[..N * 2]).map_err(KdfError::DerivationFailed)?;
168
169 let mut k_enc = Zeroizing::new([0u8; N]);
170 let mut k_mac = Zeroizing::new([0u8; N]);
171 k_enc[..].copy_from_slice(&combined[..N]);
172 k_mac[..].copy_from_slice(&combined[N..N * 2]);
173 Ok((k_enc, k_mac))
174 }
175}
176
177pub struct X963Sha3_256;
179
180impl KdfFunction for X963Sha3_256 {
181 fn derive_key<const N: usize>(ikm: &[u8], info: &[u8], _salt: Option<&[u8]>) -> Result<ZeroizingArray<N>> {
182 let mut out = Zeroizing::new([0u8; N]);
184 let mut offset = 0usize;
185 let mut counter: u32 = 1;
186 while offset < N {
187 let mut hasher = Sha3_256::new();
188 hasher.update(ikm); hasher.update(counter.to_be_bytes());
190 hasher.update(info); let block = hasher.finalize();
192 let take = core::cmp::min(block.len(), N - offset);
193 out[offset..offset + take].copy_from_slice(&block[..take]);
194 offset += take;
195 counter = counter.wrapping_add(1);
196 }
197 Ok(out)
198 }
199
200 fn derive_dynamic_key(
201 ikm: &[u8],
202 info: &[u8],
203 _salt: Option<&[u8]>,
204 key_size: usize,
205 ) -> Result<Zeroizing<Vec<u8>>> {
206 if key_size < MIN_KEY_SIZE {
207 return Err(KdfError::DerivationFailed(hkdf::InvalidLength));
208 }
209
210 let mut out = vec![0u8; key_size];
212 let mut offset = 0usize;
213 let mut counter: u32 = 1;
214 while offset < key_size {
215 let mut hasher = Sha3_256::new();
216 hasher.update(ikm);
217 hasher.update(counter.to_be_bytes());
218 hasher.update(info);
219 let block = hasher.finalize();
220 let take = core::cmp::min(block.len(), key_size - offset);
221 out[offset..offset + take].copy_from_slice(&block[..take]);
222 offset += take;
223 counter = counter.wrapping_add(1);
224 }
225 Ok(Zeroizing::new(out))
226 }
227
228 fn derive_dual_keys<const N: usize>(
229 ikm: &[u8],
230 info: &[u8],
231 _salt: Option<&[u8]>,
232 ) -> Result<(ZeroizingArray<N>, ZeroizingArray<N>)> {
233 let mut enc_info = [0u8; 256]; let mut mac_info = [0u8; 256];
237
238 let enc_info_len = core::cmp::min(enc_info.len(), info.len() + 11);
239 let mac_info_len = core::cmp::min(mac_info.len(), info.len() + 4);
240
241 enc_info[..info.len()].copy_from_slice(info);
242 enc_info[info.len()..info.len() + 11].copy_from_slice(b"-encryption");
243
244 mac_info[..info.len()].copy_from_slice(info);
245 mac_info[info.len()..info.len() + 4].copy_from_slice(b"-mac");
246
247 let k_enc = Self::derive_key::<N>(ikm, &enc_info[..enc_info_len], None)?;
248 let k_mac = Self::derive_key::<N>(ikm, &mac_info[..mac_info_len], None)?;
249 Ok((k_enc, k_mac))
250 }
251}
252
253#[cfg_attr(feature = "derive", derive(crate::Errorizable))]
255#[derive(Debug, Clone)]
256pub enum KdfError {
257 #[cfg_attr(feature = "derive", error("Key derivation failed: {0}"))]
259 DerivationFailed(hkdf::InvalidLength),
260
261 #[cfg_attr(
263 feature = "derive",
264 error("Invalid ephemeral public key length: expected 33 or 65 bytes, got {0}")
265 )]
266 InvalidPublicKeyLength(usize),
267
268 #[cfg_attr(
270 feature = "derive",
271 error("Invalid shared secret length: expected 32 bytes, got {0}")
272 )]
273 InvalidSharedSecretLength(usize),
274
275 #[cfg_attr(
277 feature = "derive",
278 error("Invalid salt length: must be at least 16 bytes, got {0}")
279 )]
280 InvalidSaltLength(usize),
281}
282
283crate::impl_error_display!(KdfError {
284 DerivationFailed(e) => "Key derivation failed: {e}",
285 InvalidPublicKeyLength(len) => "Invalid ephemeral public key length: expected 33 or 65 bytes, got {len}",
286 InvalidSharedSecretLength(len) => "Invalid shared secret length: expected 32 bytes, got {len}",
287 InvalidSaltLength(len) => "Invalid salt length: must be at least 16 bytes, got {len}",
288});
289
290#[inline]
296fn assert_valid_shared_secret(shared_secret: &[u8]) -> Result<()> {
297 if shared_secret.len() != ECDH_SHARED_SECRET_SIZE {
298 return Err(KdfError::InvalidSharedSecretLength(shared_secret.len()));
299 }
300 Ok(())
301}
302
303#[inline]
305fn assert_valid_salt(salt: Option<&[u8]>) -> Result<()> {
306 if let Some(salt_bytes) = salt {
307 if !salt_bytes.is_empty() && salt_bytes.len() < MIN_KEY_SIZE {
308 return Err(KdfError::InvalidSaltLength(salt_bytes.len()));
309 }
310 }
311 Ok(())
312}
313
314#[inline]
316fn assert_valid_ephemeral_pubkey(ephemeral_pubkey: &[u8]) -> Result<()> {
317 if ephemeral_pubkey.len() != EC_PUBKEY_COMPRESSED_SIZE && ephemeral_pubkey.len() != EC_PUBKEY_UNCOMPRESSED_SIZE {
318 return Err(KdfError::InvalidPublicKeyLength(ephemeral_pubkey.len()));
319 }
320 Ok(())
321}
322
323#[inline]
325fn assert_valid_kdf_inputs(ephemeral_pubkey: &[u8], shared_secret: &[u8], salt: Option<&[u8]>) -> Result<()> {
326 assert_valid_ephemeral_pubkey(ephemeral_pubkey)?;
327 assert_valid_shared_secret(shared_secret)?;
328 assert_valid_salt(salt)?;
329 Ok(())
330}
331
332pub fn ecies_kdf<P: KdfFunction>(
351 ephemeral_pubkey: impl AsRef<[u8]>,
352 shared_secret: SecretSlice<u8>,
353 info: impl AsRef<[u8]>,
354 salt: Option<&[u8]>,
355) -> Result<ZeroizingArray<32>> {
356 let ephemeral_pubkey = ephemeral_pubkey.as_ref();
357 let shared_secret_bytes = shared_secret.to_insecure()?;
358 let shared_secret = shared_secret_bytes.as_ref();
359
360 assert_valid_kdf_inputs(ephemeral_pubkey, shared_secret, salt)?;
361
362 let mut shared_info = Vec::with_capacity(info.as_ref().len() + 5 + ephemeral_pubkey.len());
364 shared_info.extend_from_slice(info.as_ref());
365 shared_info.extend_from_slice(b"|epk|");
366 shared_info.extend_from_slice(ephemeral_pubkey);
367 P::derive_key::<32>(shared_secret, &shared_info, salt)
368}
369
370pub fn hkdf<P: KdfFunction, const N: usize>(
383 ikm: impl AsRef<[u8]>,
384 info: impl AsRef<[u8]>,
385 salt: Option<&[u8]>,
386) -> Result<ZeroizingArray<N>> {
387 let (ikm, info) = (ikm.as_ref(), info.as_ref());
388 P::derive_key::<N>(ikm, info, salt)
389}
390
391pub fn ecies_kdf_with_size<P: KdfFunction, const N: usize>(
411 ephemeral_pubkey: impl AsRef<[u8]>,
412 shared_secret: SecretSlice<u8>,
413 info: impl AsRef<[u8]>,
414 salt: Option<&[u8]>,
415) -> Result<(ZeroizingArray<N>, ZeroizingArray<N>)> {
416 let insecure_shared_secret = shared_secret.to_insecure()?;
417 let (ephemeral_pubkey, shared_secret, info) =
418 (ephemeral_pubkey.as_ref(), insecure_shared_secret.as_ref(), info.as_ref());
419
420 assert_valid_kdf_inputs(ephemeral_pubkey, shared_secret, salt)?;
421
422 let mut shared_info = Vec::with_capacity(info.len() + 5 + ephemeral_pubkey.len());
424 shared_info.extend_from_slice(info);
425 shared_info.extend_from_slice(b"|epk|");
426 shared_info.extend_from_slice(ephemeral_pubkey);
427 P::derive_dual_keys::<N>(shared_secret, &shared_info, salt)
428}
429
430pub fn ecies_kdf_with_shared_info<P: KdfFunction>(
433 shared_secret: SecretSlice<u8>,
434 shared_info: impl AsRef<[u8]>,
435 salt: Option<&[u8]>,
436) -> Result<ZeroizingArray<32>> {
437 let insecure_shared_secret = shared_secret.to_insecure()?;
438 let (shared_secret, shared_info) = (insecure_shared_secret.as_ref(), shared_info.as_ref());
439 assert_valid_shared_secret(shared_secret)?;
440 assert_valid_salt(salt)?;
441
442 P::derive_key::<32>(shared_secret, shared_info, salt)
443}
444
445pub fn ecies_kdf_with_shared_info_and_size<P: KdfFunction, const N: usize>(
448 shared_secret: SecretSlice<u8>,
449 shared_info: impl AsRef<[u8]>,
450 salt: Option<&[u8]>,
451) -> Result<(ZeroizingArray<N>, ZeroizingArray<N>)> {
452 let insecure_shared_secret = shared_secret.to_insecure()?;
453 let (shared_secret, shared_info) = (insecure_shared_secret.as_ref(), shared_info.as_ref());
454 assert_valid_shared_secret(shared_secret)?;
455 assert_valid_salt(salt)?;
456
457 P::derive_dual_keys::<N>(shared_secret, shared_info, salt)
458}
459
460#[cfg(test)]
461mod tests {
462 use super::*;
463 use crate::crypto::secret::Secret;
464
465 #[track_caller]
467 fn assert_key_length<const N: usize>(key: &ZeroizingArray<N>, expected_len: usize) {
468 assert_eq!(key.len(), expected_len, "Key length mismatch");
469 }
470
471 #[track_caller]
472 fn assert_keys_equal<const N: usize>(key1: &ZeroizingArray<N>, key2: &ZeroizingArray<N>) {
473 assert_eq!(key1[..], key2[..], "Keys should be equal");
474 }
475
476 #[track_caller]
477 fn assert_keys_different<const N: usize>(key1: &ZeroizingArray<N>, key2: &ZeroizingArray<N>) {
478 assert_ne!(key1[..], key2[..], "Keys should be different");
479 }
480
481 macro_rules! assert_kdf_error {
483 ($result:expr, $variant:ident($value:expr)) => {
484 assert!(matches!($result, Err(KdfError::$variant(v)) if v == $value),
485 "Expected KdfError::{}({}), got {:?}", stringify!($variant), $value, $result);
486 };
487 }
488
489 macro_rules! assert_key_pair_lengths {
491 ($enc:expr, $mac:expr, $size:expr) => {
492 assert_eq!($enc.len(), $size, "Encryption key length mismatch");
493 assert_eq!($mac.len(), $size, "MAC key length mismatch");
494 };
495 }
496
497 macro_rules! assert_keys_different {
498 ($key1:expr, $key2:expr) => {
499 assert_ne!($key1[..], $key2[..], "Keys should be different");
500 };
501 }
502
503 macro_rules! assert_key_length {
504 ($key:expr, $size:expr) => {
505 assert_eq!($key.len(), $size, "Key length mismatch");
506 };
507 }
508
509 fn shared_secret_32() -> SecretSlice<u8> {
510 Secret::from(b"shared_secret_32_bytes__________".to_vec())
511 }
512
513 const EPHEMERAL_PUBKEY_33: &[u8] = b"ephemeral_public_key_33_bytes____";
515 const EPHEMERAL_PUBKEY_33_ALT: &[u8] = b"different_ephemeral_key_33_bytes_";
516 const INFO_V1: &[u8] = b"tightbeam-ecies-v1";
517 const INFO_V2: &[u8] = b"protocol-v2";
518 const SALT: &[u8] = b"random_salt_value";
519
520 #[test]
522 fn test_ecies_kdf_basic_functionality() -> crate::error::Result<()> {
523 let basic_key = ecies_kdf::<HkdfSha3_256>(EPHEMERAL_PUBKEY_33, shared_secret_32(), INFO_V1, None).unwrap();
525 let same_key = ecies_kdf::<HkdfSha3_256>(EPHEMERAL_PUBKEY_33, shared_secret_32(), INFO_V1, None).unwrap();
526 let different_pubkey =
528 ecies_kdf::<HkdfSha3_256>(EPHEMERAL_PUBKEY_33_ALT, shared_secret_32(), INFO_V1, None).unwrap();
529 let different_info = ecies_kdf::<HkdfSha3_256>(EPHEMERAL_PUBKEY_33, shared_secret_32(), INFO_V2, None).unwrap();
530 let with_salt =
531 ecies_kdf::<HkdfSha3_256>(EPHEMERAL_PUBKEY_33, shared_secret_32(), INFO_V1, Some(SALT)).unwrap();
532
533 let mut uncompressed_pubkey = [0u8; 65];
535 uncompressed_pubkey[0] = 0x04; for (i, byte) in uncompressed_pubkey.iter_mut().enumerate().skip(1) {
537 *byte = (i % 256) as u8;
538 }
539
540 let uncompressed_result = ecies_kdf::<HkdfSha3_256>(uncompressed_pubkey, shared_secret_32(), INFO_V1, None);
541
542 assert_key_length(&basic_key, 32);
544 assert_keys_equal(&basic_key, &same_key);
546 assert_keys_different(&basic_key, &different_pubkey); assert_keys_different(&basic_key, &different_info); assert_keys_different(&basic_key, &with_salt); assert!(uncompressed_result.is_ok());
552 assert_key_length(&uncompressed_result.unwrap(), 32);
553
554 Ok(())
555 }
556
557 #[test]
559 fn test_ecies_kdf_size_variations() -> crate::error::Result<()> {
560 let keys_16 =
562 ecies_kdf_with_size::<HkdfSha3_256, 16>(EPHEMERAL_PUBKEY_33, shared_secret_32(), INFO_V1, None).unwrap();
563 let keys_32 =
564 ecies_kdf_with_size::<HkdfSha3_256, 32>(EPHEMERAL_PUBKEY_33, shared_secret_32(), INFO_V1, None).unwrap();
565 let keys_64 =
566 ecies_kdf_with_size::<HkdfSha3_256, 64>(EPHEMERAL_PUBKEY_33, shared_secret_32(), INFO_V1, None).unwrap();
567
568 let (k_enc_16, k_mac_16) = keys_16;
569 let (k_enc_32, k_mac_32) = keys_32;
570 let (k_enc_64, k_mac_64) = keys_64;
571
572 assert_key_pair_lengths!(k_enc_16, k_mac_16, 16);
574 assert_key_pair_lengths!(k_enc_32, k_mac_32, 32);
575 assert_key_pair_lengths!(k_enc_64, k_mac_64, 64);
576 assert_keys_different!(k_enc_32, k_mac_32);
578
579 Ok(())
580 }
581
582 #[test]
584 fn test_ecies_kdf_input_validation() -> crate::error::Result<()> {
585 let short_pubkey_result = ecies_kdf::<HkdfSha3_256>(b"short", shared_secret_32(), INFO_V1, None);
587 let wrong_size_pubkey_result =
588 ecies_kdf::<HkdfSha3_256>(b"wrong_size_ephemeral_key_34_bytes_", shared_secret_32(), INFO_V1, None);
589 let short_secret_result =
590 ecies_kdf::<HkdfSha3_256>(EPHEMERAL_PUBKEY_33, Secret::from(b"short".to_vec()), INFO_V1, None);
591 let long_secret_result = ecies_kdf::<HkdfSha3_256>(
592 EPHEMERAL_PUBKEY_33,
593 Secret::from(b"shared_secret_that_is_too_long____".to_vec()),
594 INFO_V1,
595 None,
596 );
597
598 assert_kdf_error!(short_pubkey_result, InvalidPublicKeyLength(5));
600 assert_kdf_error!(wrong_size_pubkey_result, InvalidPublicKeyLength(34));
601 assert_kdf_error!(short_secret_result, InvalidSharedSecretLength(5));
603 assert_kdf_error!(long_secret_result, InvalidSharedSecretLength(34));
604
605 Ok(())
606 }
607
608 #[test]
610 fn test_hkdf_sha3_256_basic() -> crate::error::Result<()> {
611 let ikm = b"input_key_material";
612 let info = b"test_info";
613
614 let key_16 = hkdf::<HkdfSha3_256, 16>(ikm, info, None).unwrap();
616 let key_32 = hkdf::<HkdfSha3_256, 32>(ikm, info, None).unwrap();
617 let key_64 = hkdf::<HkdfSha3_256, 64>(ikm, info, None).unwrap();
618
619 let key_32_again = hkdf::<HkdfSha3_256, 32>(ikm, info, None).unwrap();
621 let key_different = hkdf::<HkdfSha3_256, 32>(b"different_ikm", info, None).unwrap();
623
624 assert_key_length!(key_16, 16);
626 assert_key_length!(key_32, 32);
627 assert_key_length!(key_64, 64);
628 assert_eq!(key_32[..], key_32_again[..]);
630 assert_keys_different!(key_32, key_different);
632
633 Ok(())
634 }
635
636 #[test]
638 fn test_ecies_kdf_bounds_checking() -> crate::error::Result<()> {
639 let max_size_result =
641 ecies_kdf_with_size::<HkdfSha3_256, 64>(EPHEMERAL_PUBKEY_33, shared_secret_32(), INFO_V1, None);
642 let oversized_result =
644 ecies_kdf_with_size::<HkdfSha3_256, 65>(EPHEMERAL_PUBKEY_33, shared_secret_32(), INFO_V1, None);
645
646 assert!(max_size_result.is_ok());
648 let (k_enc, k_mac) = max_size_result.unwrap();
649 assert_key_pair_lengths!(k_enc, k_mac, 64);
650
651 assert!(oversized_result.is_err());
653 assert!(matches!(oversized_result, Err(KdfError::DerivationFailed(_))));
654
655 Ok(())
656 }
657
658 #[test]
660 fn test_x963_ecies_kdf_basic() -> crate::error::Result<()> {
661 let key1 = ecies_kdf::<X963Sha3_256>(EPHEMERAL_PUBKEY_33, shared_secret_32(), INFO_V1, None).unwrap();
662 let key1_again = ecies_kdf::<X963Sha3_256>(EPHEMERAL_PUBKEY_33, shared_secret_32(), INFO_V1, None).unwrap();
663 let key_diff_info = ecies_kdf::<X963Sha3_256>(EPHEMERAL_PUBKEY_33, shared_secret_32(), INFO_V2, None).unwrap();
664 let key_with_salt =
666 ecies_kdf::<X963Sha3_256>(EPHEMERAL_PUBKEY_33, shared_secret_32(), INFO_V1, Some(SALT)).unwrap();
667
668 assert_key_length(&key1, 32);
669 assert_keys_equal(&key1, &key1_again);
670 assert_keys_different(&key1, &key_diff_info);
671 assert_keys_equal(&key1, &key_with_salt);
673 Ok(())
674 }
675
676 #[test]
677 fn test_x963_ecies_kdf_size_variations() -> crate::error::Result<()> {
678 let keys_16 =
679 ecies_kdf_with_size::<X963Sha3_256, 16>(EPHEMERAL_PUBKEY_33, shared_secret_32(), INFO_V1, None).unwrap();
680 let keys_32 =
681 ecies_kdf_with_size::<X963Sha3_256, 32>(EPHEMERAL_PUBKEY_33, shared_secret_32(), INFO_V1, None).unwrap();
682
683 let (k_enc_16, k_mac_16) = keys_16;
684 let (k_enc_32, k_mac_32) = keys_32;
685 assert_key_pair_lengths!(k_enc_16, k_mac_16, 16);
686 assert_key_pair_lengths!(k_enc_32, k_mac_32, 32);
687 assert_keys_different!(k_enc_32, k_mac_32);
688 Ok(())
689 }
690}