/rust/registry/src/index.crates.io-1949cf8c6b5b557f/aws-lc-rs-1.14.1/src/kem.rs

Source
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0 OR ISC

//! Key-Encapsulation Mechanisms (KEMs), including support for Kyber Round 3 Submission.
//!
//! # Example
//!
//! Note that this example uses the Kyber-512 Round 3 algorithm, but other algorithms can be used
//! in the exact same way by substituting
//! `kem::<desired_algorithm_here>` for `kem::KYBER512_R3`.
//!
//! ```rust
//! use aws_lc_rs::{
//!     kem::{Ciphertext, DecapsulationKey, EncapsulationKey},
//!     kem::{ML_KEM_512}
//! };
//!
//! // Alice generates their (private) decapsulation key.
//! let decapsulation_key = DecapsulationKey::generate(&ML_KEM_512)?;
//!
//! // Alices computes the (public) encapsulation key.
//! let encapsulation_key = decapsulation_key.encapsulation_key()?;
//!
//! let encapsulation_key_bytes = encapsulation_key.key_bytes()?;
//!
//! // Alice sends the encapsulation key bytes to bob through some
//! // protocol message.
//! let encapsulation_key_bytes = encapsulation_key_bytes.as_ref();
//!
//! // Bob constructs the (public) encapsulation key from the key bytes provided by Alice.
//! let retrieved_encapsulation_key = EncapsulationKey::new(&ML_KEM_512, encapsulation_key_bytes)?;
//!
//! // Bob executes the encapsulation algorithm to to produce their copy of the secret, and associated ciphertext.
//! let (ciphertext, bob_secret) = retrieved_encapsulation_key.encapsulate()?;
//!
//! // Alice receives ciphertext bytes from bob
//! let ciphertext_bytes = ciphertext.as_ref();
//!
//! // Bob sends Alice the ciphertext computed from the encapsulation algorithm, Alice runs decapsulation to derive their
//! // copy of the secret.
//! let alice_secret = decapsulation_key.decapsulate(Ciphertext::from(ciphertext_bytes))?;
//!
//! // Alice and Bob have now arrived to the same secret
//! assert_eq!(alice_secret.as_ref(), bob_secret.as_ref());
//!
//! # Ok::<(), aws_lc_rs::error::Unspecified>(())
//! ```
use crate::aws_lc::{
    EVP_PKEY_CTX_kem_set_params, EVP_PKEY_decapsulate, EVP_PKEY_encapsulate,
    EVP_PKEY_kem_new_raw_public_key, EVP_PKEY, EVP_PKEY_KEM,
};
use crate::buffer::Buffer;
use crate::encoding::generated_encodings;
use crate::error::{KeyRejected, Unspecified};
use crate::ptr::LcPtr;
use alloc::borrow::Cow;
use core::cmp::Ordering;
use zeroize::Zeroize;

const ML_KEM_512_SHARED_SECRET_LENGTH: usize = 32;
const ML_KEM_512_PUBLIC_KEY_LENGTH: usize = 800;
const ML_KEM_512_SECRET_KEY_LENGTH: usize = 1632;
const ML_KEM_512_CIPHERTEXT_LENGTH: usize = 768;

const ML_KEM_768_SHARED_SECRET_LENGTH: usize = 32;
const ML_KEM_768_PUBLIC_KEY_LENGTH: usize = 1184;
const ML_KEM_768_SECRET_KEY_LENGTH: usize = 2400;
const ML_KEM_768_CIPHERTEXT_LENGTH: usize = 1088;

const ML_KEM_1024_SHARED_SECRET_LENGTH: usize = 32;
const ML_KEM_1024_PUBLIC_KEY_LENGTH: usize = 1568;
const ML_KEM_1024_SECRET_KEY_LENGTH: usize = 3168;
const ML_KEM_1024_CIPHERTEXT_LENGTH: usize = 1568;

/// NIST FIPS 203 ML-KEM-512 algorithm.
pub const ML_KEM_512: Algorithm<AlgorithmId> = Algorithm {
    id: AlgorithmId::MlKem512,
    decapsulate_key_size: ML_KEM_512_SECRET_KEY_LENGTH,
    encapsulate_key_size: ML_KEM_512_PUBLIC_KEY_LENGTH,
    ciphertext_size: ML_KEM_512_CIPHERTEXT_LENGTH,
    shared_secret_size: ML_KEM_512_SHARED_SECRET_LENGTH,
};

/// NIST FIPS 203 ML-KEM-768 algorithm.
pub const ML_KEM_768: Algorithm<AlgorithmId> = Algorithm {
    id: AlgorithmId::MlKem768,
    decapsulate_key_size: ML_KEM_768_SECRET_KEY_LENGTH,
    encapsulate_key_size: ML_KEM_768_PUBLIC_KEY_LENGTH,
    ciphertext_size: ML_KEM_768_CIPHERTEXT_LENGTH,
    shared_secret_size: ML_KEM_768_SHARED_SECRET_LENGTH,
};

/// NIST FIPS 203 ML-KEM-1024 algorithm.
pub const ML_KEM_1024: Algorithm<AlgorithmId> = Algorithm {
    id: AlgorithmId::MlKem1024,
    decapsulate_key_size: ML_KEM_1024_SECRET_KEY_LENGTH,
    encapsulate_key_size: ML_KEM_1024_PUBLIC_KEY_LENGTH,
    ciphertext_size: ML_KEM_1024_CIPHERTEXT_LENGTH,
    shared_secret_size: ML_KEM_1024_SHARED_SECRET_LENGTH,
};

use crate::aws_lc::{NID_MLKEM1024, NID_MLKEM512, NID_MLKEM768};

/// An identifier for a KEM algorithm.
pub trait AlgorithmIdentifier:
    Copy + Clone + Debug + PartialEq + crate::sealed::Sealed + 'static
{
    /// Returns the algorithm's associated AWS-LC nid.
    fn nid(self) -> i32;
}

/// A KEM algorithm
#[derive(PartialEq)]
pub struct Algorithm<Id = AlgorithmId>
where
    Id: AlgorithmIdentifier,
{
    pub(crate) id: Id,
    pub(crate) decapsulate_key_size: usize,
    pub(crate) encapsulate_key_size: usize,
    pub(crate) ciphertext_size: usize,
    pub(crate) shared_secret_size: usize,
}

