s390-tools/rust/pv/src/openssl_extensions/ml_kem.rs

// SPDX-License-Identifier: MIT
//
// Copyright IBM Corp.

#![cfg_attr(not(test), expect(unused))]

use std::ffi::CStr;
use std::ptr::NonNull;

use foreign_types::ForeignType;
use openssl::error::ErrorStack;
use openssl::pkey::{KeyType, PKey, Private};
use openssl::pkey_ctx::PkeyCtx;

// automatically generated by rust-bindgen 0.69.5

mod ffi {
    #[repr(C)]
    #[derive(Debug, Copy, Clone)]
    pub struct evp_pkey_ctx_st {
        _unused: [u8; 0],
    }
    pub type EvpPkeyCtx = evp_pkey_ctx_st;
    #[repr(C)]
    #[derive(Debug, Copy, Clone)]
    pub struct ossl_param_st {
        pub key: *const ::std::os::raw::c_char,
        pub data_type: ::std::os::raw::c_uint,
        pub data: *mut ::std::os::raw::c_void,
        pub data_size: usize,
        pub return_size: usize,
    }
    pub type OsslParam = ossl_param_st;

    extern "C" {
        pub fn EVP_PKEY_CTX_new_from_name(
            libctx: *mut ::std::os::raw::c_void,
            name: *const ::std::os::raw::c_char,
            propquery: *const ::std::os::raw::c_char,
        ) -> *mut EvpPkeyCtx;
    }
    extern "C" {
        pub fn EVP_PKEY_keygen_init(ctx: *mut EvpPkeyCtx) -> ::std::os::raw::c_int;
    }
    extern "C" {
        pub fn EVP_PKEY_generate(
            ctx: *mut EvpPkeyCtx,
            ppkey: *mut *mut ::std::os::raw::c_void,
        ) -> ::std::os::raw::c_int;
    }
    extern "C" {
        pub fn EVP_PKEY_CTX_free(ctx: *mut EvpPkeyCtx);
    }
    extern "C" {
        pub fn EVP_PKEY_encapsulate_init(
            ctx: *mut EvpPkeyCtx,
            params: *const OsslParam,
        ) -> ::std::os::raw::c_int;
    }
    extern "C" {
        pub fn EVP_PKEY_encapsulate(
            ctx: *mut EvpPkeyCtx,
            wrappedkey: *mut ::std::os::raw::c_uchar,
            wrappedkeylen: *mut usize,
            genkey: *mut ::std::os::raw::c_uchar,
            genkeylen: *mut usize,
        ) -> ::std::os::raw::c_int;
    }
    extern "C" {
        pub fn EVP_PKEY_decapsulate_init(
            ctx: *mut EvpPkeyCtx,
            params: *const OsslParam,
        ) -> ::std::os::raw::c_int;
    }
    extern "C" {
        pub fn EVP_PKEY_decapsulate(
            ctx: *mut EvpPkeyCtx,
            unwrapped: *mut ::std::os::raw::c_uchar,
            unwrappedlen: *mut usize,
            wrapped: *const ::std::os::raw::c_uchar,
            wrappedlen: usize,
        ) -> ::std::os::raw::c_int;
    }
}

const ML_KEM_512_NAME: &CStr = c"ML-KEM-512";
const ML_KEM_768_NAME: &CStr = c"ML-KEM-768";
const ML_KEM_1024_NAME: &CStr = c"ML-KEM-1024";

fn ml_kem_name(key_type: KeyType) -> Result<&'static CStr, ErrorStack> {
    if key_type == KeyType::ML_KEM_512 {
        Ok(ML_KEM_512_NAME)
    } else if key_type == KeyType::ML_KEM_768 {
        Ok(ML_KEM_768_NAME)
    } else if key_type == KeyType::ML_KEM_1024 {
        Ok(ML_KEM_1024_NAME)
    } else {
        Err(ErrorStack::get())
    }
}

/// Encapsulate a key
pub trait PkeyEncapsulateContext {
    /// Initialize the encapsulation operation.
    ///
    /// # Errors
    /// Returns an error if the OpenSSL operation fails.
    fn encapsulate_init(&mut self) -> Result<(), ErrorStack>;

