use alloc::vec::Vec;
use crate::key::{Decapsulator, Encapsulator, Error, Secret};
use crate::rng::CryptoRngCore;
use super::{
MlKem512Ciphertext, MlKem512DecapsKey, MlKem512EncapsKey, MlKem768Ciphertext,
MlKem768DecapsKey, MlKem768EncapsKey, MlKem1024Ciphertext, MlKem1024DecapsKey,
MlKem1024EncapsKey,
};
macro_rules! ml_kem_key_impls {
($alg:ident, $dk:ident, $ek:ident, $ct:ident) => {
impl Decapsulator for $dk {
fn decapsulate(&self, ct: &[u8]) -> Result<Secret, Error> {
let bytes: [u8; <$ct>::BYTES] = ct.try_into().map_err(|_| Error::Decapsulation)?;
let ct = <$ct>::from_bytes(bytes);
let ss = self.decapsulate(&ct);
Ok(Secret::from_bytes(ss.to_vec()))
}
}
impl Encapsulator for $ek {
fn encapsulate(&self, rng: &mut dyn CryptoRngCore) -> Result<(Vec<u8>, Secret), Error> {
let mut rng = rng;
let (ct, ss) = self.encapsulate(&mut rng);
Ok((ct.to_bytes().to_vec(), Secret::from_bytes(ss.to_vec())))
}
}
};
}
ml_kem_key_impls!(
MlKem512,
MlKem512DecapsKey,
MlKem512EncapsKey,
MlKem512Ciphertext
);
ml_kem_key_impls!(
MlKem768,
MlKem768DecapsKey,
MlKem768EncapsKey,
MlKem768Ciphertext
);
ml_kem_key_impls!(
MlKem1024,
MlKem1024DecapsKey,
MlKem1024EncapsKey,
MlKem1024Ciphertext
);