use alloc::vec::Vec;
use zeroize::{Zeroize, ZeroizeOnDrop, Zeroizing};
use pakery_core::crypto::{Hash, Kdf, Mac};
use pakery_core::SharedSecret;
use crate::ciphersuite::Spake2PlusCiphersuite;
use crate::error::Spake2PlusError;
#[derive(Zeroize, ZeroizeOnDrop)]
pub struct Spake2PlusOutput {
#[zeroize(skip)]
pub session_key: SharedSecret,
}
impl Spake2PlusOutput {
#[must_use]
pub fn into_session_key(mut self) -> SharedSecret {
core::mem::replace(&mut self.session_key, SharedSecret::new(Vec::new()))
}
}
pub(crate) struct KeySchedule {
pub confirm_p: Vec<u8>,
pub confirm_v: Vec<u8>,
pub session_key: SharedSecret,
}
impl Drop for KeySchedule {
fn drop(&mut self) {
self.confirm_p.zeroize();
self.confirm_v.zeroize();
}
}
pub(crate) fn derive_key_schedule<C: Spake2PlusCiphersuite>(
tt: &[u8],
share_p: &[u8],
share_v: &[u8],
) -> Result<KeySchedule, Spake2PlusError> {
const { assert!(<C::Hash as pakery_core::crypto::Hash>::OUTPUT_SIZE >= C::NH) };
let k_main = Zeroizing::new(C::Hash::digest(tt));
let prk = C::Kdf::extract(&[], &k_main[..C::NH]);
let kc = C::Kdf::expand(&prk, b"ConfirmationKeys", 2 * C::NH)
.map_err(|_| Spake2PlusError::InternalError("KDF expand failed for ConfirmationKeys"))?;
let k_confirm_p = &kc[..C::NH];
let k_confirm_v = &kc[C::NH..2 * C::NH];
let mut k_shared = C::Kdf::expand(&prk, b"SharedKey", C::NH)
.map_err(|_| Spake2PlusError::InternalError("KDF expand failed for SharedKey"))?;
let confirm_v = C::Mac::mac(k_confirm_v, share_p)
.map_err(|_| Spake2PlusError::InternalError("MAC computation failed"))?;
let confirm_p = C::Mac::mac(k_confirm_p, share_v)
.map_err(|_| Spake2PlusError::InternalError("MAC computation failed"))?;
Ok(KeySchedule {
confirm_p,
confirm_v,
session_key: SharedSecret::new(core::mem::take(&mut *k_shared)),
})
}