    /// Perform the encapsulation operation.
    ///
    /// # Parameters
    /// - `wrappedkey`: Optional buffer to receive the wrapped key.
    /// - `genkey`: Optional buffer to receive the generated key.
    ///
    /// # Returns
    /// A tuple of `(wrappedkey_len, genkey_len)` on success.
    ///
    /// # Errors
    /// Returns an error if the OpenSSL operation fails.
    fn encapsulate(
        &mut self,
        wrappedkey: Option<&mut [u8]>,
        genkey: Option<&mut [u8]>,
    ) -> Result<(usize, usize), ErrorStack>;

    /// Convenience method to encapsulate into vectors.
    ///
    /// # Parameters
    /// - `wrappedkey`: Buffer to receive the wrapped key (ciphertext).
    /// - `genkey`: Buffer to receive the generated key (shared secret).
    ///
    /// # Returns
    /// A tuple of `(wrappedkey_len, genkey_len)` on success.
    ///
    /// # Errors
    /// Returns an error if the OpenSSL operation fails.
    fn encapsulate_to_vec(
        &mut self,
        wrappedkey: &mut Vec<u8>,
        genkey: &mut Vec<u8>,
    ) -> Result<(usize, usize), ErrorStack> {
        let wrappedkey_base = wrappedkey.len();
        let genkey_base = genkey.len();

        // Query the required output buffer sizes.
        let (wrappedkey_len, genkey_len) = self.encapsulate(None, None)?;

        wrappedkey.resize(wrappedkey_base + wrappedkey_len, 0);
        genkey.resize(genkey_base + genkey_len, 0);

        let (wrappedkey_len, genkey_len) = self.encapsulate(
            Some(&mut wrappedkey[wrappedkey_base..]),
            Some(&mut genkey[genkey_base..]),
        )?;

        wrappedkey.truncate(wrappedkey_base + wrappedkey_len);
        genkey.truncate(genkey_base + genkey_len);

        Ok((wrappedkey_len, genkey_len))
    }
}

impl<T> PkeyEncapsulateContext for PkeyCtx<T> {
    #[inline]
    fn encapsulate_init(&mut self) -> Result<(), ErrorStack> {
        // SAFETY: self.as_ptr() returns a valid EVP_PKEY_CTX pointer. Parameter-based
        // configuration is intentionally unsupported here, so a null params pointer is passed.
        let ret = unsafe {
            ffi::EVP_PKEY_encapsulate_init(self.as_ptr() as *mut ffi::EvpPkeyCtx, std::ptr::null())
        };

        if ret == 1 {
            Ok(())
        } else {
            Err(ErrorStack::get())
        }
    }

    fn encapsulate(
        &mut self,
        wrappedkey: Option<&mut [u8]>,
        genkey: Option<&mut [u8]>,
    ) -> Result<(usize, usize), ErrorStack> {
        let mut wrappedkey_len = wrappedkey.as_ref().map_or(0, |buf| buf.len());
        let mut genkey_len = genkey.as_ref().map_or(0, |buf| buf.len());

        let wrappedkey_ptr = wrappedkey
            .map(|buf| buf.as_mut_ptr())
            .unwrap_or(std::ptr::null_mut());
        let genkey_ptr = genkey
            .map(|buf| buf.as_mut_ptr())
            .unwrap_or(std::ptr::null_mut());

        // SAFETY: All pointers are either valid mutable buffers or null.
        let ret = unsafe {
            ffi::EVP_PKEY_encapsulate(
                self.as_ptr() as *mut ffi::EvpPkeyCtx,
                wrappedkey_ptr,
                &mut wrappedkey_len,
                genkey_ptr,
                &mut genkey_len,
            )
        };

        if ret == 1 {
            Ok((wrappedkey_len, genkey_len))
        } else {
            Err(ErrorStack::get())
        }
    }
}

/// Decapsulate a key
pub trait PkeyDecapsulateContext {
    /// Initialize the decapsulation operation.
    ///
    /// # Errors
    /// Returns an error if the OpenSSL operation fails.
    fn decapsulate_init(&mut self) -> Result<(), ErrorStack>;

    /// Perform the decapsulation operation.
    ///
    /// # Parameters
    /// - `wrapped`: The wrapped key (ciphertext) to decapsulate.
    /// - `unwrapped`: Optional buffer to receive the unwrapped shared secret.
    ///
    /// # Returns
    /// The length of the unwrapped shared secret on success.
    ///
    /// # Errors
    /// Returns an error if the OpenSSL operation fails.
    fn decapsulate(
        &mut self,
        wrapped: &[u8],
        unwrapped: Option<&mut [u8]>,
    ) -> Result<usize, ErrorStack>;

