quicproquo/crates/quicprochat-core/src/hybrid_crypto.rs

//! Post-quantum hybrid crypto provider for OpenMLS (M7 PoC).
//!
//! Uses X25519 + ML-KEM-768 hybrid KEM for HPKE operations where openmls
//! would use DHKEM(X25519), and delegates all other operations (AEAD, hash,
//! signatures, KDF, randomness) to `openmls_rust_crypto::RustCrypto`.
//!
//! # Key format
//!
//! When the provider sees a **hybrid public key** (length `HYBRID_PUBLIC_KEY_LEN` =
//! 32 + 1184 bytes) or **hybrid private key** (length `HYBRID_PRIVATE_KEY_LEN` =
//! 32 + 2400 bytes), it uses `hybrid_kem` for HPKE. Otherwise it delegates to
//! RustCrypto (classical X25519 HPKE).
//!
//! # MLS compatibility
//!
//! The current MLS ciphersuite (MLS_128_DHKEMX25519_AES128GCM_SHA256_Ed25519)
//! uses 32-byte X25519 init keys in the wire format. This provider can produce
//! and consume **hybrid** init keys (1216-byte public, 2432-byte private), but
//! that is a non-standard extension: other MLS implementations will not
//! accept KeyPackages with hybrid init keys unless they implement the same
//! extension. This PoC validates that the OpenMLS trait surface is satisfiable
//! with a custom HPKE backend; full interoperability would require a new
//! ciphersuite or protocol extension.

use openmls_rust_crypto::RustCrypto;
use openmls_traits::{
    crypto::OpenMlsCrypto,
    types::{
        CryptoError, ExporterSecret, HpkeCiphertext, HpkeConfig, HpkeKeyPair, HpkeKemType,
        KemOutput,
    },
    OpenMlsProvider,
};
use tls_codec::SecretVLBytes;

use crate::hybrid_kem::{
    hybrid_decapsulate_only, hybrid_decrypt, hybrid_encapsulate_only, hybrid_encrypt,
    hybrid_export, HybridKeypair, HybridPublicKey,
    HYBRID_KEM_OUTPUT_LEN, HYBRID_PRIVATE_KEY_LEN, HYBRID_PUBLIC_KEY_LEN,
};
use crate::keystore::DiskKeyStore;

// Re-export types used by OpenMlsCrypto (full path for clarity).
use openmls_traits::types::{
    AeadType, Ciphersuite, HashType, SignatureScheme,
};

/// Crypto backend that uses hybrid KEM for HPKE when keys are in hybrid format,
/// and delegates everything else to RustCrypto.
///
/// When `hybrid_enabled` is `true`, `derive_hpke_keypair` produces hybrid keys
/// (1216-byte public, 2432-byte private). When `false`, it delegates to
/// RustCrypto and produces classical 32-byte X25519 keys.
///
/// The `hpke_seal` / `hpke_open` methods always detect the key format by length,
/// so they work correctly regardless of the flag — a hybrid-length key will use
/// hybrid KEM, a classical-length key will use RustCrypto.
#[derive(Debug)]
pub struct HybridCrypto {
    rust_crypto: RustCrypto,
    /// When true, `derive_hpke_keypair` produces hybrid (X25519 + ML-KEM-768)
    /// keys. When false, it produces classical X25519 keys via RustCrypto.
    hybrid_enabled: bool,
}

impl HybridCrypto {
    /// Create a hybrid-enabled crypto backend (derive_hpke_keypair produces hybrid keys).
    pub fn new() -> Self {
        Self {
            rust_crypto: RustCrypto::default(),
            hybrid_enabled: true,
        }
    }

    /// Alias for `new()` — hybrid mode enabled.
    pub fn new_hybrid() -> Self {
        Self::new()
    }

    /// Create a classical crypto backend (derive_hpke_keypair produces standard
    /// X25519 keys, but seal/open still accept hybrid keys by length detection).
    pub fn new_classical() -> Self {
        Self {
            rust_crypto: RustCrypto::default(),
            hybrid_enabled: false,
        }
    }

