use hmac::{Hmac, KeyInit, Mac};
use sha2::Sha256;
use zeroize::Zeroize;
use crate::kdf::hkdf_sha256;
use super::aead::{chacha20_poly1305_encrypt, xchacha20_poly1305_encrypt};
use super::errors::{EciesSealedPoeError, EciesSealedPoeErrorCode};
use super::kem::{
mlkem768x25519_encapsulate, x25519_ecdh, x25519_public_key, MLKEM768X25519_ENC_LENGTH,
MLKEM768X25519_ESEED_LENGTH, MLKEM768X25519_PUBLIC_KEY_LENGTH,
};
use super::slots::{
chunk_kem_ct, slots_to_mac_cbor, Mlkem768X25519Slot, SealedEnvelope, SealedPoeOutput,
SealedSlots, X25519Slot, AEAD_XCHACHA20_POLY1305, KEM_MLKEM768X25519, KEM_X25519,
};
pub const CARDANO_POE_HKDF_INFO_KEK: &[u8] = b"cardano-poe-kek-v1";
pub const CARDANO_POE_HKDF_INFO_KEK_MLKEM768X25519: &[u8] = b"cardano-poe-kek-mlkem768x25519-v1";
pub const CARDANO_POE_HKDF_INFO_SLOTS_MAC: &[u8] = b"cardano-poe-slots-mac-v1";
const ZERO_NONCE_12: [u8; 12] = [0u8; 12];
const X25519_KEY_LENGTH: usize = 32;
const CEK_LENGTH: usize = 32;
const NONCE_LENGTH: usize = 24;
const WRAP_LENGTH: usize = 48;
const SLOTS_MAC_LENGTH: usize = 32;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SealedKem {
X25519,
Mlkem768X25519,
}
impl SealedKem {
#[must_use]
pub const fn as_str(self) -> &'static str {
match self {
SealedKem::X25519 => KEM_X25519,
SealedKem::Mlkem768X25519 => KEM_MLKEM768X25519,
}
}
}
pub type RandomSource<'r> = &'r mut dyn FnMut(&mut [u8]);
#[derive(Default)]
pub struct WrapArgs<'a> {
pub plaintext: &'a [u8],
pub recipient_public_keys: &'a [Vec<u8>],
pub kem: Option<SealedKem>,
pub cek: Option<&'a [u8]>,
pub nonce: Option<&'a [u8]>,
pub ephemeral_secrets: Option<&'a [Vec<u8>]>,
pub eseeds: Option<&'a [Vec<u8>]>,
pub skip_shuffle: bool,
}
#[must_use]
pub fn uniform_index_ceiling(m: u32) -> u64 {
assert!(m != 0, "uniform_index_ceiling: modulus must be positive");
let two_pow_32: u64 = 1 << 32;
two_pow_32 - (two_pow_32 % u64::from(m))
}
fn uniform_index_below(fill: &mut dyn FnMut(&mut [u8]), m: u32) -> u32 {
let limit = uniform_index_ceiling(m);
loop {
let mut buf = [0u8; 4];
fill(&mut buf);
let x = u64::from(u32::from_le_bytes(buf));
if x < limit {
return (x % u64::from(m)) as u32;
}
}
}
fn csprng_shuffle<T>(arr: &mut [T], fill: &mut dyn FnMut(&mut [u8])) {
if arr.len() < 2 {
return;
}
for i in (1..arr.len()).rev() {
let j = uniform_index_below(fill, (i + 1) as u32) as usize;
arr.swap(i, j);
}
}
fn wrap_slot_x25519(
pub_r: &[u8],
priv_eph: Option<&[u8]>,
cek: &[u8],
slot_idx: usize,
fill: &mut dyn FnMut(&mut [u8]),
) -> Result<X25519Slot, EciesSealedPoeError> {
let mut owned_eph = [0u8; X25519_KEY_LENGTH];
let priv_eph: &[u8] = match priv_eph {
Some(eph) => {
if eph.len() != X25519_KEY_LENGTH {
return Err(EciesSealedPoeError::new(
EciesSealedPoeErrorCode::InvalidEphemeralSecretLength,
format!(
"ephemeral_secrets[{slot_idx}] MUST be exactly {X25519_KEY_LENGTH} bytes, got {}",
eph.len()
),
));
}
eph
}
None => {
fill(&mut owned_eph);
&owned_eph
}
};
let epk =
x25519_public_key(priv_eph).expect("ephemeral scalar is exactly 32 bytes, validated above");
let mut shared = x25519_ecdh(priv_eph, pub_r).map_err(|e| {
EciesSealedPoeError::new(
EciesSealedPoeErrorCode::KemEpkLengthMismatch,
format!("recipient_public_keys[{slot_idx}] X25519 ECDH failed: {e}"),
)
})?;
let mut salt = Vec::with_capacity(epk.len() + pub_r.len());
salt.extend_from_slice(&epk);
salt.extend_from_slice(pub_r);
let mut kek = hkdf_sha256(&shared, &salt, CARDANO_POE_HKDF_INFO_KEK, 32)
.expect("32-byte HKDF output is within the RFC 5869 maximum");
shared.zeroize();
let wrap = chacha20_poly1305_encrypt(&kek, &ZERO_NONCE_12, CARDANO_POE_HKDF_INFO_KEK, cek);
kek.zeroize();
owned_eph.zeroize();
debug_assert_eq!(wrap.len(), WRAP_LENGTH);
Ok(X25519Slot {
epk: epk.to_vec(),
wrap,
})
}
fn wrap_slot_mlkem768x25519(
pub_r: &[u8],
eseed: &[u8],
cek: &[u8],
slot_idx: usize,
) -> Result<Mlkem768X25519Slot, EciesSealedPoeError> {
let encaps = mlkem768x25519_encapsulate(pub_r, eseed).map_err(|e| {
EciesSealedPoeError::new(
EciesSealedPoeErrorCode::KemEpkLengthMismatch,
format!("recipient_public_keys[{slot_idx}] X-Wing encapsulation failed: {e}"),
)
})?;
debug_assert_eq!(encaps.enc.len(), MLKEM768X25519_ENC_LENGTH);
let mut kek = hkdf_sha256(
&encaps.ss,
&[],
CARDANO_POE_HKDF_INFO_KEK_MLKEM768X25519,
32,
)
.expect("32-byte HKDF output is within the RFC 5869 maximum");
let wrap = chacha20_poly1305_encrypt(
&kek,
&ZERO_NONCE_12,
CARDANO_POE_HKDF_INFO_KEK_MLKEM768X25519,
cek,
);
kek.zeroize();
debug_assert_eq!(wrap.len(), WRAP_LENGTH);
Ok(Mlkem768X25519Slot {
kem_ct: chunk_kem_ct(&encaps.enc),
wrap,
})
}
fn compute_slots_mac(cek: &[u8], slots: &SealedSlots) -> [u8; SLOTS_MAC_LENGTH] {
let mut hmac_key = hkdf_sha256(cek, &[], CARDANO_POE_HKDF_INFO_SLOTS_MAC, 32)
.expect("32-byte HKDF output is within the RFC 5869 maximum");
let slots_cbor = slots_to_mac_cbor(slots);
let mut mac =
<Hmac<Sha256>>::new_from_slice(&hmac_key).expect("HMAC accepts a key of any length");
mac.update(&slots_cbor);
let out: [u8; SLOTS_MAC_LENGTH] = mac.finalize().into_bytes().into();
hmac_key.zeroize();
out
}
pub fn ecies_sealed_poe_wrap_secure(
args: WrapArgs<'_>,
) -> Result<SealedPoeOutput, EciesSealedPoeError> {
let mut rng_error: Option<EciesSealedPoeError> = None;
let mut fill = |buf: &mut [u8]| {
if rng_error.is_some() {
return;
}
if let Err(e) = getrandom::getrandom(buf) {
rng_error = Some(EciesSealedPoeError::new(
EciesSealedPoeErrorCode::RngUnavailable,
format!("operating-system CSPRNG is unavailable: {e}"),
));
}
};
let result = ecies_sealed_poe_wrap_with_rng(args, &mut fill);
if let Some(e) = rng_error {
return Err(e);
}
result
}
pub fn ecies_sealed_poe_wrap_with_rng(
args: WrapArgs<'_>,
rng: RandomSource<'_>,
) -> Result<SealedPoeOutput, EciesSealedPoeError> {
let kem = args.kem.unwrap_or(SealedKem::X25519);
let n = args.recipient_public_keys.len();
if n < 1 {
return Err(EciesSealedPoeError::new(
EciesSealedPoeErrorCode::EncSlotsEmpty,
format!("recipient_public_keys.len()={n} must be >= 1"),
));
}
let expected_pub_len = match kem {
SealedKem::X25519 => X25519_KEY_LENGTH,
SealedKem::Mlkem768X25519 => MLKEM768X25519_PUBLIC_KEY_LENGTH,
};
for (i, pub_key) in args.recipient_public_keys.iter().enumerate() {
if pub_key.len() != expected_pub_len {
return Err(EciesSealedPoeError::new(
EciesSealedPoeErrorCode::KemEpkLengthMismatch,
format!(
"recipient_public_keys[{i}] MUST be exactly {expected_pub_len} bytes for kem='{}'",
kem.as_str()
),
));
}
}
match kem {
SealedKem::X25519 => {
if args.eseeds.is_some() {
return Err(EciesSealedPoeError::new(
EciesSealedPoeErrorCode::EphemeralSecretsCountMismatch,
"eseeds is an X-Wing override and MUST NOT be supplied for kem='x25519'",
));
}
if let Some(eph) = args.ephemeral_secrets {
if eph.len() != n {
return Err(EciesSealedPoeError::new(
EciesSealedPoeErrorCode::EphemeralSecretsCountMismatch,
format!(
"ephemeral_secrets.len()={} must match recipient_public_keys.len()={n}",
eph.len()
),
));
}
}
}
SealedKem::Mlkem768X25519 => {
if args.ephemeral_secrets.is_some() {
return Err(EciesSealedPoeError::new(
EciesSealedPoeErrorCode::EphemeralSecretsCountMismatch,
"ephemeral_secrets is an X25519 override and MUST NOT be supplied for kem='mlkem768x25519'",
));
}
if let Some(eseeds) = args.eseeds {
if eseeds.len() != n {
return Err(EciesSealedPoeError::new(
EciesSealedPoeErrorCode::EphemeralSecretsCountMismatch,
format!(
"eseeds.len()={} must match recipient_public_keys.len()={n}",
eseeds.len()
),
));
}
for (i, eseed) in eseeds.iter().enumerate() {
if eseed.len() != MLKEM768X25519_ESEED_LENGTH {
return Err(EciesSealedPoeError::new(
EciesSealedPoeErrorCode::InvalidEphemeralSecretLength,
format!(
"eseeds[{i}] MUST be exactly {MLKEM768X25519_ESEED_LENGTH} bytes, got {}",
eseed.len()
),
));
}
}
}
}
}
let mut owned_cek = [0u8; CEK_LENGTH];
let cek: &[u8] = match args.cek {
Some(c) => {
if c.len() != CEK_LENGTH {
return Err(EciesSealedPoeError::new(
EciesSealedPoeErrorCode::InvalidCekLength,
format!("cek MUST be exactly {CEK_LENGTH} bytes, got {}", c.len()),
));
}
c
}
None => {
rng(&mut owned_cek);
&owned_cek
}
};
let mut owned_nonce = [0u8; NONCE_LENGTH];
let nonce: &[u8] = match args.nonce {
Some(nc) => {
if nc.len() != NONCE_LENGTH {
return Err(EciesSealedPoeError::new(
EciesSealedPoeErrorCode::NonceLengthMismatch,
format!(
"nonce MUST be exactly {NONCE_LENGTH} bytes, got {}",
nc.len()
),
));
}
nc
}
None => {
rng(&mut owned_nonce);
&owned_nonce
}
};
let slots = match kem {
SealedKem::X25519 => {
let mut slots = Vec::with_capacity(n);
for (i, pub_r) in args.recipient_public_keys.iter().enumerate() {
let priv_eph = args.ephemeral_secrets.map(|e| e[i].as_slice());
slots.push(wrap_slot_x25519(pub_r, priv_eph, cek, i, rng)?);
}
if !args.skip_shuffle {
csprng_shuffle(&mut slots, rng);
}
SealedSlots::X25519(slots)
}
SealedKem::Mlkem768X25519 => {
let mut slots = Vec::with_capacity(n);
for (i, pub_r) in args.recipient_public_keys.iter().enumerate() {
let mut fresh = [0u8; MLKEM768X25519_ESEED_LENGTH];
let eseed: &[u8] = match args.eseeds {
Some(e) => e[i].as_slice(),
None => {
rng(&mut fresh);
&fresh
}
};
let slot = wrap_slot_mlkem768x25519(pub_r, eseed, cek, i);
fresh.zeroize();
slots.push(slot?);
}
if !args.skip_shuffle {
csprng_shuffle(&mut slots, rng);
}
SealedSlots::Mlkem768X25519(slots)
}
};
let slots_mac = compute_slots_mac(cek, &slots);
let mut ad_content = Vec::with_capacity(NONCE_LENGTH + SLOTS_MAC_LENGTH);
ad_content.extend_from_slice(nonce);
ad_content.extend_from_slice(&slots_mac);
let ciphertext = xchacha20_poly1305_encrypt(cek, nonce, &ad_content, args.plaintext);
owned_cek.zeroize();
Ok(SealedPoeOutput {
envelope: SealedEnvelope {
scheme: 1,
aead: AEAD_XCHACHA20_POLY1305.to_string(),
kem: kem.as_str().to_string(),
nonce: nonce.to_vec(),
slots,
slots_mac: slots_mac.to_vec(),
},
ciphertext,
})
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn info_label_byte_lengths_match_the_protocol() {
assert_eq!(CARDANO_POE_HKDF_INFO_KEK.len(), 18);
assert_eq!(CARDANO_POE_HKDF_INFO_KEK_MLKEM768X25519.len(), 33);
assert_eq!(CARDANO_POE_HKDF_INFO_SLOTS_MAC.len(), 24);
}
#[test]
fn uniform_index_ceiling_is_a_multiple_of_m_and_two_pow_32_for_powers_of_two() {
let two_pow_32: u64 = 1 << 32;
for m in [2u32, 3, 4, 5, 6, 7, 8, 17, 64, 100, 256, 257, 1000] {
let limit = uniform_index_ceiling(m);
assert_eq!(limit % u64::from(m), 0);
let is_power_of_two = m.is_power_of_two();
assert_eq!(limit == two_pow_32, is_power_of_two);
}
}
#[test]
fn uniform_index_below_stays_in_range() {
let mut ctr: u32 = 0;
let mut fill = |buf: &mut [u8]| {
ctr = ctr.wrapping_add(0x9e37_79b9);
for (i, b) in buf.iter_mut().enumerate() {
*b = (ctr >> (8 * (i % 4))) as u8;
}
};
for m in [1u32, 2, 3, 5, 7, 17, 100, 257] {
for _ in 0..200 {
let v = uniform_index_below(&mut fill, m);
assert!(v < m);
}
}
}
#[test]
fn empty_recipients_is_rejected() {
let mut rng = |_: &mut [u8]| panic!("deterministic path must not draw randomness");
let err = ecies_sealed_poe_wrap_with_rng(
WrapArgs {
plaintext: b"",
recipient_public_keys: &[],
..Default::default()
},
&mut rng,
)
.unwrap_err();
assert_eq!(err.code(), "ENC_SLOTS_EMPTY");
}
#[test]
fn secure_wrap_draws_fresh_random_material_each_call() {
let recipient = x25519_public_key(&[3u8; X25519_KEY_LENGTH]).unwrap();
let recipients = vec![recipient.to_vec()];
let mk = || {
ecies_sealed_poe_wrap_secure(WrapArgs {
plaintext: b"hello sealed poe",
recipient_public_keys: &recipients,
..Default::default()
})
.unwrap()
};
let a = mk();
let b = mk();
assert_ne!(a.ciphertext, b.ciphertext, "fresh CEK/nonce per wrap");
assert_ne!(a.envelope.nonce, b.envelope.nonce, "fresh nonce per wrap");
assert_ne!(
a.envelope.slots_mac, b.envelope.slots_mac,
"slots_mac is keyed by the fresh CEK"
);
assert_eq!(a.envelope.nonce.len(), NONCE_LENGTH);
assert_eq!(a.envelope.slots_mac.len(), SLOTS_MAC_LENGTH);
}
}