hpke-core 0.0.1

//! Extended KEM functions for HPKE as defined in RFC 9180.
//!
//! Here we implement the `LabeledExtract` and `LabeledExpand` functions.

#![allow(clippy::missing_errors_doc)]

pub mod dhkem;

use hpke_crypto::{
    Crypto, EncapsulatedSecret, EncapsulatedSecretRef, HpkeKemId, HpkePrivateKey,
    HpkePrivateKeyRef, HpkePublicKey, HpkePublicKeyRef, IkmRef, SharedSecret,
};

use crate::error::Error;

/// `GenerateKeyPair()`: Randomized algorithm to generate a key pair (skX, pkX).
pub fn generate_key_pair<C: Crypto>(
    crypto_backend: &mut C,
    alg: HpkeKemId,
) -> Result<(HpkePrivateKey, HpkePublicKey), Error> {
    // Currently, we only support DHKEM.
    dhkem::generate_key_pair(crypto_backend, alg)
}

/// `DeriveKeyPair(ikm)`: Deterministic algorithm to derive a key pair(skX, pkX)
/// from the byte string `ikm`, where `ikm` SHOULD have at least `Nsk` bytes of
/// entropy (see [RFC 9180, Section 7.1.3] for discussion).
pub fn derive_key_pair<C: Crypto>(
    crypto_backend: &mut C,
    alg: HpkeKemId,
    ikm: IkmRef<'_>,
) -> Result<(HpkePrivateKey, HpkePublicKey), Error> {
    // Currently, we only support DHKEM.
    dhkem::derive_key_pair(crypto_backend, alg, ikm)
}

/// `Encap(pkR)`: Randomized algorithm to generate an ephemeral,
/// fixed-length symmetric key (the KEM shared secret) and a
/// fixed-length encapsulation of that key that can be decapsulated by
/// the holder of the private key corresponding to `pkR`.
///
/// See [RFC 9180, Section 4] for details.
///
/// [RFC 9180, Section 4]: https://www.rfc-editor.org/rfc/rfc9180.html#section-4
pub fn encap<C: Crypto>(
    crypto_backend: &mut C,
    alg: HpkeKemId,
    pk_r: HpkePublicKeyRef<'_>,
) -> Result<(SharedSecret, EncapsulatedSecret), Error> {
    // Currently, we only support DHKEM.
    dhkem::encap(crypto_backend, alg, pk_r)
}

/// `Decap(enc, skR)`: Deterministic algorithm using the private key `skR` to
/// recover the ephemeral symmetric key (the KEM shared secret) from its
/// encapsulated representation `enc`.
///
/// See [RFC 9180, Section 4] for details.
///
/// [RFC 9180, Section 4]: https://www.rfc-editor.org/rfc/rfc9180.html#section-4
pub fn decap<C: Crypto>(
    crypto_backend: &mut C,
    alg: HpkeKemId,
    enc: EncapsulatedSecretRef<'_>,
    sk_r: HpkePrivateKeyRef<'_>,
) -> Result<SharedSecret, Error> {
    // Currently, we only support DHKEM.
    dhkem::decap(crypto_backend, alg, enc, sk_r)
}

/// `AuthEncap(pkR, skS)` (optional): Same as `Encap()`, and the outputs encode
/// an assurance that the KEM shared secret was generated by the holder of the
/// private key `skS`.
///
/// See [RFC 9180, Section 4] for details.
///
/// [RFC 9180, Section 4]: https://www.rfc-editor.org/rfc/rfc9180.html#section-4
pub fn auth_encap<C: Crypto>(
    crypto_backend: &mut C,
    alg: HpkeKemId,
    pk_r: HpkePublicKeyRef<'_>,
    sk_s: HpkePrivateKeyRef<'_>,
) -> Result<(SharedSecret, EncapsulatedSecret), Error> {
    // Currently, we only support DHKEM.
    dhkem::auth_encap(crypto_backend, alg, pk_r, sk_s)
}