impl<Id> Algorithm<Id>
where
    Id: AlgorithmIdentifier,
{
    /// Returns the identifier for this algorithm.
    #[must_use]
    pub fn id(&self) -> Id {
        self.id
    }

    #[inline]
    #[allow(dead_code)]
    pub(crate) fn decapsulate_key_size(&self) -> usize {
        self.decapsulate_key_size
    }

    #[inline]
    pub(crate) fn encapsulate_key_size(&self) -> usize {
        self.encapsulate_key_size
    }

    #[inline]
    pub(crate) fn ciphertext_size(&self) -> usize {
        self.ciphertext_size
    }

    #[inline]
    pub(crate) fn shared_secret_size(&self) -> usize {
        self.shared_secret_size
    }
}

impl<Id> Debug for Algorithm<Id>
where
    Id: AlgorithmIdentifier,
{
    fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
        Debug::fmt(&self.id, f)
    }
}

/// A serializable decapulsation key usable with KEMs. This can be randomly generated with `DecapsulationKey::generate`.
pub struct DecapsulationKey<Id = AlgorithmId>
where
    Id: AlgorithmIdentifier,
{
    algorithm: &'static Algorithm<Id>,
    evp_pkey: LcPtr<EVP_PKEY>,
}

/// Identifier for a KEM algorithm.
#[non_exhaustive]
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum AlgorithmId {
    /// NIST FIPS 203 ML-KEM-512 algorithm.
    MlKem512,

    /// NIST FIPS 203 ML-KEM-768 algorithm.
    MlKem768,

    /// NIST FIPS 203 ML-KEM-1024 algorithm.
    MlKem1024,
}

impl AlgorithmIdentifier for AlgorithmId {
    fn nid(self) -> i32 {
        match self {
            AlgorithmId::MlKem512 => NID_MLKEM512,
            AlgorithmId::MlKem768 => NID_MLKEM768,
            AlgorithmId::MlKem1024 => NID_MLKEM1024,
        }
    }
}

impl crate::sealed::Sealed for AlgorithmId {}

impl<Id> DecapsulationKey<Id>
where
    Id: AlgorithmIdentifier,
{
    /// Generate a new KEM decapsulation key for the given algorithm.
    ///
    /// # Errors
    /// `error::Unspecified` when operation fails due to internal error.
    pub fn generate(alg: &'static Algorithm<Id>) -> Result<Self, Unspecified> {
        let kyber_key = kem_key_generate(alg.id.nid())?;
        Ok(DecapsulationKey {
            algorithm: alg,
            evp_pkey: kyber_key,
        })
    }

    /// Return the algorithm associated with the given KEM decapsulation key.
    #[must_use]
    pub fn algorithm(&self) -> &'static Algorithm<Id> {
        self.algorithm
    }

    /// Computes the KEM encapsulation key from the KEM decapsulation key.
    ///
    /// # Errors
    /// `error::Unspecified` when operation fails due to internal error.
    #[allow(clippy::missing_panics_doc)]
    pub fn encapsulation_key(&self) -> Result<EncapsulationKey<Id>, Unspecified> {
        let evp_pkey = self.evp_pkey.clone();

        Ok(EncapsulationKey {
            algorithm: self.algorithm,
            evp_pkey,
        })
    }

    /// Performs the decapsulate operation using this KEM decapsulation key on the given ciphertext.
    ///
    /// `ciphertext` is the ciphertext generated by the encapsulate operation using the KEM encapsulation key
    /// associated with this KEM decapsulation key.
    ///
    /// # Errors
    /// `Unspecified` when operation fails due to internal error.
    #[allow(clippy::needless_pass_by_value)]
    pub fn decapsulate(&self, ciphertext: Ciphertext<'_>) -> Result<SharedSecret, Unspecified> {
        let mut shared_secret_len = self.algorithm.shared_secret_size();
        let mut shared_secret: Vec<u8> = vec![0u8; shared_secret_len];

        let mut ctx = self.evp_pkey.create_EVP_PKEY_CTX()?;

        let ciphertext = ciphertext.as_ref();

        if 1 != unsafe {
            EVP_PKEY_decapsulate(
                *ctx.as_mut(),
                shared_secret.as_mut_ptr(),
                &mut shared_secret_len,
                // AWS-LC incorrectly has this as an unqualified `uint8_t *`, it should be qualified with const
                ciphertext.as_ptr().cast_mut(),
                ciphertext.len(),
            )
        } {
            return Err(Unspecified);
        }

        // This is currently pedantic but done for safety in-case the shared_secret buffer
        // size changes in the future. `EVP_PKEY_decapsulate` updates `shared_secret_len` with
        // the length of the shared secret in the event the buffer provided was larger then the secret.
        // This truncates the buffer to the proper length to match the shared secret written.
        debug_assert_eq!(shared_secret_len, shared_secret.len());
        shared_secret.truncate(shared_secret_len);

        Ok(SharedSecret(shared_secret.into_boxed_slice()))
    }
}

unsafe impl<Id> Send for DecapsulationKey<Id> where Id: AlgorithmIdentifier {}

unsafe impl<Id> Sync for DecapsulationKey<Id> where Id: AlgorithmIdentifier {}

impl<Id> Debug for DecapsulationKey<Id>
where
    Id: AlgorithmIdentifier,
{
    fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("DecapsulationKey")
            .field("algorithm", &self.algorithm)
            .finish_non_exhaustive()
    }
}

generated_encodings!((EncapsulationKeyBytes, EncapsulationKeyBytesType));

/// A serializable encapsulation key usable with KEM algorithms. Constructed
/// from either a `DecapsulationKey` or raw bytes.
pub struct EncapsulationKey<Id = AlgorithmId>
where
    Id: AlgorithmIdentifier,
{
    algorithm: &'static Algorithm<Id>,
    evp_pkey: LcPtr<EVP_PKEY>,
}

