use alloc::vec::Vec;
use hkdf::Hkdf;
use rand::CryptoRng;
use sha2::{Sha256, digest::OutputSizeUser};
use thiserror::Error;
use crate::utils::{
Deserializable, Serializable,
zeroize::{Zeroize, ZeroizeOnDrop},
};
pub mod k256;
pub mod x25519;
pub(crate) trait KeyAgreementScheme {
type EphemeralSecretKey: ZeroizeOnDrop;
type EphemeralPublicKey: Serializable + Deserializable;
type SecretKey;
type PublicKey: Clone;
type SharedSecret: AsRef<[u8]> + Zeroize + ZeroizeOnDrop;
fn generate_ephemeral_keypair<R: CryptoRng>(
rng: &mut R,
) -> (Self::EphemeralSecretKey, Self::EphemeralPublicKey);
fn exchange_ephemeral_static(
ephemeral_sk: Self::EphemeralSecretKey,
static_pk: &Self::PublicKey,
) -> Result<Self::SharedSecret, KeyAgreementError>;
fn exchange_static_ephemeral(
static_sk: &Self::SecretKey,
ephemeral_pk: &Self::EphemeralPublicKey,
) -> Result<Self::SharedSecret, KeyAgreementError>;
fn extract_key_material(
shared_secret: &Self::SharedSecret,
length: usize,
info: &[u8],
) -> Result<Vec<u8>, KeyAgreementError>;
}
pub fn extract_key_material(
shared_secret: &[u8],
salt: Option<&[u8]>,
length: usize,
info: &[u8],
) -> Result<Vec<u8>, KeyAgreementError> {
if length > 255 * Sha256::output_size() {
return Err(KeyAgreementError::HkdfExpansionFailed);
}
let hkdf = Hkdf::<Sha256>::new(salt, shared_secret);
let mut buf = vec![0_u8; length];
hkdf.expand(info, &mut buf)
.map_err(|_| KeyAgreementError::HkdfExpansionFailed)?;
Ok(buf)
}
#[derive(Debug, Error)]
pub enum KeyAgreementError {
#[error("hkdf expansion failed")]
HkdfExpansionFailed,
#[error("shared secret is invalid")]
InvalidSharedSecret,
}