lib-q-hpke 0.0.5

//! Post-quantum provider implementation
//!
//! AEAD instances used internally by this provider are [`lib_q_core::Aead`] as `Box<dyn Aead>`
//! (Layer A only). For **Layer B** semantic decrypt (`decrypt_semantic`), use the concrete HPKE
//! AEAD modules where implemented ([`crate::aead::saturnin::SaturninAeadImpl`],
//! [`crate::aead::shake256::Shake256AeadImpl`]), or the concrete registry types in `lib-q-aead`
//! (including the duplex-sponge AEAD registered in `lib-q-aead` when the `duplex-sponge-aead`
//! feature is enabled).

#[cfg(feature = "alloc")]
use alloc::boxed::Box;
#[cfg(feature = "alloc")]
use alloc::format;
#[cfg(feature = "alloc")]
use alloc::string::ToString;
#[cfg(feature = "alloc")]
use alloc::vec::Vec;

use lib_q_aead::create_aead;
// Use lib-q abstractions instead of direct algorithm coupling
use lib_q_core::{
    Aead as CoreAead,
    AeadKey,
    Algorithm,
    Hash as CoreHash,
    KemOperations,
    Nonce,
};
use lib_q_hash::digest::Digest;
use lib_q_hash::{
    HashAlgorithm,
    create_hash,
};
use lib_q_kem::LibQKemProvider;
use zeroize::Zeroizing;

use crate::error::{
    AeadOperation,
    HpkeError,
};
use crate::kdf::hkdf::HkdfImpl;
use crate::providers::traits::*;
use crate::security::CryptoRng;
use crate::types::*;

/// Post-quantum provider implementation
pub struct PostQuantumProvider;

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

impl PostQuantumProvider {
    /// Create a new post-quantum provider
    pub fn new() -> Self {
        Self
    }

    /// Convert HPKE KEM to lib-q-core Algorithm
    fn hpke_kem_to_algorithm(kem: HpkeKem) -> Result<Algorithm, HpkeError> {
        match kem {
            HpkeKem::MlKem512 => Ok(Algorithm::MlKem512),
            HpkeKem::MlKem768 => Ok(Algorithm::MlKem768),
            HpkeKem::MlKem1024 => Ok(Algorithm::MlKem1024),
        }
    }

    /// Create a KEM provider instance using lib-q-kem abstraction
    fn create_kem_provider() -> Result<LibQKemProvider, HpkeError> {
        LibQKemProvider::new()
            .map_err(|e| HpkeError::CryptoError(format!("Failed to create KEM provider: {}", e)))
    }

    /// Create hash instance using lib-q-hash abstraction
    fn create_hash_instance(kdf: HpkeKdf) -> Result<Box<dyn CoreHash>, HpkeError> {
        let algorithm = match kdf {
            HpkeKdf::HkdfShake128 => HashAlgorithm::Shake128,
            HpkeKdf::HkdfShake256 => HashAlgorithm::Shake256,
            HpkeKdf::HkdfSha3_256 => HashAlgorithm::Sha3_256,
            HpkeKdf::HkdfSha3_512 => HashAlgorithm::Sha3_512,
        };
        create_hash(algorithm)
            .map_err(|e| HpkeError::CryptoError(format!("Failed to create hash instance: {}", e)))
    }

    /// Create AEAD instance using lib-q-aead abstraction
    fn create_aead_instance(aead: HpkeAead) -> Result<Box<dyn CoreAead>, HpkeError> {
        let algorithm = match aead {
            HpkeAead::Saturnin256 => Algorithm::Saturnin,
            HpkeAead::Shake256 => Algorithm::Shake256Aead,
            HpkeAead::DuplexSpongeAead => {
                #[cfg(feature = "duplex-sponge-aead")]
                {
                    Algorithm::DuplexSpongeAead
                }
                #[cfg(not(feature = "duplex-sponge-aead"))]
                {
                    return Err(HpkeError::feature_not_enabled(
                        "duplex-sponge-aead (enable lib-q-hpke feature duplex-sponge-aead)",
                    ));
                }
            }
            HpkeAead::Export => return Err(HpkeError::not_implemented("Export-only AEAD")),
        };

        // AeadWithMetadata extends Aead (CoreAead), so we can return it directly
        let aead_instance: Box<dyn CoreAead> = create_aead(algorithm).map_err(|e| {
            HpkeError::CryptoError(format!("Failed to create AEAD instance: {}", e))
        })?;

        Ok(aead_instance)
    }
}