    /// Whether this backend produces hybrid keys from `derive_hpke_keypair`.
    pub fn is_hybrid_enabled(&self) -> bool {
        self.hybrid_enabled
    }

    /// Expose the underlying RustCrypto for rand() and delegation.
    pub fn rust_crypto(&self) -> &RustCrypto {
        &self.rust_crypto
    }

    fn is_hybrid_public_key(pk_r: &[u8]) -> bool {
        pk_r.len() == HYBRID_PUBLIC_KEY_LEN
    }

    fn is_hybrid_private_key(sk_r: &[u8]) -> bool {
        sk_r.len() == HYBRID_PRIVATE_KEY_LEN
    }
}

impl Default for HybridCrypto {
    fn default() -> Self {
        Self::new()
    }
}

impl OpenMlsCrypto for HybridCrypto {
    fn supports(&self, ciphersuite: Ciphersuite) -> Result<(), CryptoError> {
        self.rust_crypto.supports(ciphersuite)
    }

    fn supported_ciphersuites(&self) -> Vec<Ciphersuite> {
        self.rust_crypto.supported_ciphersuites()
    }

    fn hkdf_extract(
        &self,
        hash_type: HashType,
        salt: &[u8],
        ikm: &[u8],
    ) -> Result<SecretVLBytes, CryptoError> {
        self.rust_crypto.hkdf_extract(hash_type, salt, ikm)
    }

    fn hmac(
        &self,
        hash_type: HashType,
        key: &[u8],
        message: &[u8],
    ) -> Result<SecretVLBytes, CryptoError> {
        self.rust_crypto.hmac(hash_type, key, message)
    }

    fn hkdf_expand(
        &self,
        hash_type: HashType,
        prk: &[u8],
        info: &[u8],
        okm_len: usize,
    ) -> Result<SecretVLBytes, CryptoError> {
        self.rust_crypto.hkdf_expand(hash_type, prk, info, okm_len)
    }

    fn hash(&self, hash_type: HashType, data: &[u8]) -> Result<Vec<u8>, CryptoError> {
        self.rust_crypto.hash(hash_type, data)
    }

    fn aead_encrypt(
        &self,
        alg: AeadType,
        key: &[u8],
        data: &[u8],
        nonce: &[u8],
        aad: &[u8],
    ) -> Result<Vec<u8>, CryptoError> {
        self.rust_crypto.aead_encrypt(alg, key, data, nonce, aad)
    }

    fn aead_decrypt(
        &self,
        alg: AeadType,
        key: &[u8],
        ct_tag: &[u8],
        nonce: &[u8],
        aad: &[u8],
    ) -> Result<Vec<u8>, CryptoError> {
        self.rust_crypto.aead_decrypt(alg, key, ct_tag, nonce, aad)
    }

    fn signature_key_gen(&self, alg: SignatureScheme) -> Result<(Vec<u8>, Vec<u8>), CryptoError> {
        self.rust_crypto.signature_key_gen(alg)
    }

    fn verify_signature(
        &self,
        alg: SignatureScheme,
        data: &[u8],
        pk: &[u8],
        signature: &[u8],
    ) -> Result<(), CryptoError> {
        self.rust_crypto.verify_signature(alg, data, pk, signature)
    }

    fn sign(&self, alg: SignatureScheme, data: &[u8], key: &[u8]) -> Result<Vec<u8>, CryptoError> {
        self.rust_crypto.sign(alg, data, key)
    }