    /// Convenience method to decapsulate into a vector.
    ///
    /// # Parameters
    /// - `wrapped`: The wrapped key (ciphertext) to decapsulate.
    /// - `unwrapped`: Buffer to receive the unwrapped shared secret.
    ///
    /// # Returns
    /// The length of the unwrapped shared secret on success.
    ///
    /// # Errors
    /// Returns an error if the OpenSSL operation fails.
    fn decapsulate_to_vec(
        &mut self,
        wrapped: &[u8],
        unwrapped: &mut Vec<u8>,
    ) -> Result<usize, ErrorStack> {
        let unwrapped_base = unwrapped.len();

        // Query the required output buffer size.
        let unwrapped_len = self.decapsulate(wrapped, None)?;

        unwrapped.resize(unwrapped_base + unwrapped_len, 0);

        let unwrapped_len = self.decapsulate(wrapped, Some(&mut unwrapped[unwrapped_base..]))?;

        unwrapped.truncate(unwrapped_base + unwrapped_len);

        Ok(unwrapped_len)
    }
}

impl<T> PkeyDecapsulateContext for PkeyCtx<T> {
    #[inline]
    fn decapsulate_init(&mut self) -> Result<(), ErrorStack> {
        // SAFETY: self.as_ptr() returns a valid EVP_PKEY_CTX pointer. Parameter-based
        // configuration is intentionally unsupported here, so a null params pointer is passed.
        let ret = unsafe {
            ffi::EVP_PKEY_decapsulate_init(self.as_ptr() as *mut ffi::EvpPkeyCtx, std::ptr::null())
        };

        if ret == 1 {
            Ok(())
        } else {
            Err(ErrorStack::get())
        }
    }

    fn decapsulate(
        &mut self,
        wrapped: &[u8],
        unwrapped: Option<&mut [u8]>,
    ) -> Result<usize, ErrorStack> {
        let mut unwrapped_len = unwrapped.as_ref().map_or(0, |buf| buf.len());

        let unwrapped_ptr = unwrapped
            .map(|buf| buf.as_mut_ptr())
            .unwrap_or(std::ptr::null_mut());

        // SAFETY: All pointers are either valid mutable buffers or null, and wrapped is a valid
        // slice.
        let ret = unsafe {
            ffi::EVP_PKEY_decapsulate(
                self.as_ptr() as *mut ffi::EvpPkeyCtx,
                unwrapped_ptr,
                &mut unwrapped_len,
                wrapped.as_ptr(),
                wrapped.len(),
            )
        };

        if ret == 1 {
            Ok(unwrapped_len)
        } else {
            Err(ErrorStack::get())
        }
    }
}

/// ML-KEM key generation context.
///
/// The inner pointer is guaranteed to be non-null for the lifetime of the context.
pub struct MlKemKeyCtx {
    ctx: NonNull<ffi::EvpPkeyCtx>,
}

impl MlKemKeyCtx {
    /// Create a new ML-KEM key generation context.
    ///
    /// # Parameters
    /// - `key_type`: The ML-KEM key type (`KeyType::ML_KEM_512`, `KeyType::ML_KEM_768`, or
    ///   `KeyType::ML_KEM_1024`).
    ///
    /// # Errors
    /// Returns an error if the context cannot be created or the key type is unsupported.
    pub fn new(key_type: KeyType) -> Result<Self, ErrorStack> {
        let name = ml_kem_name(key_type)?;

        // SAFETY: name is a valid null-terminated C string.
        let ctx = unsafe {
            ffi::EVP_PKEY_CTX_new_from_name(std::ptr::null_mut(), name.as_ptr(), std::ptr::null())
        };
        let Some(ctx) = NonNull::new(ctx) else {
            return Err(ErrorStack::get());
        };

        // SAFETY: ctx is a valid non-null pointer.
        let ret = unsafe { ffi::EVP_PKEY_keygen_init(ctx.as_ptr()) };
        if ret != 1 {
            // SAFETY: ctx is a valid non-null pointer.
            unsafe { ffi::EVP_PKEY_CTX_free(ctx.as_ptr()) };
            return Err(ErrorStack::get());
        }

        Ok(Self { ctx })
    }