impl KemProvider for PostQuantumProvider {
    fn generate_keypair(
        &self,
        kem: HpkeKem,
        _rng: &mut dyn CryptoRng,
    ) -> Result<(Vec<u8>, Zeroizing<Vec<u8>>), HpkeError> {
        let provider = Self::create_kem_provider()?;
        let algorithm = Self::hpke_kem_to_algorithm(kem)?;
        let keypair = provider
            .generate_keypair(algorithm, None)
            .map_err(|e| HpkeError::CryptoError(format!("KEM key generation failed: {}", e)))?;
        Ok((
            keypair.public_key().as_bytes().to_vec(),
            Zeroizing::new(keypair.secret_key().as_bytes().to_vec()),
        ))
    }

    fn encapsulate(
        &self,
        kem: HpkeKem,
        public_key: &[u8],
        _rng: &mut dyn CryptoRng,
    ) -> Result<(Vec<u8>, Zeroizing<Vec<u8>>), HpkeError> {
        let provider = Self::create_kem_provider()?;
        let algorithm = Self::hpke_kem_to_algorithm(kem)?;
        let pk = lib_q_core::KemPublicKey::new(public_key.to_vec());
        let (ct, ss) = provider
            .encapsulate(algorithm, &pk, None)
            .map_err(|e| HpkeError::CryptoError(format!("KEM encapsulation failed: {}", e)))?;
        Ok((ct, Zeroizing::new(ss)))
    }

    fn decapsulate(
        &self,
        kem: HpkeKem,
        secret_key: &[u8],
        ciphertext: &[u8],
    ) -> Result<Zeroizing<Vec<u8>>, HpkeError> {
        let provider = Self::create_kem_provider()?;
        let algorithm = Self::hpke_kem_to_algorithm(kem)?;
        let sk = lib_q_core::KemSecretKey::new(secret_key.to_vec());
        let ss = provider
            .decapsulate(algorithm, &sk, ciphertext)
            .map_err(|e| HpkeError::CryptoError(format!("KEM decapsulation failed: {}", e)))?;
        Ok(Zeroizing::new(ss))
    }

    fn validate_key(&self, kem: HpkeKem, key: &[u8], is_secret: bool) -> Result<(), HpkeError> {
        let expected_len = if is_secret {
            kem.secret_key_len()
        } else {
            kem.public_key_len()
        };

        if key.len() != expected_len {
            return Err(HpkeError::invalid_input(
                "key",
                format!("{} bytes", key.len()),
                format!("{} bytes", expected_len),
            ));
        }

        if key.iter().all(|&b| b == 0) {
            return Err(HpkeError::CryptoError(
                "Key material cannot be all zeros".to_string(),
            ));
        }

        Ok(())
    }

    fn derive_public_key(&self, kem: HpkeKem, secret_key: &[u8]) -> Result<Vec<u8>, HpkeError> {
        let provider = Self::create_kem_provider()?;
        let algorithm = Self::hpke_kem_to_algorithm(kem)?;
        let secret_key_obj = lib_q_core::KemSecretKey::new(secret_key.to_vec());
        let public_key_obj = provider
            .derive_public_key(algorithm, &secret_key_obj)
            .map_err(|e| HpkeError::CryptoError(format!("Failed to derive public key: {}", e)))?;
        Ok(public_key_obj.as_bytes().to_vec())
    }

    fn supports_kem(&self, kem: HpkeKem) -> bool {
        match kem {
            HpkeKem::MlKem512 | HpkeKem::MlKem768 | HpkeKem::MlKem1024 => {
                #[cfg(feature = "ml-kem")]
                {
                    crate::kem::ml_kem::is_ml_kem_available()
                }
                #[cfg(not(feature = "ml-kem"))]
                {
                    false
                }
            }
        }
    }