impl<Id> EncapsulationKey<Id>
where
    Id: AlgorithmIdentifier,
{
    /// Return the algorithm associated with the given KEM encapsulation key.
    #[must_use]
    pub fn algorithm(&self) -> &'static Algorithm<Id> {
        self.algorithm
    }

    /// Performs the encapsulate operation using this KEM encapsulation key, generating a ciphertext
    /// and associated shared secret.
    ///
    /// # Errors
    /// `error::Unspecified` when operation fails due to internal error.
    pub fn encapsulate(&self) -> Result<(Ciphertext<'static>, SharedSecret), Unspecified> {
        let mut ciphertext_len = self.algorithm.ciphertext_size();
        let mut shared_secret_len = self.algorithm.shared_secret_size();
        let mut ciphertext: Vec<u8> = vec![0u8; ciphertext_len];
        let mut shared_secret: Vec<u8> = vec![0u8; shared_secret_len];

        let mut ctx = self.evp_pkey.create_EVP_PKEY_CTX()?;

        if 1 != unsafe {
            EVP_PKEY_encapsulate(
                *ctx.as_mut(),
                ciphertext.as_mut_ptr(),
                &mut ciphertext_len,
                shared_secret.as_mut_ptr(),
                &mut shared_secret_len,
            )
        } {
            return Err(Unspecified);
        }

        // The following two steps are currently pedantic but done for safety in-case the buffer allocation
        // sizes change in the future. `EVP_PKEY_encapsulate` updates `ciphertext_len` and `shared_secret_len` with
        // the length of the ciphertext and shared secret respectivly in the event the buffer provided for each was
        // larger then the actual values. Thus these two steps truncate the buffers to the proper length to match the
        // value lengths written.
        debug_assert_eq!(ciphertext_len, ciphertext.len());
        ciphertext.truncate(ciphertext_len);
        debug_assert_eq!(shared_secret_len, shared_secret.len());
        shared_secret.truncate(shared_secret_len);

        Ok((
            Ciphertext::new(ciphertext),
            SharedSecret::new(shared_secret.into_boxed_slice()),
        ))
    }

    /// Returns the `EnscapsulationKey` bytes.
    ///
    /// # Errors
    /// * `Unspecified`: Any failure to retrieve the `EnscapsulationKey` bytes.
    pub fn key_bytes(&self) -> Result<EncapsulationKeyBytes<'static>, Unspecified> {
        let mut encapsulate_bytes = vec![0u8; self.algorithm.encapsulate_key_size()];
        let encapsulate_key_size = self
            .evp_pkey
            .as_const()
            .marshal_raw_public_to_buffer(&mut encapsulate_bytes)?;

        debug_assert_eq!(encapsulate_key_size, encapsulate_bytes.len());
        encapsulate_bytes.truncate(encapsulate_key_size);

        Ok(EncapsulationKeyBytes::new(encapsulate_bytes))
    }

    /// Creates a new KEM encapsulation key from raw bytes. This method MUST NOT be used to generate
    /// a new encapsulation key, rather it MUST be used to construct `EncapsulationKey` previously serialized
    /// to raw bytes.
    ///
    /// `alg` is the [`Algorithm`] to be associated with the generated `EncapsulationKey`.
    ///
    /// `bytes` is a slice of raw bytes representing a `EncapsulationKey`.
    ///
    /// # Errors
    /// `error::KeyRejected` when operation fails during key creation.
    pub fn new(alg: &'static Algorithm<Id>, bytes: &[u8]) -> Result<Self, KeyRejected> {
        match bytes.len().cmp(&alg.encapsulate_key_size()) {
            Ordering::Less => Err(KeyRejected::too_small()),
            Ordering::Greater => Err(KeyRejected::too_large()),
            Ordering::Equal => Ok(()),
        }?;
        let pubkey = LcPtr::new(unsafe {
            EVP_PKEY_kem_new_raw_public_key(alg.id.nid(), bytes.as_ptr(), bytes.len())
        })?;
        Ok(EncapsulationKey {
            algorithm: alg,
            evp_pkey: pubkey,
        })
    }
}

unsafe impl<Id> Send for EncapsulationKey<Id> where Id: AlgorithmIdentifier {}

unsafe impl<Id> Sync for EncapsulationKey<Id> where Id: AlgorithmIdentifier {}

impl<Id> Debug for EncapsulationKey<Id>
where
    Id: AlgorithmIdentifier,
{
    fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("EncapsulationKey")
            .field("algorithm", &self.algorithm)
            .finish_non_exhaustive()
    }
}

/// A set of encrypted bytes produced by [`EncapsulationKey::encapsulate`],
/// and used as an input to [`DecapsulationKey::decapsulate`].
pub struct Ciphertext<'a>(Cow<'a, [u8]>);

impl<'a> Ciphertext<'a> {
    fn new(value: Vec<u8>) -> Ciphertext<'a> {
        Self(Cow::Owned(value))
    }
}

impl Drop for Ciphertext<'_> {
    fn drop(&mut self) {
        if let Cow::Owned(ref mut v) = self.0 {
            v.zeroize();
        }
    }
}

impl AsRef<[u8]> for Ciphertext<'_> {
    fn as_ref(&self) -> &[u8] {
        match self.0 {
            Cow::Borrowed(v) => v,
            Cow::Owned(ref v) => v.as_ref(),
        }
    }
}

impl<'a> From<&'a [u8]> for Ciphertext<'a> {
    fn from(value: &'a [u8]) -> Self {
        Self(Cow::Borrowed(value))
    }
}

/// The cryptographic shared secret output from the KEM encapsulate / decapsulate process.
pub struct SharedSecret(Box<[u8]>);

impl SharedSecret {
    fn new(value: Box<[u8]>) -> Self {
        Self(value)
    }
}

impl Drop for SharedSecret {
    fn drop(&mut self) {
        self.0.zeroize();
    }
}

impl AsRef<[u8]> for SharedSecret {
    fn as_ref(&self) -> &[u8] {
        self.0.as_ref()
    }
}

// Returns an LcPtr to an EVP_PKEY
#[inline]
fn kem_key_generate(nid: i32) -> Result<LcPtr<EVP_PKEY>, Unspecified> {
    let params_fn = |ctx| {
        if 1 == unsafe { EVP_PKEY_CTX_kem_set_params(ctx, nid) } {
            Ok(())
        } else {
            Err(())
        }
    };

    LcPtr::<EVP_PKEY>::generate(EVP_PKEY_KEM, Some(params_fn))
}

#[cfg(test)]
mod tests {
    use super::{Ciphertext, DecapsulationKey, EncapsulationKey, SharedSecret};
    use crate::error::KeyRejected;