/// `AuthDecap(enc, skR, pkS)` (optional): Same as `Decap()`, and the recipient
/// is assured that the KEM shared secret was generated by the holder of the
/// private key `skS`.
///
/// See [RFC 9180, Section 4] for details.
///
/// [RFC 9180, Section 4]: https://www.rfc-editor.org/rfc/rfc9180.html#section-4
pub fn auth_decap<C: Crypto>(
    crypto_backend: &mut C,
    alg: HpkeKemId,
    enc: EncapsulatedSecretRef<'_>,
    sk_r: HpkePrivateKeyRef<'_>,
    pk_s: HpkePublicKeyRef<'_>,
) -> Result<SharedSecret, Error> {
    // Currently, we only support DHKEM.
    dhkem::auth_decap(crypto_backend, alg, enc, sk_r, pk_s)
}

#[cfg(test)]
mod tests {
    use alloc::format;
    use alloc::vec::Vec;
    use std::panic::catch_unwind;
    use std::println;

    use super::*;
    use crate::{CryptoError, HPKE_TEST_VECTORS};

    #[test_case::test_matrix(
        [
            hpke_crypto::backend::HpkeCryptoAwsLc::new,
            hpke_crypto::backend::HpkeCryptoGraviola::new,
        ]
    )]
    fn kat_derive_key_pair<C: Crypto + Send + Sync + core::panic::UnwindSafe, F>(crypto_backend: F)
    where
        F: Fn() -> Result<C, CryptoError>,
    {
        let mut rets = Vec::new();

        for (idx, test_case) in HPKE_TEST_VECTORS.iter().enumerate() {
            {
                let crypto_backend = crypto_backend().unwrap();
                let ret = catch_unwind(move || {
                    test_derive_key_pair_each(
                        crypto_backend,
                        idx,
                        "R",
                        test_case.kem_id,
                        &test_case.ikm_r,
                        &test_case.sk_rm,
                        &test_case.pk_rm,
                    );
                });

                rets.push((format!("R-{}({})", test_case.kem_id, idx), ret));
            }

            {
                let crypto_backend = crypto_backend().unwrap();
                let ret = catch_unwind(move || {
                    test_derive_key_pair_each(
                        crypto_backend,
                        idx,
                        "E",
                        test_case.kem_id,
                        &test_case.ikm_e,
                        &test_case.sk_em,
                        &test_case.pk_em,
                    );
                });

                rets.push((format!("E-{}({})", test_case.kem_id, idx), ret));
            }

            if test_case.ikm_s.is_some() {
                let crypto_backend = crypto_backend().unwrap();
                let ret = catch_unwind(move || {
                    test_derive_key_pair_each(
                        crypto_backend,
                        idx,
                        "E",
                        test_case.kem_id,
                        &test_case.ikm_s.as_ref().unwrap(),
                        &test_case.sk_sm.as_ref().unwrap(),
                        &test_case.pk_sm.as_ref().unwrap(),
                    );
                });

                rets.push((format!("S-{}({})", test_case.kem_id, idx), ret));
            }
        }

        let errors: Vec<_> = rets
            .iter()
            .filter(|(_, ret)| ret.is_err())
            .collect();

        if !errors.is_empty() {
            for (name, err) in &errors {
                println!("[FAILED] {name}: {err:?}");
            }

            panic!("{} test cases failed", errors.len());
        }
    }

    fn test_derive_key_pair_each<C: Crypto>(
        mut crypto_backend: C,
        idx: usize,
        role: &'static str,
        alg: HpkeKemId,
        ikm: &[u8],
        expected_sk: &[u8],
        expected_pk: &[u8],
    ) {
        if !crypto_backend.is_kem_supported(&alg) {
            // Skip unsupported KEMs.
            println!(
                "[{name}][{idx}][{role}] Skipping, unsupported KEM {alg:?}",
                name = core::any::type_name::<C>(),
            );
            return;
        }

        let (sk, pk) = derive_key_pair(&mut crypto_backend, alg, ikm.into()).unwrap();

        assert_eq!(&*sk, expected_sk);
        assert_eq!(&*pk, expected_pk);
    }
}