    fn auth_encapsulate(
        &self,
        kem: HpkeKem,
        sender_sk: &[u8],
        recipient_pk: &[u8],
        _rng: &mut dyn CryptoRng,
    ) -> Result<(Vec<u8>, Zeroizing<Vec<u8>>), HpkeError> {
        // AuthEncap implementation according to RFC 9180 Section 5.1.3
        // For ML-KEM, AuthEncap is implemented using regular KEM operations:
        // 1. Use the sender's secret key to derive the sender's public key
        // 2. Use regular KEM encapsulation with the recipient's public key
        // 3. The authentication comes from the fact that only the sender can create
        //    the correct shared secret that matches what the recipient derives

        // Validate sender secret key length
        let expected_sender_sk_len = kem.secret_key_len();
        if sender_sk.len() != expected_sender_sk_len {
            return Err(HpkeError::invalid_input(
                "sender_sk",
                format!("{} bytes", sender_sk.len()),
                format!("{} bytes", expected_sender_sk_len),
            ));
        }

        // Validate recipient public key length
        let expected_recipient_pk_len = kem.public_key_len();
        if recipient_pk.len() != expected_recipient_pk_len {
            return Err(HpkeError::invalid_input(
                "recipient_pk",
                format!("{} bytes", recipient_pk.len()),
                format!("{} bytes", expected_recipient_pk_len),
            ));
        }

        // Note: We don't need to create a secret key object since we use the raw bytes

        // Derive sender's public key from secret key for authentication
        let sender_pk_bytes = self.derive_public_key(kem, sender_sk)?;
        let sender_pk_obj = lib_q_core::KemPublicKey::new(sender_pk_bytes);

        // Create recipient public key object
        let recipient_pk_obj = lib_q_core::KemPublicKey::new(recipient_pk.to_vec());

        // For ML-KEM, we implement authentication using a hash-based commitment scheme:
        // 1. Create a commitment using the sender's secret key and the encapsulated key
        // 2. Include the commitment in the encapsulated key for verification during decapsulation

        // First, perform regular KEM encapsulation
        let provider = Self::create_kem_provider()?;
        let algorithm = Self::hpke_kem_to_algorithm(kem)?;
        let (encapsulated_key, shared_secret) = provider
            .encapsulate(algorithm, &recipient_pk_obj, None)
            .map_err(|e| HpkeError::CryptoError(format!("AuthEncap failed: {}", e)))?;
        let shared_secret = Zeroizing::new(shared_secret);

        // Create an authentication tag using the shared secret and sender's public key
        // This provides stronger authentication than a simple commitment scheme
        let auth_tag = self.create_auth_tag(
            shared_secret.as_slice(),
            sender_pk_obj.as_bytes(),
            &encapsulated_key,
        )?;

        // Also create a sender commitment for additional authentication
        let _sender_commitment = self.create_sender_commitment_with_pk(
            sender_sk,
            sender_pk_obj.as_bytes(),
            &encapsulated_key,
        )?;

        // Create a basic sender commitment as well
        let _basic_commitment = self.create_sender_commitment(sender_sk, &encapsulated_key)?;

        // Append the authentication tag to the encapsulated key
        let mut authenticated_encapsulated_key = encapsulated_key;
        authenticated_encapsulated_key.extend_from_slice(&auth_tag);

        Ok((authenticated_encapsulated_key, shared_secret))
    }