    fn hpke_seal(
        &self,
        config: HpkeConfig,
        pk_r: &[u8],
        info: &[u8],
        aad: &[u8],
        ptxt: &[u8],
    ) -> Result<HpkeCiphertext, CryptoError> {
        if Self::is_hybrid_public_key(pk_r) {
            let recipient_pk = HybridPublicKey::from_bytes(pk_r)
                .map_err(|_| CryptoError::CryptoLibraryError)?;
            let envelope = hybrid_encrypt(&recipient_pk, ptxt, info, aad)
                .map_err(|_| CryptoError::CryptoLibraryError)?;
            let kem_output = envelope[..HYBRID_KEM_OUTPUT_LEN].to_vec();
            let ciphertext = envelope[HYBRID_KEM_OUTPUT_LEN..].to_vec();
            Ok(HpkeCiphertext {
                kem_output: kem_output.into(),
                ciphertext: ciphertext.into(),
            })
        } else {
            self.rust_crypto.hpke_seal(config, pk_r, info, aad, ptxt)
        }
    }

    fn hpke_open(
        &self,
        config: HpkeConfig,
        input: &HpkeCiphertext,
        sk_r: &[u8],
        info: &[u8],
        aad: &[u8],
    ) -> Result<Vec<u8>, CryptoError> {
        if Self::is_hybrid_private_key(sk_r) {
            let keypair = HybridKeypair::from_private_bytes(sk_r)
                .map_err(|_| CryptoError::HpkeDecryptionError)?;
            let envelope: Vec<u8> = input
                .kem_output.as_slice()
                .iter()
                .chain(input.ciphertext.as_slice())
                .copied()
                .collect();
            // Pass HPKE info and aad through for proper context binding (RFC 9180).
            hybrid_decrypt(&keypair, &envelope, info, aad)
                .map_err(|_| CryptoError::HpkeDecryptionError)
        } else {
            self.rust_crypto.hpke_open(config, input, sk_r, info, aad)
        }
    }

    fn hpke_setup_sender_and_export(
        &self,
        config: HpkeConfig,
        pk_r: &[u8],
        info: &[u8],
        exporter_context: &[u8],
        exporter_length: usize,
    ) -> Result<(KemOutput, ExporterSecret), CryptoError> {
        if Self::is_hybrid_public_key(pk_r) {
            // A key that passes the hybrid length check but fails deserialization
            // is corrupted — return an error instead of silently downgrading to
            // classical crypto (which would defeat PQ protection).
            let recipient_pk = HybridPublicKey::from_bytes(pk_r)
                .map_err(|_| CryptoError::SenderSetupError)?;
            let (kem_output, shared_secret) =
                hybrid_encapsulate_only(&recipient_pk).map_err(|_| CryptoError::SenderSetupError)?;
            let exported = hybrid_export(&shared_secret, exporter_context, exporter_length);
            Ok((kem_output, exported.into()))
        } else {
            self.rust_crypto.hpke_setup_sender_and_export(
                config, pk_r, info, exporter_context, exporter_length,
            )
        }
    }

    fn hpke_setup_receiver_and_export(
        &self,
        config: HpkeConfig,
        enc: &[u8],
        sk_r: &[u8],
        info: &[u8],
        exporter_context: &[u8],
        exporter_length: usize,
    ) -> Result<ExporterSecret, CryptoError> {
        if Self::is_hybrid_private_key(sk_r) {
            let keypair = HybridKeypair::from_private_bytes(sk_r)
                .map_err(|_| CryptoError::ReceiverSetupError)?;
            let shared_secret =
                hybrid_decapsulate_only(&keypair, enc).map_err(|_| CryptoError::ReceiverSetupError)?;
            let exported = hybrid_export(&shared_secret, exporter_context, exporter_length);
            Ok(exported.into())
        } else {
            self.rust_crypto.hpke_setup_receiver_and_export(
                config, enc, sk_r, info, exporter_context, exporter_length,
            )
        }
    }

    fn derive_hpke_keypair(&self, config: HpkeConfig, ikm: &[u8]) -> Result<HpkeKeyPair, CryptoError> {
        if self.hybrid_enabled && config.0 == HpkeKemType::DhKem25519 {
            let kp = HybridKeypair::derive_from_ikm(ikm);
            let private_bytes = kp.private_to_bytes();
            Ok(HpkeKeyPair {
                private: private_bytes.as_slice().into(),
                public: kp.public_key().to_bytes(),
            })
        } else {
            self.rust_crypto.derive_hpke_keypair(config, ikm)
        }
    }
}