    use crate::kem::{ML_KEM_1024, ML_KEM_512, ML_KEM_768};

    #[test]
    fn ciphertext() {
        let ciphertext_bytes = vec![42u8; 4];
        let ciphertext = Ciphertext::from(ciphertext_bytes.as_ref());
        assert_eq!(ciphertext.as_ref(), &[42, 42, 42, 42]);
        drop(ciphertext);

        let ciphertext_bytes = vec![42u8; 4];
        let ciphertext = Ciphertext::<'static>::new(ciphertext_bytes);
        assert_eq!(ciphertext.as_ref(), &[42, 42, 42, 42]);
    }

    #[test]
    fn shared_secret() {
        let secret_bytes = vec![42u8; 4];
        let shared_secret = SharedSecret::new(secret_bytes.into_boxed_slice());
        assert_eq!(shared_secret.as_ref(), &[42, 42, 42, 42]);
    }

    #[test]
    fn test_kem_serialize() {
        for algorithm in [&ML_KEM_512, &ML_KEM_768, &ML_KEM_1024] {
            let priv_key = DecapsulationKey::generate(algorithm).unwrap();
            assert_eq!(priv_key.algorithm(), algorithm);

            let pub_key = priv_key.encapsulation_key().unwrap();
            let pubkey_raw_bytes = pub_key.key_bytes().unwrap();
            let pub_key_from_bytes =
                EncapsulationKey::new(algorithm, pubkey_raw_bytes.as_ref()).unwrap();

            assert_eq!(
                pub_key.key_bytes().unwrap().as_ref(),
                pub_key_from_bytes.key_bytes().unwrap().as_ref()
            );
            assert_eq!(pub_key.algorithm(), pub_key_from_bytes.algorithm());
        }
    }

    #[test]
    fn test_kem_wrong_sizes() {
        for algorithm in [&ML_KEM_512, &ML_KEM_768, &ML_KEM_1024] {
            let too_long_bytes = vec![0u8; algorithm.encapsulate_key_size() + 1];
            let long_pub_key_from_bytes = EncapsulationKey::new(algorithm, &too_long_bytes);
            assert_eq!(
                long_pub_key_from_bytes.err(),
                Some(KeyRejected::too_large())
            );

            let too_short_bytes = vec![0u8; algorithm.encapsulate_key_size() - 1];
            let short_pub_key_from_bytes = EncapsulationKey::new(algorithm, &too_short_bytes);
            assert_eq!(
                short_pub_key_from_bytes.err(),
                Some(KeyRejected::too_small())
            );
        }
    }

    #[test]
    fn test_kem_e2e() {
        for algorithm in [&ML_KEM_512, &ML_KEM_768, &ML_KEM_1024] {
            let priv_key = DecapsulationKey::generate(algorithm).unwrap();
            assert_eq!(priv_key.algorithm(), algorithm);

            let pub_key = priv_key.encapsulation_key().unwrap();

            let (alice_ciphertext, alice_secret) =
                pub_key.encapsulate().expect("encapsulate successful");

            let bob_secret = priv_key
                .decapsulate(alice_ciphertext)
                .expect("decapsulate successful");

            assert_eq!(alice_secret.as_ref(), bob_secret.as_ref());
        }
    }

    #[test]
    fn test_serialized_kem_e2e() {
        for algorithm in [&ML_KEM_512, &ML_KEM_768, &ML_KEM_1024] {
            let priv_key = DecapsulationKey::generate(algorithm).unwrap();
            assert_eq!(priv_key.algorithm(), algorithm);

            let pub_key = priv_key.encapsulation_key().unwrap();

            // Generate public key bytes to send to bob
            let pub_key_bytes = pub_key.key_bytes().unwrap();

            // Test that priv_key's EVP_PKEY isn't entirely freed since we remove this pub_key's reference.
            drop(pub_key);

            let retrieved_pub_key =
                EncapsulationKey::new(algorithm, pub_key_bytes.as_ref()).unwrap();
            let (ciphertext, bob_secret) = retrieved_pub_key
                .encapsulate()
                .expect("encapsulate successful");

            let alice_secret = priv_key
                .decapsulate(ciphertext)
                .expect("decapsulate successful");

            assert_eq!(alice_secret.as_ref(), bob_secret.as_ref());
        }
    }

    #[test]
    fn test_debug_fmt() {
        let private = DecapsulationKey::generate(&ML_KEM_512).expect("successful generation");
        assert_eq!(
            format!("{private:?}"),
            "DecapsulationKey { algorithm: MlKem512, .. }"
        );
        assert_eq!(
            format!(
                "{:?}",
                private.encapsulation_key().expect("public key retrievable")
            ),
            "EncapsulationKey { algorithm: MlKem512, .. }"
        );
    }
}

Coverage Report

Created: 2025-11-16 06:35