    /// Generate an ML-KEM keypair.
    ///
    /// # Returns
    /// A `PKey` containing both the private and public key.
    ///
    /// # Errors
    /// Returns an error if key generation fails.
    pub fn generate(&mut self) -> Result<PKey<Private>, ErrorStack> {
        let mut pkey: *mut std::os::raw::c_void = std::ptr::null_mut();

        // SAFETY: self.ctx is valid, and pkey is a valid mutable pointer.
        let ret = unsafe { ffi::EVP_PKEY_generate(self.ctx.as_ptr(), &mut pkey) };

        if ret != 1 || pkey.is_null() {
            return Err(ErrorStack::get());
        }

        // SAFETY: pkey is a valid EVP_PKEY pointer created by OpenSSL.
        unsafe { Ok(PKey::from_ptr(pkey as *mut openssl_sys::EVP_PKEY)) }
    }
}

impl Drop for MlKemKeyCtx {
    fn drop(&mut self) {
        // SAFETY: self.ctx is guaranteed to be non-null and owned by this context.
        unsafe {
            ffi::EVP_PKEY_CTX_free(self.ctx.as_ptr());
        }
    }
}

/// Generate an ML-KEM keypair.
///
/// # Parameters
/// - `key_type`: The ML-KEM key type (`KeyType::ML_KEM_512`, `KeyType::ML_KEM_768`, or
///   `KeyType::ML_KEM_1024`).
///
/// # Returns
/// A `PKey` containing both the private and public key.
///
/// # Errors
/// Returns an error if key generation fails.
pub fn generate_ml_kem(key_type: KeyType) -> Result<PKey<Private>, ErrorStack> {
    let mut ctx = MlKemKeyCtx::new(key_type)?;
    ctx.generate()
}

#[cfg(test)]
mod tests {
    use super::*;

    fn roundtrip_ml_kem(key_type: KeyType, expected_ciphertext_len: usize) {
        let keypair = generate_ml_kem(key_type).expect("Failed to generate ML-KEM keypair");

        // Extract public key for encapsulation
        let public_key_der = keypair
            .public_key_to_der()
            .expect("Failed to export public key");
        let public_key =
            PKey::public_key_from_der(&public_key_der).expect("Failed to import public key");

        // Encapsulate: generate ciphertext and shared secret using public key
        let mut ctx_enc =
            PkeyCtx::new(&public_key).expect("Failed to create encapsulation context");
        ctx_enc
            .encapsulate_init()
            .expect("Failed to initialize encapsulation");

        let mut ciphertext = Vec::new();
        let mut shared_secret_enc = Vec::new();
        ctx_enc
            .encapsulate_to_vec(&mut ciphertext, &mut shared_secret_enc)
            .expect("Failed to encapsulate");

        assert!(!ciphertext.is_empty(), "Ciphertext should not be empty");
        assert!(
            !shared_secret_enc.is_empty(),
            "Shared secret should not be empty"
        );

        // Decapsulate: recover shared secret from ciphertext using private key
        let mut ctx_dec = PkeyCtx::new(&keypair).expect("Failed to create decapsulation context");
        ctx_dec
            .decapsulate_init()
            .expect("Failed to initialize decapsulation");

        let mut shared_secret_dec = Vec::new();
        ctx_dec
            .decapsulate_to_vec(&ciphertext, &mut shared_secret_dec)
            .expect("Failed to decapsulate");

        // Verify that the shared secrets match
        assert_eq!(
            shared_secret_enc, shared_secret_dec,
            "Shared secrets from encapsulation and decapsulation should match"
        );
        assert_eq!(ciphertext.len(), expected_ciphertext_len);
        assert_eq!(shared_secret_enc.len(), 32);
    }

    #[test]
    fn test_encapsulate_decapsulate_roundtrip_ml_kem_512() {
        roundtrip_ml_kem(KeyType::ML_KEM_512, 768);
    }

    #[test]
    fn test_encapsulate_decapsulate_roundtrip_ml_kem_768() {
        roundtrip_ml_kem(KeyType::ML_KEM_768, 1088);
    }

    #[test]
    fn test_encapsulate_decapsulate_roundtrip_ml_kem_1024() {
        roundtrip_ml_kem(KeyType::ML_KEM_1024, 1568);
    }

    #[test]
    fn test_ml_kem_key_ctx_new_accepts_supported_key_types() {
        MlKemKeyCtx::new(KeyType::ML_KEM_512).expect("ML-KEM-512 context creation must succeed");
        MlKemKeyCtx::new(KeyType::ML_KEM_768).expect("ML-KEM-768 context creation must succeed");
        MlKemKeyCtx::new(KeyType::ML_KEM_1024).expect("ML-KEM-1024 context creation must succeed");
    }