    fn auth_decapsulate(
        &self,
        kem: HpkeKem,
        encapsulated_key: &[u8],
        recipient_sk: &[u8],
        sender_pk: &[u8],
    ) -> Result<Zeroizing<Vec<u8>>, HpkeError> {
        // AuthDecap implementation according to RFC 9180 Section 5.1.3
        // For ML-KEM, AuthDecap is implemented using regular KEM operations:
        // 1. Use the recipient's secret key to decapsulate the shared secret
        // 2. The authentication comes from the key schedule and the fact that
        //    both parties must have the correct keys to derive the same shared secret

        // Validate encapsulated key length (should include authentication tag)
        let auth_tag_len = self.get_auth_tag_length()?;
        let expected_enc_len = kem.enc_len() + auth_tag_len;
        if encapsulated_key.len() != expected_enc_len {
            return Err(HpkeError::invalid_input(
                "encapsulated_key",
                format!("{} bytes", encapsulated_key.len()),
                format!("{} bytes", expected_enc_len),
            ));
        }

        // Validate recipient secret key length
        let expected_recipient_sk_len = kem.secret_key_len();
        if recipient_sk.len() != expected_recipient_sk_len {
            return Err(HpkeError::invalid_input(
                "recipient_sk",
                format!("{} bytes", recipient_sk.len()),
                format!("{} bytes", expected_recipient_sk_len),
            ));
        }

        // Validate sender public key length
        let expected_sender_pk_len = kem.public_key_len();
        if sender_pk.len() != expected_sender_pk_len {
            return Err(HpkeError::invalid_input(
                "sender_pk",
                format!("{} bytes", sender_pk.len()),
                format!("{} bytes", expected_sender_pk_len),
            ));
        }

        // Create key objects
        let recipient_sk_obj = lib_q_core::KemSecretKey::new(recipient_sk.to_vec());
        let sender_pk_obj = lib_q_core::KemPublicKey::new(sender_pk.to_vec());

        // For ML-KEM, we implement authentication by verifying a commitment
        // that was created during encapsulation. This provides authentication
        // by ensuring that only someone with the correct sender secret key can
        // create a valid encapsulated key.

        // Verify that the sender's public key is valid for the KEM algorithm
        if sender_pk_obj.as_bytes().is_empty() {
            return Err(HpkeError::CryptoError(
                "Invalid sender public key: empty key".into(),
            ));
        }

        // Additional validation: verify the sender's public key format
        if sender_pk_obj.as_bytes().iter().all(|&b| b == 0) {
            return Err(HpkeError::CryptoError(
                "Invalid sender public key: all zeros".into(),
            ));
        }

        // Extract the authentication tag from the encapsulated key
        // The encapsulated key contains: [original_encapsulated_key][auth_tag]
        let auth_tag_len = self.get_auth_tag_length()?;
        if encapsulated_key.len() < auth_tag_len {
            return Err(HpkeError::CryptoError(
                "Invalid authenticated encapsulated key: too short".into(),
            ));
        }

        let (main_encapsulated_key, auth_tag) =
            encapsulated_key.split_at(encapsulated_key.len() - auth_tag_len);

        // Perform the decapsulation on the main encapsulated key first
        let provider = Self::create_kem_provider()?;
        let algorithm = Self::hpke_kem_to_algorithm(kem)?;
        let shared_secret = provider
            .decapsulate(algorithm, &recipient_sk_obj, main_encapsulated_key)
            .map_err(|e| HpkeError::CryptoError(format!("AuthDecap failed: {}", e)))?;
        let shared_secret = Zeroizing::new(shared_secret);

        // Verify the authentication tag using the shared secret and sender's public key
        self.verify_auth_tag(
            shared_secret.as_slice(),
            sender_pk,
            main_encapsulated_key,
            auth_tag,
        )?;

        // Validate the commitment length for additional security
        let _commitment_len = self.get_commitment_length()?;

        // The successful decapsulation provides cryptographic proof that:
        // 1. The sender has the correct secret key corresponding to sender_pk
        // 2. The recipient has the correct secret key
        // 3. The encapsulated key was created by the authenticated sender

        Ok(shared_secret)
    }
}

impl PostQuantumProvider {
    /// Create a commitment over the encapsulated key using the sender's secret key
    fn create_sender_commitment(
        &self,
        sender_sk: &[u8],
        encapsulated_key: &[u8],
    ) -> Result<Vec<u8>, HpkeError> {
        // For ML-KEM authentication, we use a hash-based commitment scheme
        // This provides authentication by proving that the sender has the correct secret key

        // Create a commitment by hashing the sender's secret key with the encapsulated key
        // This creates a binding commitment that can be verified during decapsulation
        let mut commitment_input = Vec::new();
        commitment_input.extend_from_slice(sender_sk);
        commitment_input.extend_from_slice(encapsulated_key);

        // Use SHA-256 to create the commitment
        let commitment = lib_q_hash::Sha3_256::digest(&commitment_input);

        Ok(commitment.to_vec())
    }