Line	Count	Source
1		// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2		// SPDX-License-Identifier: Apache-2.0 OR ISC
3
4		//! Key-Encapsulation Mechanisms (KEMs), including support for Kyber Round 3 Submission.
5		//!
6		//! # Example
7		//!
8		//! Note that this example uses the Kyber-512 Round 3 algorithm, but other algorithms can be used
9		//! in the exact same way by substituting
10		//! `kem::<desired_algorithm_here>` for `kem::KYBER512_R3`.
11		//!
12		//! ```rust
13		//! use aws_lc_rs::{
14		//! kem::{Ciphertext, DecapsulationKey, EncapsulationKey},
15		//! kem::{ML_KEM_512}
16		//! };
17		//!
18		//! // Alice generates their (private) decapsulation key.
19		//! let decapsulation_key = DecapsulationKey::generate(&ML_KEM_512)?;
20		//!
21		//! // Alices computes the (public) encapsulation key.
22		//! let encapsulation_key = decapsulation_key.encapsulation_key()?;
23		//!
24		//! let encapsulation_key_bytes = encapsulation_key.key_bytes()?;
25		//!
26		//! // Alice sends the encapsulation key bytes to bob through some
27		//! // protocol message.
28		//! let encapsulation_key_bytes = encapsulation_key_bytes.as_ref();
29		//!
30		//! // Bob constructs the (public) encapsulation key from the key bytes provided by Alice.
31		//! let retrieved_encapsulation_key = EncapsulationKey::new(&ML_KEM_512, encapsulation_key_bytes)?;
32		//!
33		//! // Bob executes the encapsulation algorithm to to produce their copy of the secret, and associated ciphertext.
34		//! let (ciphertext, bob_secret) = retrieved_encapsulation_key.encapsulate()?;
35		//!
36		//! // Alice receives ciphertext bytes from bob
37		//! let ciphertext_bytes = ciphertext.as_ref();
38		//!
39		//! // Bob sends Alice the ciphertext computed from the encapsulation algorithm, Alice runs decapsulation to derive their
40		//! // copy of the secret.
41		//! let alice_secret = decapsulation_key.decapsulate(Ciphertext::from(ciphertext_bytes))?;
42		//!
43		//! // Alice and Bob have now arrived to the same secret
44		//! assert_eq!(alice_secret.as_ref(), bob_secret.as_ref());
45		//!
46		//! # Ok::<(), aws_lc_rs::error::Unspecified>(())
47		//! ```
48		use crate::aws_lc::{
49		EVP_PKEY_CTX_kem_set_params, EVP_PKEY_decapsulate, EVP_PKEY_encapsulate,
50		EVP_PKEY_kem_new_raw_public_key, EVP_PKEY, EVP_PKEY_KEM,
51		};
52		use crate::buffer::Buffer;
53		use crate::encoding::generated_encodings;
54		use crate::error::{KeyRejected, Unspecified};
55		use crate::ptr::LcPtr;
56		use alloc::borrow::Cow;
57		use core::cmp::Ordering;
58		use zeroize::Zeroize;
59
60		const ML_KEM_512_SHARED_SECRET_LENGTH: usize = 32;
61		const ML_KEM_512_PUBLIC_KEY_LENGTH: usize = 800;
62		const ML_KEM_512_SECRET_KEY_LENGTH: usize = 1632;
63		const ML_KEM_512_CIPHERTEXT_LENGTH: usize = 768;
64
65		const ML_KEM_768_SHARED_SECRET_LENGTH: usize = 32;
66		const ML_KEM_768_PUBLIC_KEY_LENGTH: usize = 1184;
67		const ML_KEM_768_SECRET_KEY_LENGTH: usize = 2400;
68		const ML_KEM_768_CIPHERTEXT_LENGTH: usize = 1088;
69
70		const ML_KEM_1024_SHARED_SECRET_LENGTH: usize = 32;
71		const ML_KEM_1024_PUBLIC_KEY_LENGTH: usize = 1568;
72		const ML_KEM_1024_SECRET_KEY_LENGTH: usize = 3168;
73		const ML_KEM_1024_CIPHERTEXT_LENGTH: usize = 1568;
74
75		/// NIST FIPS 203 ML-KEM-512 algorithm.
76		pub const ML_KEM_512: Algorithm<AlgorithmId> = Algorithm {
77		id: AlgorithmId::MlKem512,
78		decapsulate_key_size: ML_KEM_512_SECRET_KEY_LENGTH,
79		encapsulate_key_size: ML_KEM_512_PUBLIC_KEY_LENGTH,
80		ciphertext_size: ML_KEM_512_CIPHERTEXT_LENGTH,
81		shared_secret_size: ML_KEM_512_SHARED_SECRET_LENGTH,
82		};
83
84		/// NIST FIPS 203 ML-KEM-768 algorithm.
85		pub const ML_KEM_768: Algorithm<AlgorithmId> = Algorithm {
86		id: AlgorithmId::MlKem768,
87		decapsulate_key_size: ML_KEM_768_SECRET_KEY_LENGTH,
88		encapsulate_key_size: ML_KEM_768_PUBLIC_KEY_LENGTH,
89		ciphertext_size: ML_KEM_768_CIPHERTEXT_LENGTH,
90		shared_secret_size: ML_KEM_768_SHARED_SECRET_LENGTH,
91		};
92
93		/// NIST FIPS 203 ML-KEM-1024 algorithm.
94		pub const ML_KEM_1024: Algorithm<AlgorithmId> = Algorithm {
95		id: AlgorithmId::MlKem1024,
96		decapsulate_key_size: ML_KEM_1024_SECRET_KEY_LENGTH,
97		encapsulate_key_size: ML_KEM_1024_PUBLIC_KEY_LENGTH,
98		ciphertext_size: ML_KEM_1024_CIPHERTEXT_LENGTH,
99		shared_secret_size: ML_KEM_1024_SHARED_SECRET_LENGTH,
100		};
101
102		use crate::aws_lc::{NID_MLKEM1024, NID_MLKEM512, NID_MLKEM768};
103
104		/// An identifier for a KEM algorithm.
105		pub trait AlgorithmIdentifier:
106		Copy + Clone + Debug + PartialEq + crate::sealed::Sealed + 'static
107		{
108		/// Returns the algorithm's associated AWS-LC nid.
109		fn nid(self) -> i32;
110		}
111
112		/// A KEM algorithm
113		#[derive(PartialEq)]
114		pub struct Algorithm<Id = AlgorithmId>
115		where
116		Id: AlgorithmIdentifier,
117		{
118		pub(crate) id: Id,
119		pub(crate) decapsulate_key_size: usize,
120		pub(crate) encapsulate_key_size: usize,
121		pub(crate) ciphertext_size: usize,
122		pub(crate) shared_secret_size: usize,
123		}
124
125		impl<Id> Algorithm<Id>
126		where
127		Id: AlgorithmIdentifier,
128		{
129		/// Returns the identifier for this algorithm.
130		#[must_use]
131	0	pub fn id(&self) -> Id {
132	0	self.id
133	0	}
134
135		#[inline]
136		#[allow(dead_code)]
137	0	pub(crate) fn decapsulate_key_size(&self) -> usize {
138	0	self.decapsulate_key_size
139	0	}
140
141		#[inline]
142	0	pub(crate) fn encapsulate_key_size(&self) -> usize {
143	0	self.encapsulate_key_size
144	0	}
145
146		#[inline]
147	0	pub(crate) fn ciphertext_size(&self) -> usize {
148	0	self.ciphertext_size
149	0	}
150
151		#[inline]
152	0	pub(crate) fn shared_secret_size(&self) -> usize {
153	0	self.shared_secret_size
154	0	}
155		}
156
157		impl<Id> Debug for Algorithm<Id>
158		where
159		Id: AlgorithmIdentifier,
160		{
161	0	fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
162	0	Debug::fmt(&self.id, f)
163	0	}
164		}
165
166		/// A serializable decapulsation key usable with KEMs. This can be randomly generated with `DecapsulationKey::generate`.
167		pub struct DecapsulationKey<Id = AlgorithmId>
168		where
169		Id: AlgorithmIdentifier,
170		{
171		algorithm: &'static Algorithm<Id>,
172		evp_pkey: LcPtr<EVP_PKEY>,
173		}
174
175		/// Identifier for a KEM algorithm.
176		#[non_exhaustive]
177		#[derive(Clone, Copy, Debug, PartialEq)]
178		pub enum AlgorithmId {
179		/// NIST FIPS 203 ML-KEM-512 algorithm.
180		MlKem512,
181
182		/// NIST FIPS 203 ML-KEM-768 algorithm.
183		MlKem768,
184
185		/// NIST FIPS 203 ML-KEM-1024 algorithm.
186		MlKem1024,
187		}
188
189		impl AlgorithmIdentifier for AlgorithmId {
190	0	fn nid(self) -> i32 {
191	0	match self {
192	0	AlgorithmId::MlKem512 => NID_MLKEM512,
193	0	AlgorithmId::MlKem768 => NID_MLKEM768,
194	0	AlgorithmId::MlKem1024 => NID_MLKEM1024,
195		}
196	0	}
197		}
198
199		impl crate::sealed::Sealed for AlgorithmId {}
200
201		impl<Id> DecapsulationKey<Id>
202		where
203		Id: AlgorithmIdentifier,
204		{
205		/// Generate a new KEM decapsulation key for the given algorithm.
206		///
207		/// # Errors
208		/// `error::Unspecified` when operation fails due to internal error.
209	0	pub fn generate(alg: &'static Algorithm<Id>) -> Result<Self, Unspecified> {
210	0	let kyber_key = kem_key_generate(alg.id.nid())?;
211	0	Ok(DecapsulationKey {
212	0	algorithm: alg,
213	0	evp_pkey: kyber_key,
214	0	})
215	0	}
216
217		/// Return the algorithm associated with the given KEM decapsulation key.
218		#[must_use]
219	0	pub fn algorithm(&self) -> &'static Algorithm<Id> {
220	0	self.algorithm
221	0	}
222
223		/// Computes the KEM encapsulation key from the KEM decapsulation key.
224		///
225		/// # Errors
226		/// `error::Unspecified` when operation fails due to internal error.
227		#[allow(clippy::missing_panics_doc)]
228	0	pub fn encapsulation_key(&self) -> Result<EncapsulationKey<Id>, Unspecified> {
229	0	let evp_pkey = self.evp_pkey.clone();
230
231	0	Ok(EncapsulationKey {
232	0	algorithm: self.algorithm,
233	0	evp_pkey,
234	0	})
235	0	}
236
237		/// Performs the decapsulate operation using this KEM decapsulation key on the given ciphertext.
238		///
239		/// `ciphertext` is the ciphertext generated by the encapsulate operation using the KEM encapsulation key
240		/// associated with this KEM decapsulation key.
241		///
242		/// # Errors
243		/// `Unspecified` when operation fails due to internal error.
244		#[allow(clippy::needless_pass_by_value)]
245	0	pub fn decapsulate(&self, ciphertext: Ciphertext<'_>) -> Result<SharedSecret, Unspecified> {
246	0	let mut shared_secret_len = self.algorithm.shared_secret_size();
247	0	let mut shared_secret: Vec<u8> = vec![0u8; shared_secret_len];
248
249	0	let mut ctx = self.evp_pkey.create_EVP_PKEY_CTX()?;
250
251	0	let ciphertext = ciphertext.as_ref();
252
253	0	if 1 != unsafe {
254	0	EVP_PKEY_decapsulate(
255	0	*ctx.as_mut(),
256	0	shared_secret.as_mut_ptr(),
257	0	&mut shared_secret_len,
258	0	// AWS-LC incorrectly has this as an unqualified `uint8_t *`, it should be qualified with const
259	0	ciphertext.as_ptr().cast_mut(),
260	0	ciphertext.len(),
261	0	)
262	0	} {
263	0	return Err(Unspecified);
264	0	}
265
266		// This is currently pedantic but done for safety in-case the shared_secret buffer
267		// size changes in the future. `EVP_PKEY_decapsulate` updates `shared_secret_len` with
268		// the length of the shared secret in the event the buffer provided was larger then the secret.
269		// This truncates the buffer to the proper length to match the shared secret written.
270	0	debug_assert_eq!(shared_secret_len, shared_secret.len());
271	0	shared_secret.truncate(shared_secret_len);
272
273	0	Ok(SharedSecret(shared_secret.into_boxed_slice()))
274	0	}
275		}
276
277		unsafe impl<Id> Send for DecapsulationKey<Id> where Id: AlgorithmIdentifier {}
278
279		unsafe impl<Id> Sync for DecapsulationKey<Id> where Id: AlgorithmIdentifier {}
280
281		impl<Id> Debug for DecapsulationKey<Id>
282		where
283		Id: AlgorithmIdentifier,
284		{
285	0	fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
286	0	f.debug_struct("DecapsulationKey")
287	0	.field("algorithm", &self.algorithm)
288	0	.finish_non_exhaustive()
289	0	}
290		}
291
292		generated_encodings!((EncapsulationKeyBytes, EncapsulationKeyBytesType));
293
294		/// A serializable encapsulation key usable with KEM algorithms. Constructed
295		/// from either a `DecapsulationKey` or raw bytes.
296		pub struct EncapsulationKey<Id = AlgorithmId>
297		where
298		Id: AlgorithmIdentifier,
299		{
300		algorithm: &'static Algorithm<Id>,
301		evp_pkey: LcPtr<EVP_PKEY>,
302		}
303
304		impl<Id> EncapsulationKey<Id>
305		where
306		Id: AlgorithmIdentifier,
307		{
308		/// Return the algorithm associated with the given KEM encapsulation key.
309		#[must_use]
310	0	pub fn algorithm(&self) -> &'static Algorithm<Id> {
311	0	self.algorithm
312	0	}
313
314		/// Performs the encapsulate operation using this KEM encapsulation key, generating a ciphertext
315		/// and associated shared secret.
316		///
317		/// # Errors
318		/// `error::Unspecified` when operation fails due to internal error.
319	0	pub fn encapsulate(&self) -> Result<(Ciphertext<'static>, SharedSecret), Unspecified> {
320	0	let mut ciphertext_len = self.algorithm.ciphertext_size();
321	0	let mut shared_secret_len = self.algorithm.shared_secret_size();
322	0	let mut ciphertext: Vec<u8> = vec![0u8; ciphertext_len];
323	0	let mut shared_secret: Vec<u8> = vec![0u8; shared_secret_len];
324
325	0	let mut ctx = self.evp_pkey.create_EVP_PKEY_CTX()?;
326
327	0	if 1 != unsafe {
328	0	EVP_PKEY_encapsulate(
329	0	*ctx.as_mut(),
330	0	ciphertext.as_mut_ptr(),
331	0	&mut ciphertext_len,
332	0	shared_secret.as_mut_ptr(),
333	0	&mut shared_secret_len,
334	0	)
335	0	} {
336	0	return Err(Unspecified);
337	0	}
338
339		// The following two steps are currently pedantic but done for safety in-case the buffer allocation
340		// sizes change in the future. `EVP_PKEY_encapsulate` updates `ciphertext_len` and `shared_secret_len` with
341		// the length of the ciphertext and shared secret respectivly in the event the buffer provided for each was
342		// larger then the actual values. Thus these two steps truncate the buffers to the proper length to match the
343		// value lengths written.
344	0	debug_assert_eq!(ciphertext_len, ciphertext.len());
345	0	ciphertext.truncate(ciphertext_len);
346	0	debug_assert_eq!(shared_secret_len, shared_secret.len());
347	0	shared_secret.truncate(shared_secret_len);
348
349	0	Ok((
350	0	Ciphertext::new(ciphertext),
351	0	SharedSecret::new(shared_secret.into_boxed_slice()),
352	0	))
353	0	}
354
355		/// Returns the `EnscapsulationKey` bytes.
356		///
357		/// # Errors
358		/// * `Unspecified`: Any failure to retrieve the `EnscapsulationKey` bytes.
359	0	pub fn key_bytes(&self) -> Result<EncapsulationKeyBytes<'static>, Unspecified> {
360	0	let mut encapsulate_bytes = vec![0u8; self.algorithm.encapsulate_key_size()];
361	0	let encapsulate_key_size = self
362	0	.evp_pkey
363	0	.as_const()
364	0	.marshal_raw_public_to_buffer(&mut encapsulate_bytes)?;
365
366	0	debug_assert_eq!(encapsulate_key_size, encapsulate_bytes.len());
367	0	encapsulate_bytes.truncate(encapsulate_key_size);
368
369	0	Ok(EncapsulationKeyBytes::new(encapsulate_bytes))
370	0	}
371
372		/// Creates a new KEM encapsulation key from raw bytes. This method MUST NOT be used to generate
373		/// a new encapsulation key, rather it MUST be used to construct `EncapsulationKey` previously serialized
374		/// to raw bytes.
375		///
376		/// `alg` is the [`Algorithm`] to be associated with the generated `EncapsulationKey`.
377		///
378		/// `bytes` is a slice of raw bytes representing a `EncapsulationKey`.
379		///
380		/// # Errors
381		/// `error::KeyRejected` when operation fails during key creation.
382	0	pub fn new(alg: &'static Algorithm<Id>, bytes: &[u8]) -> Result<Self, KeyRejected> {
383	0	match bytes.len().cmp(&alg.encapsulate_key_size()) {
384	0	Ordering::Less => Err(KeyRejected::too_small()),
385	0	Ordering::Greater => Err(KeyRejected::too_large()),
386	0	Ordering::Equal => Ok(()),
387	0	}?;
388	0	let pubkey = LcPtr::new(unsafe {
389	0	EVP_PKEY_kem_new_raw_public_key(alg.id.nid(), bytes.as_ptr(), bytes.len())
390	0	})?;
391	0	Ok(EncapsulationKey {
392	0	algorithm: alg,
393	0	evp_pkey: pubkey,
394	0	})
395	0	}
396		}
397
398		unsafe impl<Id> Send for EncapsulationKey<Id> where Id: AlgorithmIdentifier {}
399
400		unsafe impl<Id> Sync for EncapsulationKey<Id> where Id: AlgorithmIdentifier {}
401
402		impl<Id> Debug for EncapsulationKey<Id>
403		where
404		Id: AlgorithmIdentifier,
405		{
406	0	fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
407	0	f.debug_struct("EncapsulationKey")
408	0	.field("algorithm", &self.algorithm)
409	0	.finish_non_exhaustive()
410	0	}
411		}
412
413		/// A set of encrypted bytes produced by [`EncapsulationKey::encapsulate`],
414		/// and used as an input to [`DecapsulationKey::decapsulate`].
415		pub struct Ciphertext<'a>(Cow<'a, [u8]>);
416
417		impl<'a> Ciphertext<'a> {
418	0	fn new(value: Vec<u8>) -> Ciphertext<'a> {
419	0	Self(Cow::Owned(value))
420	0	}
421		}
422
423		impl Drop for Ciphertext<'_> {
424	0	fn drop(&mut self) {
425	0	if let Cow::Owned(ref mut v) = self.0 {
426	0	v.zeroize();
427	0	}
428	0	}
429		}
430
431		impl AsRef<[u8]> for Ciphertext<'_> {
432	0	fn as_ref(&self) -> &[u8] {
433	0	match self.0 {
434	0	Cow::Borrowed(v) => v,
435	0	Cow::Owned(ref v) => v.as_ref(),
436		}
437	0	}
438		}
439
440		impl<'a> From<&'a [u8]> for Ciphertext<'a> {
441	0	fn from(value: &'a [u8]) -> Self {
442	0	Self(Cow::Borrowed(value))
443	0	}
444		}
445
446		/// The cryptographic shared secret output from the KEM encapsulate / decapsulate process.
447		pub struct SharedSecret(Box<[u8]>);
448
449		impl SharedSecret {
450	0	fn new(value: Box<[u8]>) -> Self {
451	0	Self(value)
452	0	}
453		}
454
455		impl Drop for SharedSecret {
456	0	fn drop(&mut self) {
457	0	self.0.zeroize();
458	0	}
459		}
460
461		impl AsRef<[u8]> for SharedSecret {
462	0	fn as_ref(&self) -> &[u8] {
463	0	self.0.as_ref()
464	0	}
465		}
466
467		// Returns an LcPtr to an EVP_PKEY
468		#[inline]
469	0	fn kem_key_generate(nid: i32) -> Result<LcPtr<EVP_PKEY>, Unspecified> {
470	0	let params_fn = \|ctx\| {
471	0	if 1 == unsafe { EVP_PKEY_CTX_kem_set_params(ctx, nid) } {
472	0	Ok(())
473		} else {
474	0	Err(())
475		}
476	0	};
477
478	0	LcPtr::<EVP_PKEY>::generate(EVP_PKEY_KEM, Some(params_fn))
479	0	}
480
481		#[cfg(test)]
482		mod tests {
483		use super::{Ciphertext, DecapsulationKey, EncapsulationKey, SharedSecret};
484		use crate::error::KeyRejected;
485
486		use crate::kem::{ML_KEM_1024, ML_KEM_512, ML_KEM_768};
487
488		#[test]
489		fn ciphertext() {
490		let ciphertext_bytes = vec![42u8; 4];
491		let ciphertext = Ciphertext::from(ciphertext_bytes.as_ref());
492		assert_eq!(ciphertext.as_ref(), &[42, 42, 42, 42]);
493		drop(ciphertext);
494
495		let ciphertext_bytes = vec![42u8; 4];
496		let ciphertext = Ciphertext::<'static>::new(ciphertext_bytes);
497		assert_eq!(ciphertext.as_ref(), &[42, 42, 42, 42]);
498		}
499
500		#[test]
501		fn shared_secret() {
502		let secret_bytes = vec![42u8; 4];
503		let shared_secret = SharedSecret::new(secret_bytes.into_boxed_slice());
504		assert_eq!(shared_secret.as_ref(), &[42, 42, 42, 42]);
505		}
506
507		#[test]
508		fn test_kem_serialize() {
509		for algorithm in [&ML_KEM_512, &ML_KEM_768, &ML_KEM_1024] {
510		let priv_key = DecapsulationKey::generate(algorithm).unwrap();
511		assert_eq!(priv_key.algorithm(), algorithm);
512
513		let pub_key = priv_key.encapsulation_key().unwrap();
514		let pubkey_raw_bytes = pub_key.key_bytes().unwrap();
515		let pub_key_from_bytes =
516		EncapsulationKey::new(algorithm, pubkey_raw_bytes.as_ref()).unwrap();
517
518		assert_eq!(
519		pub_key.key_bytes().unwrap().as_ref(),
520		pub_key_from_bytes.key_bytes().unwrap().as_ref()
521		);
522		assert_eq!(pub_key.algorithm(), pub_key_from_bytes.algorithm());
523		}
524		}
525
526		#[test]
527		fn test_kem_wrong_sizes() {
528		for algorithm in [&ML_KEM_512, &ML_KEM_768, &ML_KEM_1024] {
529		let too_long_bytes = vec![0u8; algorithm.encapsulate_key_size() + 1];
530		let long_pub_key_from_bytes = EncapsulationKey::new(algorithm, &too_long_bytes);
531		assert_eq!(
532		long_pub_key_from_bytes.err(),
533		Some(KeyRejected::too_large())
534		);
535
536		let too_short_bytes = vec![0u8; algorithm.encapsulate_key_size() - 1];
537		let short_pub_key_from_bytes = EncapsulationKey::new(algorithm, &too_short_bytes);
538		assert_eq!(
539		short_pub_key_from_bytes.err(),
540		Some(KeyRejected::too_small())
541		);
542		}
543		}
544
545		#[test]
546		fn test_kem_e2e() {
547		for algorithm in [&ML_KEM_512, &ML_KEM_768, &ML_KEM_1024] {
548		let priv_key = DecapsulationKey::generate(algorithm).unwrap();
549		assert_eq!(priv_key.algorithm(), algorithm);
550
551		let pub_key = priv_key.encapsulation_key().unwrap();
552
553		let (alice_ciphertext, alice_secret) =
554		pub_key.encapsulate().expect("encapsulate successful");
555
556		let bob_secret = priv_key
557		.decapsulate(alice_ciphertext)
558		.expect("decapsulate successful");
559
560		assert_eq!(alice_secret.as_ref(), bob_secret.as_ref());
561		}
562		}
563
564		#[test]
565		fn test_serialized_kem_e2e() {
566		for algorithm in [&ML_KEM_512, &ML_KEM_768, &ML_KEM_1024] {
567		let priv_key = DecapsulationKey::generate(algorithm).unwrap();
568		assert_eq!(priv_key.algorithm(), algorithm);
569
570		let pub_key = priv_key.encapsulation_key().unwrap();
571
572		// Generate public key bytes to send to bob
573		let pub_key_bytes = pub_key.key_bytes().unwrap();
574
575		// Test that priv_key's EVP_PKEY isn't entirely freed since we remove this pub_key's reference.
576		drop(pub_key);
577
578		let retrieved_pub_key =
579		EncapsulationKey::new(algorithm, pub_key_bytes.as_ref()).unwrap();
580		let (ciphertext, bob_secret) = retrieved_pub_key
581		.encapsulate()
582		.expect("encapsulate successful");
583
584		let alice_secret = priv_key
585		.decapsulate(ciphertext)
586		.expect("decapsulate successful");
587
588		assert_eq!(alice_secret.as_ref(), bob_secret.as_ref());
589		}
590		}
591
592		#[test]
593		fn test_debug_fmt() {
594		let private = DecapsulationKey::generate(&ML_KEM_512).expect("successful generation");
595		assert_eq!(
596		format!("{private:?}"),
597		"DecapsulationKey { algorithm: MlKem512, .. }"
598		);
599		assert_eq!(
600		format!(
601		"{:?}",
602		private.encapsulation_key().expect("public key retrievable")
603		),
604		"EncapsulationKey { algorithm: MlKem512, .. }"
605		);
606		}
607		}