/// OpenMLS crypto provider that uses hybrid KEM for HPKE (when keys are in
/// hybrid format) and delegates the rest to RustCrypto.
#[derive(Debug)]
pub struct HybridCryptoProvider {
    crypto: HybridCrypto,
    key_store: DiskKeyStore,
}

impl HybridCryptoProvider {
    /// Create a hybrid-enabled provider (KeyPackages will contain hybrid init keys).
    pub fn new(key_store: DiskKeyStore) -> Self {
        Self {
            crypto: HybridCrypto::new_hybrid(),
            key_store,
        }
    }

    /// Alias for `new()` — hybrid mode enabled.
    pub fn new_hybrid(key_store: DiskKeyStore) -> Self {
        Self::new(key_store)
    }

    /// Create a classical-mode provider (KeyPackages use standard X25519 init keys,
    /// but seal/open still accept hybrid keys by length detection).
    pub fn new_classical(key_store: DiskKeyStore) -> Self {
        Self {
            crypto: HybridCrypto::new_classical(),
            key_store,
        }
    }

    /// Whether this provider produces hybrid keys from `derive_hpke_keypair`.
    pub fn is_hybrid_enabled(&self) -> bool {
        self.crypto.is_hybrid_enabled()
    }
}

impl Default for HybridCryptoProvider {
    fn default() -> Self {
        Self::new(DiskKeyStore::ephemeral())
    }
}

impl OpenMlsProvider for HybridCryptoProvider {
    type CryptoProvider = HybridCrypto;
    type RandProvider = RustCrypto;
    type StorageProvider = DiskKeyStore;

    fn crypto(&self) -> &Self::CryptoProvider {
        &self.crypto
    }

    fn rand(&self) -> &Self::RandProvider {
        self.crypto.rust_crypto()
    }

    fn storage(&self) -> &Self::StorageProvider {
        &self.key_store
    }
}

// ── Tests ───────────────────────────────────────────────────────────────────

#[cfg(test)]
#[allow(clippy::unwrap_used)]
mod tests {
    use super::*;
    use openmls_traits::types::HpkeKdfType;

    fn hpke_config_dhkem_x25519() -> HpkeConfig {
        HpkeConfig(
            HpkeKemType::DhKem25519,
            HpkeKdfType::HkdfSha256,
            openmls_traits::types::HpkeAeadType::AesGcm128,
        )
    }

    /// HPKE path with hybrid keys: derive_hpke_keypair (hybrid) -> hpke_seal -> hpke_open.
    #[test]
    fn hybrid_hpke_seal_open_round_trip() {
        let crypto = HybridCrypto::new();
        let ikm = b"test-ikm-for-hybrid-hpke-keypair";

        let keypair = crypto.derive_hpke_keypair(hpke_config_dhkem_x25519(), ikm).unwrap();
        assert_eq!(keypair.public.len(), HYBRID_PUBLIC_KEY_LEN);
        assert_eq!(keypair.private.as_ref().len(), HYBRID_PRIVATE_KEY_LEN);

        let plaintext = b"hello post-quantum MLS";
        let info = b"mls 1.0 test";
        let aad = b"additional data";

        let ct = crypto.hpke_seal(
            hpke_config_dhkem_x25519(),
            &keypair.public,
            info,
            aad,
            plaintext,
        ).unwrap();
        assert!(!ct.kem_output.as_slice().is_empty());
        assert!(!ct.ciphertext.as_slice().is_empty());

        let decrypted = crypto
            .hpke_open(
                hpke_config_dhkem_x25519(),
                &ct,
                keypair.private.as_ref(),
                info,
                aad,
            )
            .expect("hpke_open with hybrid keys");
        assert_eq!(decrypted.as_slice(), plaintext);
    }