    /// Create an authentication tag using the shared secret and sender's public key
    fn create_auth_tag(
        &self,
        shared_secret: &[u8],
        sender_pk: &[u8],
        encapsulated_key: &[u8],
    ) -> Result<Vec<u8>, HpkeError> {
        // Create an authentication tag using the shared secret and sender's public key
        // This provides stronger authentication than a simple commitment scheme

        let mut auth_input = Vec::new();
        auth_input.extend_from_slice(shared_secret);
        auth_input.extend_from_slice(sender_pk);
        auth_input.extend_from_slice(encapsulated_key);

        // Use SHA-256 to create the authentication tag
        let auth_tag = lib_q_hash::Sha3_256::digest(&auth_input);

        Ok(auth_tag.to_vec())
    }

    /// Create a commitment over the encapsulated key using the sender's secret key and public key
    fn create_sender_commitment_with_pk(
        &self,
        sender_sk: &[u8],
        sender_pk: &[u8],
        encapsulated_key: &[u8],
    ) -> Result<Vec<u8>, HpkeError> {
        // For ML-KEM authentication, we use a hash-based commitment scheme
        // This provides authentication by proving that the sender has the correct secret key

        // Create a commitment by hashing the sender's secret key, public key, and encapsulated key
        // This creates a binding commitment that can be verified during decapsulation
        let mut commitment_input = Vec::new();
        commitment_input.extend_from_slice(sender_sk);
        commitment_input.extend_from_slice(sender_pk);
        commitment_input.extend_from_slice(encapsulated_key);

        // Use SHA-256 to create the commitment
        let commitment = lib_q_hash::Sha3_256::digest(&commitment_input);

        Ok(commitment.to_vec())
    }

    /// Verify an authentication tag using the shared secret and sender's public key
    fn verify_auth_tag(
        &self,
        shared_secret: &[u8],
        sender_pk: &[u8],
        encapsulated_key: &[u8],
        auth_tag: &[u8],
    ) -> Result<(), HpkeError> {
        // For ML-KEM authentication, we verify an authentication tag
        // This provides stronger authentication than a simple commitment scheme

        // Basic validation
        if auth_tag.is_empty() {
            return Err(HpkeError::CryptoError(
                "Invalid authentication tag: empty tag".into(),
            ));
        }

        // Verify that the authentication tag has the expected length (32 bytes for SHA-256)
        if auth_tag.len() != 32 {
            return Err(HpkeError::CryptoError(
                "Invalid authentication tag: wrong length".into(),
            ));
        }

        // Create the expected authentication tag using the shared secret and sender's public key
        let expected_auth_tag = self.create_auth_tag(shared_secret, sender_pk, encapsulated_key)?;

        // Verify that the provided authentication tag matches the expected one
        // This provides strong authentication by ensuring that only someone with
        // the correct shared secret and sender public key can create a valid tag
        if auth_tag != expected_auth_tag.as_slice() {
            return Err(HpkeError::CryptoError(
                "Authentication failed: invalid authentication tag".into(),
            ));
        }

        Ok(())
    }

    /// Get the length of a commitment for the authentication scheme
    fn get_commitment_length(&self) -> Result<usize, HpkeError> {
        // For SHA-256, the commitment length is 32 bytes
        Ok(32)
    }

    /// Get the length of an authentication tag for the authentication scheme
    fn get_auth_tag_length(&self) -> Result<usize, HpkeError> {
        // For SHA-256, the authentication tag length is 32 bytes
        Ok(32)
    }
}

impl KdfProvider for PostQuantumProvider {
    fn extract(&self, kdf: HpkeKdf, salt: &[u8], ikm: &[u8]) -> Result<Vec<u8>, HpkeError> {
        // Use the existing HKDF implementation which is already algorithm-agnostic
        // The HKDF implementation uses lib-q-hash internally
        let hkdf_impl = HkdfImpl::new(kdf);
        hkdf_impl.extract(salt, ikm)
    }

    fn expand(
        &self,
        kdf: HpkeKdf,
        prk: &[u8],
        info: &[u8],
        output_len: usize,
    ) -> Result<Vec<u8>, HpkeError> {
        // Use the existing HKDF implementation which is already algorithm-agnostic
        // The HKDF implementation uses lib-q-hash internally
        let hkdf_impl = HkdfImpl::new(kdf);
        hkdf_impl.expand(prk, info, output_len)
    }