    #[test]
    fn test_ml_kem_key_ctx_new_rejects_unsupported_key_type() {
        assert!(
            MlKemKeyCtx::new(KeyType::RSA).is_err(),
            "Unsupported key type must fail"
        );
    }

    #[test]
    fn test_ml_kem_key_ctx_generate_produces_ml_kem_keypair() {
        let mut ctx = MlKemKeyCtx::new(KeyType::ML_KEM_768)
            .expect("ML-KEM-768 context creation must succeed");
        let keypair = ctx
            .generate()
            .expect("ML-KEM-768 key generation must succeed");

        let public_key_der = keypair
            .public_key_to_der()
            .expect("Generated keypair must export a public key");
        let public_key =
            PKey::public_key_from_der(&public_key_der).expect("Generated public key must reimport");

        let mut ctx_enc =
            PkeyCtx::new(&public_key).expect("Failed to create encapsulation context");
        ctx_enc
            .encapsulate_init()
            .expect("Failed to initialize encapsulation");

        let mut ciphertext = Vec::new();
        let mut shared_secret = Vec::new();
        ctx_enc
            .encapsulate_to_vec(&mut ciphertext, &mut shared_secret)
            .expect("Generated keypair must support encapsulation");

        assert_eq!(ciphertext.len(), 1088);
        assert_eq!(shared_secret.len(), 32);
    }

    #[test]
    fn test_generate_ml_kem_rejects_unsupported_key_type() {
        generate_ml_kem(KeyType::RSA).expect_err("Unsupported key type must fail");
    }

    #[test]
    fn test_encapsulate_requires_init() {
        let keypair =
            generate_ml_kem(KeyType::ML_KEM_512).expect("Failed to generate ML-KEM-512 keypair");
        let public_key_der = keypair
            .public_key_to_der()
            .expect("Failed to export public key");
        let public_key =
            PKey::public_key_from_der(&public_key_der).expect("Failed to import public key");

        let mut ctx_enc =
            PkeyCtx::new(&public_key).expect("Failed to create encapsulation context");
        let err = ctx_enc
            .encapsulate_to_vec(&mut Vec::new(), &mut Vec::new())
            .expect_err("Encapsulation without init must fail");

        assert!(
            !err.errors().is_empty(),
            "OpenSSL should report an error when encapsulate_init was not called"
        );
    }

    #[test]
    fn test_decapsulate_requires_init() {
        let keypair =
            generate_ml_kem(KeyType::ML_KEM_512).expect("Failed to generate ML-KEM-512 keypair");
        let mut ctx_dec = PkeyCtx::new(&keypair).expect("Failed to create decapsulation context");
        let err = ctx_dec
            .decapsulate_to_vec(&[0_u8; 768], &mut Vec::new())
            .expect_err("Decapsulation without init must fail");

        assert!(
            !err.errors().is_empty(),
            "OpenSSL should report an error when decapsulate_init was not called"
        );
    }

    #[test]
    fn test_decapsulate_rejects_undersized_output_buffer() {
        let keypair =
            generate_ml_kem(KeyType::ML_KEM_768).expect("Failed to generate ML-KEM-768 keypair");

        let public_key_der = keypair
            .public_key_to_der()
            .expect("Failed to export public key");
        let public_key =
            PKey::public_key_from_der(&public_key_der).expect("Failed to import public key");

        let mut ctx_enc =
            PkeyCtx::new(&public_key).expect("Failed to create encapsulation context");
        ctx_enc
            .encapsulate_init()
            .expect("Failed to initialize encapsulation");

        let mut ciphertext = Vec::new();
        let mut shared_secret = Vec::new();
        ctx_enc
            .encapsulate_to_vec(&mut ciphertext, &mut shared_secret)
            .expect("Failed to encapsulate");

        let mut ctx_dec = PkeyCtx::new(&keypair).expect("Failed to create decapsulation context");
        ctx_dec
            .decapsulate_init()
            .expect("Failed to initialize decapsulation");

        let mut undersized = [0_u8; 31];
        let err = ctx_dec
            .decapsulate(&ciphertext, Some(&mut undersized))
            .expect_err("Undersized output buffer must fail");

        assert!(
            !err.errors().is_empty(),
            "OpenSSL should report an error for an undersized output buffer"
        );
    }
}