qudag-crypto 0.5.1

//! ML-KEM (Kyber) implementation for post-quantum key encapsulation

// mod ml_kem;
// pub use ml_kem::MlKem768Impl as MlKem768;

use subtle::ConstantTimeEq;
use thiserror::Error;
use zeroize::{Zeroize, ZeroizeOnDrop};

/// Errors that can occur during KEM operations
#[derive(Debug, Error)]
pub enum KEMError {
    #[error("Key generation failed")]
    KeyGenerationError,
    #[error("Encapsulation failed")]
    EncapsulationError,
    #[error("Decapsulation failed")]
    DecapsulationError,
    #[error("Invalid length")]
    InvalidLength,
    #[error("Invalid key format")]
    InvalidKey,
    #[error("Internal error")]
    InternalError,
}

/// ML-KEM public key.
#[derive(Debug, Clone, Zeroize, ZeroizeOnDrop)]
pub struct PublicKey(Vec<u8>);

impl PublicKey {
    pub fn from_bytes(bytes: &[u8]) -> Result<Self, KEMError> {
        Ok(Self(bytes.to_vec()))
    }

    pub fn as_bytes(&self) -> &[u8] {
        &self.0
    }
}

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

impl PartialEq for PublicKey {
    fn eq(&self, other: &Self) -> bool {
        self.0.ct_eq(&other.0).into()
    }
}

impl Eq for PublicKey {}

/// ML-KEM secret key.
#[derive(Debug, Clone, Zeroize, ZeroizeOnDrop)]
pub struct SecretKey(Vec<u8>);

impl SecretKey {
    pub fn from_bytes(bytes: &[u8]) -> Result<Self, KEMError> {
        Ok(Self(bytes.to_vec()))
    }

    pub fn as_bytes(&self) -> &[u8] {
        &self.0
    }
}

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

impl PartialEq for SecretKey {
    fn eq(&self, other: &Self) -> bool {
        self.0.ct_eq(&other.0).into()
    }
}

impl Eq for SecretKey {}

/// ML-KEM ciphertext.
#[derive(Debug, Clone, Zeroize, ZeroizeOnDrop)]
pub struct Ciphertext(Vec<u8>);

impl Ciphertext {
    pub fn from_bytes(bytes: &[u8]) -> Result<Self, KEMError> {
        Ok(Self(bytes.to_vec()))
    }

    pub fn as_bytes(&self) -> &[u8] {
        &self.0
    }
}

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

impl PartialEq for Ciphertext {
    fn eq(&self, other: &Self) -> bool {
        self.0.ct_eq(&other.0).into()
    }
}

impl Eq for Ciphertext {}

/// ML-KEM shared secret.
#[derive(Debug, Clone, Zeroize, ZeroizeOnDrop)]
pub struct SharedSecret(Vec<u8>);

impl SharedSecret {
    pub fn from_bytes(bytes: &[u8]) -> Result<Self, KEMError> {
        Ok(Self(bytes.to_vec()))
    }

    pub fn as_bytes(&self) -> &[u8] {
        &self.0
    }
}

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

impl PartialEq for SharedSecret {
    fn eq(&self, other: &Self) -> bool {
        self.0.ct_eq(&other.0).into()
    }
}

impl Eq for SharedSecret {}

/// ML-KEM key encapsulation trait.
pub trait KeyEncapsulation {
    /// Generate a new key pair.
    fn keygen() -> Result<(PublicKey, SecretKey), KEMError>;

    /// Encapsulate a shared secret using a public key.
    fn encapsulate(public_key: &PublicKey) -> Result<(Ciphertext, SharedSecret), KEMError>;

    /// Decapsulate a shared secret using a secret key and ciphertext.
    fn decapsulate(
        secret_key: &SecretKey,
        ciphertext: &Ciphertext,
    ) -> Result<SharedSecret, KEMError>;
}

/// ML-KEM key pair
#[derive(Debug, ZeroizeOnDrop)]
pub struct KeyPair {
    pub public_key: Vec<u8>,
    pub secret_key: Vec<u8>,
}

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

impl KeyPair {
    /// Create a new key pair (placeholder implementation)
    pub fn new() -> Self {
        Self {
            public_key: vec![0u8; 32], // Placeholder
            secret_key: vec![0u8; 32], // Placeholder
        }
    }

    /// Get public key reference
    pub fn public_key(&self) -> &[u8] {
        &self.public_key
    }

    /// Get secret key reference  
    pub fn secret_key(&self) -> &[u8] {
        &self.secret_key
    }
}