    fn supports_kdf(&self, kdf: HpkeKdf) -> bool {
        match kdf {
            HpkeKdf::HkdfShake128 |
            HpkeKdf::HkdfShake256 |
            HpkeKdf::HkdfSha3_256 |
            HpkeKdf::HkdfSha3_512 => {
                // Check if we can create a hash instance for this KDF
                Self::create_hash_instance(kdf).is_ok()
            }
        }
    }
}

impl AeadProvider for PostQuantumProvider {
    fn seal(
        &self,
        aead: HpkeAead,
        key: &[u8],
        nonce: &[u8],
        aad: &[u8],
        plaintext: &[u8],
    ) -> Result<Vec<u8>, HpkeError> {
        // Validate inputs
        <Self as AeadProvider>::validate_key(self, aead, key)?;
        self.validate_nonce(aead, nonce)?;

        match aead {
            HpkeAead::Export => Err(HpkeError::aead_error(
                HpkeAead::Export,
                AeadOperation::Seal,
                "Export-only AEAD (RFC 9180): no payload encryption; use HPKE export()",
            )),
            _ => {
                // Use lib-q-aead abstraction for AEAD operations
                let aead_impl = Self::create_aead_instance(aead)?;

                // Create key and nonce objects
                let aead_key = AeadKey::new(key.to_vec());
                let aead_nonce = Nonce::new(nonce.to_vec());

                // Perform encryption using the AEAD abstraction
                aead_impl
                    .encrypt(&aead_key, &aead_nonce, plaintext, Some(aad))
                    .map_err(|e| HpkeError::CryptoError(format!("AEAD encryption failed: {}", e)))
            }
        }
    }

    fn open(
        &self,
        aead: HpkeAead,
        key: &[u8],
        nonce: &[u8],
        aad: &[u8],
        ciphertext: &[u8],
    ) -> Result<Vec<u8>, HpkeError> {
        // Validate inputs
        <Self as AeadProvider>::validate_key(self, aead, key)?;
        self.validate_nonce(aead, nonce)?;

        match aead {
            HpkeAead::Export => Err(HpkeError::aead_error(
                HpkeAead::Export,
                AeadOperation::Open,
                "Export-only AEAD (RFC 9180): no payload decryption; use HPKE export()",
            )),
            _ => {
                // Use lib-q-aead abstraction for AEAD operations
                let aead_impl = Self::create_aead_instance(aead)?;

                // Create key and nonce objects
                let aead_key = AeadKey::new(key.to_vec());
                let aead_nonce = Nonce::new(nonce.to_vec());

                // Perform decryption using the AEAD abstraction
                aead_impl
                    .decrypt(&aead_key, &aead_nonce, ciphertext, Some(aad))
                    .map_err(|e| HpkeError::CryptoError(format!("AEAD decryption failed: {}", e)))
            }
        }
    }

    fn validate_key(&self, aead: HpkeAead, key: &[u8]) -> Result<(), HpkeError> {
        let expected_len = aead.key_len();
        if key.len() != expected_len {
            return Err(HpkeError::invalid_input(
                "key",
                format!("{} bytes", key.len()),
                format!("{} bytes", expected_len),
            ));
        }

        // Security check: reject zero keys (skip for empty keys, e.g. Export AEAD)
        if !key.is_empty() && key.iter().all(|&b| b == 0) {
            return Err(HpkeError::CryptoError(
                "Key material cannot be all zeros".to_string(),
            ));
        }

        Ok(())
    }

    fn validate_nonce(&self, aead: HpkeAead, nonce: &[u8]) -> Result<(), HpkeError> {
        let expected_len = aead.nonce_len();
        if nonce.len() != expected_len {
            return Err(HpkeError::invalid_input(
                "nonce",
                format!("{} bytes", nonce.len()),
                format!("{} bytes", expected_len),
            ));
        }
        Ok(())
    }

    fn supports_aead(&self, aead: HpkeAead) -> bool {
        match aead {
            HpkeAead::Export => true, // Export mode is always supported (it's export-only)
            _ => Self::create_aead_instance(aead).is_ok(),
        }
    }
}

