use crate::cbor::{encode_canonical_cbor, CborValue};
pub const KEM_X25519: &str = "x25519";
pub const KEM_MLKEM768X25519: &str = "mlkem768x25519";
pub const AEAD_XCHACHA20_POLY1305: &str = "xchacha20-poly1305";
const CHUNK_MAX_BYTES: usize = 64;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct X25519Slot {
pub epk: Vec<u8>,
pub wrap: Vec<u8>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Mlkem768X25519Slot {
pub kem_ct: Vec<Vec<u8>>,
pub wrap: Vec<u8>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum SealedSlots {
X25519(Vec<X25519Slot>),
Mlkem768X25519(Vec<Mlkem768X25519Slot>),
}
impl SealedSlots {
#[must_use]
pub fn len(&self) -> usize {
match self {
SealedSlots::X25519(s) => s.len(),
SealedSlots::Mlkem768X25519(s) => s.len(),
}
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.len() == 0
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SealedEnvelope {
pub scheme: i64,
pub aead: String,
pub kem: String,
pub nonce: Vec<u8>,
pub slots: SealedSlots,
pub slots_mac: Vec<u8>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SealedPoeOutput {
pub envelope: SealedEnvelope,
pub ciphertext: Vec<u8>,
}
#[must_use]
pub fn chunk_kem_ct(value: &[u8]) -> Vec<Vec<u8>> {
assert!(
!value.is_empty(),
"chunk_kem_ct: refusing to chunk an empty byte string",
);
value.chunks(CHUNK_MAX_BYTES).map(<[u8]>::to_vec).collect()
}
#[must_use]
pub fn join_kem_ct(chunks: &[Vec<u8>]) -> Vec<u8> {
chunks.iter().flatten().copied().collect()
}
#[must_use]
pub fn slots_to_mac_cbor(slots: &SealedSlots) -> Vec<u8> {
let value = match slots {
SealedSlots::X25519(slots) => CborValue::Array(
slots
.iter()
.map(|s| {
CborValue::Map(vec![
(CborValue::text("epk"), CborValue::Bytes(s.epk.clone())),
(CborValue::text("wrap"), CborValue::Bytes(s.wrap.clone())),
])
})
.collect(),
),
SealedSlots::Mlkem768X25519(slots) => CborValue::Array(
slots
.iter()
.map(|s| {
let canonical = chunk_kem_ct(&join_kem_ct(&s.kem_ct));
let chunks = CborValue::Array(
canonical
.iter()
.map(|c| CborValue::Bytes(c.clone()))
.collect(),
);
CborValue::Map(vec![
(CborValue::text("kem_ct"), chunks),
(CborValue::text("wrap"), CborValue::Bytes(s.wrap.clone())),
])
})
.collect(),
),
};
encode_canonical_cbor(&value)
.expect("slot byte strings never collide as duplicate map keys, so encoding cannot fail")
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn chunk_then_join_is_the_identity_for_1120_bytes() {
let enc: Vec<u8> = (0..1120u32).map(|i| (i & 0xff) as u8).collect();
let chunks = chunk_kem_ct(&enc);
assert_eq!(chunks.len(), 18);
for c in &chunks[..17] {
assert_eq!(c.len(), 64);
}
assert_eq!(chunks[17].len(), 32);
assert_eq!(join_kem_ct(&chunks), enc);
}
#[test]
#[should_panic(expected = "refusing to chunk an empty byte string")]
fn chunk_rejects_empty() {
let _ = chunk_kem_ct(&[]);
}
#[test]
fn hybrid_mac_cbor_orders_wrap_before_kem_ct() {
let slots = SealedSlots::Mlkem768X25519(vec![Mlkem768X25519Slot {
kem_ct: vec![vec![0xaa; 4]],
wrap: vec![0xbb; 48],
}]);
let bytes = slots_to_mac_cbor(&slots);
assert_eq!(bytes[0], 0x81);
assert_eq!(bytes[1], 0xa2);
assert_eq!(&bytes[2..7], b"\x64wrap");
}
}