use crate::kem::{Compress, Error, SharedSecret, IBKEM, SS_BYTES};
use core::slice::Iter;
use rand::{CryptoRng, Rng};
use subtle::CtOption;
use aes_gcm::aead::{Nonce, Tag};
use aes_gcm::{AeadInPlace, Aes128Gcm, KeyInit};
#[cfg(feature = "cgwfo")]
use crate::kem::cgw_fo::CGWFO;
#[cfg(feature = "cgwkv")]
use crate::kem::cgw_kv::CGWKV;
#[cfg(feature = "kv1")]
use crate::kem::kiltz_vahlis_one::KV1;
const TAG_SIZE: usize = 16;
const NONCE_SIZE: usize = 12;
const KEY_SIZE: usize = 16;
impl SharedSecret {
fn random<R: Rng + CryptoRng>(r: &mut R) -> Self {
let mut ss_bytes = [0u8; SS_BYTES];
r.fill_bytes(&mut ss_bytes);
SharedSecret(ss_bytes)
}
}
#[derive(Debug, Clone)]
pub struct Ciphertext<K: IBKEM> {
ct_asymm: K::Ct,
ct_symm: [u8; SS_BYTES],
tag: Tag<Aes128Gcm>,
nonce: Nonce<Aes128Gcm>,
}
#[derive(Debug)]
pub struct Ciphertexts<'a, K: IBKEM, R> {
ss: SharedSecret,
pk: &'a K::Pk,
ids: Iter<'a, K::Id>,
rng: &'a mut R,
}
impl<'a, K, R> Iterator for Ciphertexts<'a, K, R>
where
K: IBKEM,
R: Rng + CryptoRng,
{
type Item = Ciphertext<K>;
fn next(&mut self) -> Option<Self::Item> {
let id = self.ids.next()?;
let (ct_asymm, kek) = <K as IBKEM>::encaps(self.pk, id, self.rng);
let aead = Aes128Gcm::new_from_slice(&kek.0[..KEY_SIZE]).unwrap();
let nonce_bytes = self.rng.gen::<[u8; NONCE_SIZE]>();
let nonce = Nonce::<Aes128Gcm>::from_slice(&nonce_bytes);
let mut shared_key = self.ss.0;
let tag = aead
.encrypt_in_place_detached(nonce, b"", &mut shared_key)
.unwrap();
Some(Ciphertext::<K> {
ct_asymm,
ct_symm: shared_key,
nonce: *nonce,
tag,
})
}
}
pub trait MultiRecipient: IBKEM {
fn multi_encaps<'a, R: Rng + CryptoRng>(
pk: &'a <Self as IBKEM>::Pk,
ids: impl IntoIterator<IntoIter = Iter<'a, Self::Id>>,
rng: &'a mut R,
) -> (Ciphertexts<'a, Self, R>, SharedSecret) {
let ss = SharedSecret::random(rng);
let cts = Ciphertexts {
ss,
pk,
rng,
ids: ids.into_iter(),
};
(cts, ss)
}
fn multi_decaps(
mpk: Option<&Self::Pk>,
usk: &Self::Usk,
ct: &Ciphertext<Self>,
) -> Result<SharedSecret, Error> {
let kek = <Self as IBKEM>::decaps(mpk, usk, &ct.ct_asymm)?;
let aead = Aes128Gcm::new_from_slice(&kek.0[..KEY_SIZE]).unwrap();
let mut shared_key = ct.ct_symm;
aead.decrypt_in_place_detached(&ct.nonce, b"", &mut shared_key, &ct.tag)
.map_err(|_e| Error)?;
Ok(SharedSecret(shared_key))
}
}
macro_rules! impl_mkemct_compress {
($scheme: ident) => {
impl Compress for Ciphertext<$scheme> {
const OUTPUT_SIZE: usize = $scheme::CT_BYTES + SS_BYTES + TAG_SIZE + NONCE_SIZE;
type Output = [u8; Self::OUTPUT_SIZE];
fn to_bytes(&self) -> Self::Output {
use arrayref::mut_array_refs;
let mut res = [0u8; Self::OUTPUT_SIZE];
let (ct_asymm, ct_symm, tag, nonce) =
mut_array_refs![&mut res, $scheme::CT_BYTES, SS_BYTES, TAG_SIZE, NONCE_SIZE];
*ct_asymm = self.ct_asymm.to_bytes();
*ct_symm = self.ct_symm;
*tag = self.tag.into();
*nonce = self.nonce.into();
res
}
fn from_bytes(output: &Self::Output) -> CtOption<Self> {
use arrayref::array_refs;
let (ct_asymm, ct_symm, tag, nonce) =
array_refs![&output, $scheme::CT_BYTES, SS_BYTES, TAG_SIZE, NONCE_SIZE];
let ct_asymm = <$scheme as IBKEM>::Ct::from_bytes(ct_asymm);
let tag = Tag::<Aes128Gcm>::from_slice(tag);
let nonce = Nonce::<Aes128Gcm>::from_slice(nonce);
ct_asymm.map(|ct_asymm| Ciphertext {
ct_asymm,
ct_symm: *ct_symm,
tag: *tag,
nonce: *nonce,
})
}
}
};
}
#[cfg(feature = "cgwkv")]
impl_mkemct_compress!(CGWKV);
#[cfg(feature = "cgwfo")]
impl_mkemct_compress!(CGWFO);
#[cfg(feature = "kv1")]
impl_mkemct_compress!(KV1);