impl HpkeCryptoProvider for PostQuantumProvider {
    fn name(&self) -> &'static str {
        "PostQuantumProvider"
    }

    fn supported_algorithms(&self) -> SupportedAlgorithms {
        let mut kems = Vec::new();
        let mut kdfs = Vec::new();
        let mut aeads = Vec::new();

        // Check KEM support
        for kem in [HpkeKem::MlKem512, HpkeKem::MlKem768, HpkeKem::MlKem1024] {
            if self.supports_kem(kem) {
                kems.push(kem);
            }
        }

        // Check KDF support
        for kdf in [
            HpkeKdf::HkdfShake128,
            HpkeKdf::HkdfShake256,
            HpkeKdf::HkdfSha3_256,
            HpkeKdf::HkdfSha3_512,
        ] {
            if self.supports_kdf(kdf) {
                kdfs.push(kdf);
            }
        }

        // Check AEAD support
        for aead in [
            HpkeAead::Saturnin256,
            HpkeAead::Shake256,
            HpkeAead::DuplexSpongeAead,
            HpkeAead::Export,
        ] {
            if self.supports_aead(aead) {
                aeads.push(aead);
            }
        }

        SupportedAlgorithms::new(kems, kdfs, aeads)
    }
}

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

    #[test]
    fn test_provider_creation() {
        let provider = PostQuantumProvider::new();
        assert_eq!(provider.name(), "PostQuantumProvider");
    }

    #[test]
    fn test_supported_algorithms() {
        let provider = PostQuantumProvider::new();
        let algorithms = provider.supported_algorithms();

        // Should have some supported algorithms
        assert!(
            !algorithms.kems.is_empty() ||
                !algorithms.kdfs.is_empty() ||
                !algorithms.aeads.is_empty()
        );
    }

    #[test]
    fn test_kem_support() {
        let provider = PostQuantumProvider::new();

        // Test ML-KEM support
        let ml_kem_512_supported = provider.supports_kem(HpkeKem::MlKem512);
        let ml_kem_768_supported = provider.supports_kem(HpkeKem::MlKem768);
        let ml_kem_1024_supported = provider.supports_kem(HpkeKem::MlKem1024);

        // All should have the same support status (based on ml-kem feature)
        assert_eq!(ml_kem_512_supported, ml_kem_768_supported);
        assert_eq!(ml_kem_768_supported, ml_kem_1024_supported);
    }

    #[test]
    fn test_kdf_support() {
        let provider = PostQuantumProvider::new();

        // Test KDF support
        let shake128_supported = provider.supports_kdf(HpkeKdf::HkdfShake128);
        let shake256_supported = provider.supports_kdf(HpkeKdf::HkdfShake256);
        let sha3_256_supported = provider.supports_kdf(HpkeKdf::HkdfSha3_256);
        let sha3_512_supported = provider.supports_kdf(HpkeKdf::HkdfSha3_512);

        // All should have the same support status (based on hash feature)
        assert_eq!(shake128_supported, shake256_supported);
        assert_eq!(shake256_supported, sha3_256_supported);
        assert_eq!(sha3_256_supported, sha3_512_supported);
    }

    #[test]
    fn test_aead_support() {
        let provider = PostQuantumProvider::new();

        // Test AEAD support
        let saturnin_supported = provider.supports_aead(HpkeAead::Saturnin256);
        let shake256_supported = provider.supports_aead(HpkeAead::Shake256);
        let duplex_supported = provider.supports_aead(HpkeAead::DuplexSpongeAead);
        let export_supported = provider.supports_aead(HpkeAead::Export);

        // Export should always be supported
        assert!(export_supported);

        // Others depend on features
        #[cfg(feature = "saturnin")]
        assert!(saturnin_supported);
        #[cfg(not(feature = "saturnin"))]
        assert!(!saturnin_supported);

        // SHAKE256 AEAD is now implemented in lib-q-aead
        // This test reflects the current reality - SHAKE256 AEAD is supported
        assert!(
            shake256_supported,
            "SHAKE256 AEAD should be supported after migration to lib-q-aead"
        );

        #[cfg(feature = "duplex-sponge-aead")]
        assert!(
            duplex_supported,
            "Duplex-sponge AEAD should be supported when feature is enabled"
        );
        #[cfg(not(feature = "duplex-sponge-aead"))]
        assert!(!duplex_supported);
    }
}