    /// HPKE exporter path: setup_sender_and_export then setup_receiver_and_export.
    #[test]
    fn hybrid_hpke_setup_sender_receiver_export() {
        let crypto = HybridCrypto::new();
        let ikm = b"exporter-ikm";

        let keypair = crypto.derive_hpke_keypair(hpke_config_dhkem_x25519(), ikm).unwrap();
        let info = b"";
        let exporter_context = b"MLS 1.0 external init";
        let exporter_length = 32;

        let (kem_output, sender_exported) = crypto
            .hpke_setup_sender_and_export(
                hpke_config_dhkem_x25519(),
                &keypair.public,
                info,
                exporter_context,
                exporter_length,
            )
            .expect("sender and export");

        assert_eq!(kem_output.len(), HYBRID_KEM_OUTPUT_LEN);
        assert_eq!(sender_exported.as_ref().len(), exporter_length);

        let receiver_exported = crypto
            .hpke_setup_receiver_and_export(
                hpke_config_dhkem_x25519(),
                &kem_output,
                keypair.private.as_ref(),
                info,
                exporter_context,
                exporter_length,
            )
            .expect("receiver and export");

        assert_eq!(sender_exported.as_ref(), receiver_exported.as_ref());
    }

    /// Classical mode: derive_hpke_keypair produces standard 32-byte X25519 keys.
    #[test]
    fn classical_mode_produces_standard_keys() {
        let crypto = HybridCrypto::new_classical();
        let ikm = b"test-ikm-for-classical-hpke";

        let keypair = crypto.derive_hpke_keypair(hpke_config_dhkem_x25519(), ikm).unwrap();
        // Classical X25519 keys are 32 bytes
        assert_eq!(keypair.public.len(), 32);
        assert_eq!(keypair.private.as_ref().len(), 32);
    }

    /// Classical mode round-trip: seal/open works with classical keys.
    #[test]
    fn classical_mode_seal_open_round_trip() {
        let crypto = HybridCrypto::new_classical();
        let ikm = b"test-ikm-for-classical-round-trip";

        let keypair = crypto.derive_hpke_keypair(hpke_config_dhkem_x25519(), ikm).unwrap();
        assert_eq!(keypair.public.len(), 32); // classical key

        let plaintext = b"hello classical MLS";
        let info = b"mls 1.0 test";
        let aad = b"additional data";

        let ct = crypto.hpke_seal(
            hpke_config_dhkem_x25519(),
            &keypair.public,
            info,
            aad,
            plaintext,
        ).unwrap();
        assert!(!ct.kem_output.as_slice().is_empty());

        let decrypted = crypto
            .hpke_open(
                hpke_config_dhkem_x25519(),
                &ct,
                keypair.private.as_ref(),
                info,
                aad,
            )
            .expect("hpke_open with classical keys");
        assert_eq!(decrypted.as_slice(), plaintext);
    }

    /// KeyPackage generation with HybridCryptoProvider (validates full HPKE path in MLS).
    #[test]
    fn key_package_generation_with_hybrid_provider() {
        use openmls::prelude::{
            BasicCredential, CredentialWithKey, KeyPackage,
        };
        use std::sync::Arc;
        use tls_codec::Serialize;

        use crate::identity::IdentityKeypair;

        const CIPHERSUITE: Ciphersuite =
            Ciphersuite::MLS_128_DHKEMX25519_AES128GCM_SHA256_Ed25519;

        let provider = HybridCryptoProvider::default();
        let identity = Arc::new(IdentityKeypair::generate());

        let credential: openmls::prelude::Credential =
            BasicCredential::new(identity.public_key_bytes().to_vec()).into();
        let credential_with_key = CredentialWithKey {
            credential,
            signature_key: identity.public_key_bytes().to_vec().into(),
        };

        let key_package_bundle = KeyPackage::builder()
            .build(
                CIPHERSUITE,
                &provider,
                identity.as_ref(),
                credential_with_key,
            )
            .expect("KeyPackage with hybrid HPKE");

        let bytes = key_package_bundle
            .key_package()
            .tls_serialize_detached()
            .expect("serialize KeyPackage");
        assert!(!bytes.is_empty());
    }
}