[ML-KEM] Match mainline APIs

libcrux-iot/ml-kem/Cargo.toml

@@ -31,11 +31,18 @@ libcrux-iot-sha3 = { path = "../sha3" , features = [
 ] }
 hax-lib.workspace = true
 libcrux-secrets.workspace = true
+libcrux-traits = { version = "0.0.4", features = ["generic-tests"]}
+tls_codec = { version = "0.4.2", features = [
+    "derive",
+], default-features = false, optional = true }
 
 [features]
 # By default all variants and std are enabled.
 default = ["mlkem512", "mlkem768", "mlkem1024"]
 
+# Serialization & Deserialization using tls_codec
+codec = ["dep:tls_codec"]
+
 # Features for the different key sizes of ML-KEM
 mlkem512 = []
 mlkem768 = []
@@ -48,9 +55,18 @@ kyber = []
 rand = ["dep:rand"]
 
 # std support
-std = []
+std = ["alloc", "rand/std", "tls_codec/std"]
+alloc = []
+
+
+check-secret-independence = [
+                          "libcrux-secrets/check-secret-independence",
+                          "libcrux-iot-sha3/check-secret-independence",
+                          "libcrux-traits/check-secret-independence"
+                          ]
 
-check-secret-independence = ["libcrux-secrets/check-secret-independence", "libcrux-iot-sha3/check-secret-independence"]
+# PQCP Common APIs
+pqcp = []
 
 [dev-dependencies]
 rand = { version = "0.9" }

libcrux-iot/ml-kem/hax.py

@@ -61,10 +61,6 @@ def __call__(self, parser, args, values, option_string=None) -> None:
         cargo_hax_into = [
             "cargo",
             "hax",
-            "-C",
-            "--features",
-            "std",
-            ";",
             "into",
             "-i",
             include_str,

libcrux-iot/ml-kem/src/lib.rs

@@ -72,7 +72,7 @@
 // Enable doc cfg feature for doc builds. They use nightly.
 #![cfg_attr(doc_cfg, feature(doc_cfg))]
 
-#[cfg(feature = "std")]
+#[cfg(any(feature = "std", hax))]
 extern crate std;
 
 /// Feature gating helper macros
@@ -108,6 +108,9 @@ mod utils;
 mod variant;
 mod vector;
 
+#[cfg(feature = "pqcp")]
+pub(crate) mod pqcp;
+
 #[cfg(feature = "mlkem512")]
 #[cfg_attr(docsrs, doc(cfg(feature = "mlkem512")))]
 pub mod mlkem512;
@@ -133,6 +136,7 @@ cfg_kyber! {
     pub mod kyber512 {
         //! Kyber 512 (NIST PQC Round 3)
         cfg_no_eurydice! {
+            pub use crate::mlkem512::kyber::Kyber512;
             pub use crate::mlkem512::kyber::generate_key_pair;
             pub use crate::mlkem512::kyber::decapsulate;
             pub use crate::mlkem512::kyber::encapsulate;
@@ -146,6 +150,7 @@ cfg_kyber! {
     pub mod kyber768 {
         //! Kyber 768 (NIST PQC Round 3)
         cfg_no_eurydice! {
+            pub use crate::mlkem768::kyber::Kyber768;
             pub use crate::mlkem768::kyber::generate_key_pair;
             pub use crate::mlkem768::kyber::decapsulate;
             pub use crate::mlkem768::kyber::encapsulate;
@@ -159,6 +164,7 @@ cfg_kyber! {
     pub mod kyber1024 {
         //! Kyber 1024 (NIST PQC Round 3)
         cfg_no_eurydice! {
+            pub use crate::mlkem1024::kyber::Kyber1024;
             pub use crate::mlkem1024::kyber::generate_key_pair;
             pub use crate::mlkem1024::kyber::decapsulate;
             pub use crate::mlkem1024::kyber::encapsulate;
@@ -167,3 +173,69 @@ cfg_kyber! {
         }
     }
 }
+
+#[cfg(not(hax))]
+macro_rules! impl_kem_trait {
+    ($variant:ty, $pk:ty, $sk:ty, $ct:ty) => {
+        impl
+            libcrux_traits::kem::arrayref::Kem<
+                CPA_PKE_PUBLIC_KEY_SIZE,
+                SECRET_KEY_SIZE,
+                CPA_PKE_CIPHERTEXT_SIZE,
+                SHARED_SECRET_SIZE,
+                KEY_GENERATION_SEED_SIZE,
+                SHARED_SECRET_SIZE,
+            > for $variant
+        {
+            fn keygen(
+                ek: &mut [u8; CPA_PKE_PUBLIC_KEY_SIZE],
+                dk: &mut [libcrux_secrets::U8; SECRET_KEY_SIZE],
+                rand: &[libcrux_secrets::U8; KEY_GENERATION_SEED_SIZE],
+            ) -> Result<(), libcrux_traits::kem::owned::KeyGenError> {
+                let key_pair = generate_key_pair(*rand);
+                ek.copy_from_slice(key_pair.pk());
+                dk.copy_from_slice(key_pair.sk());
+
+                Ok(())
+            }
+
+            fn encaps(
+                ct: &mut [u8; CPA_PKE_CIPHERTEXT_SIZE],
+                ss: &mut [libcrux_secrets::U8; SHARED_SECRET_SIZE],
+                ek: &[u8; CPA_PKE_PUBLIC_KEY_SIZE],
+                rand: &[libcrux_secrets::U8; SHARED_SECRET_SIZE],
+            ) -> Result<(), libcrux_traits::kem::owned::EncapsError> {
+                let public_key: $pk = ek.into();
+
+                let (ct_, ss_) = encapsulate(&public_key, *rand);
+                ct.copy_from_slice(ct_.as_slice());
+                ss.copy_from_slice(ss_.as_slice());
+
+                Ok(())
+            }
+
+            fn decaps(
+                ss: &mut [libcrux_secrets::U8; SHARED_SECRET_SIZE],
+                ct: &[u8; CPA_PKE_CIPHERTEXT_SIZE],
+                dk: &[libcrux_secrets::U8; SECRET_KEY_SIZE],
+            ) -> Result<(), libcrux_traits::kem::owned::DecapsError> {
+                let secret_key: $sk = dk.into();
+                let ciphertext: $ct = ct.into();
+
+                let ss_ = decapsulate(&secret_key, &ciphertext);
+
+                ss.copy_from_slice(ss_.as_slice());
+
+                Ok(())
+            }
+        }
+
+    libcrux_traits::kem::slice::impl_trait!($variant =>
+        CPA_PKE_PUBLIC_KEY_SIZE, SECRET_KEY_SIZE,
+        CPA_PKE_CIPHERTEXT_SIZE, SHARED_SECRET_SIZE,
+        KEY_GENERATION_SEED_SIZE, SHARED_SECRET_SIZE);
+    };
+}
+
+#[cfg(not(hax))]
+use impl_kem_trait;

libcrux-iot/ml-kem/src/mlkem1024.rs

@@ -29,6 +29,25 @@ const PRF_OUTPUT_SIZE2: usize = ETA2_RANDOMNESS_SIZE * RANK;
 
 const IMPLICIT_REJECTION_HASH_INPUT_SIZE: usize = SHARED_SECRET_SIZE + CPA_PKE_CIPHERTEXT_SIZE;
 
+/// The ML-KEM 1024 algorithms
+pub struct MlKem1024;
+
+#[cfg(not(any(hax, eurydice)))]
+crate::impl_kem_trait!(
+    MlKem1024,
+    MlKem1024PublicKey,
+    MlKem1024PrivateKey,
+    MlKem1024Ciphertext
+);
+
+// Provide the (packed) PQCP APIs
+#[cfg(feature = "pqcp")]
+crate::pqcp::pqcp_api!(
+    "use libcrux_iot_ml_kem::mlkem1024::pqcp::*;",
+    MlKem1024,
+    " 1024 "
+);
+
 /// An ML-KEM 1024 Ciphertext
 pub type MlKem1024Ciphertext = MlKemCiphertext<CPA_PKE_CIPHERTEXT_SIZE>;
 /// An ML-KEM 1024 Private key
@@ -410,6 +429,16 @@ pub(crate) mod kyber {
 
     use super::*;
 
+    /// The Kyber 1024 algorithms
+    pub struct Kyber1024;
+
+    crate::impl_kem_trait!(
+        Kyber1024,
+        MlKem1024PublicKey,
+        MlKem1024PrivateKey,
+        MlKem1024Ciphertext
+    );
+
     /// Generate Kyber 1024 Key Pair
     ///
     /// Generate an ML-KEM key pair. The input is a byte array of size

libcrux-iot/ml-kem/src/mlkem512.rs

@@ -31,6 +31,25 @@ const PRF_OUTPUT_SIZE2: usize = ETA2_RANDOMNESS_SIZE * RANK;
 
 const IMPLICIT_REJECTION_HASH_INPUT_SIZE: usize = SHARED_SECRET_SIZE + CPA_PKE_CIPHERTEXT_SIZE;
 
+/// The ML-KEM 512 algorithms
+pub struct MlKem512;
+
+#[cfg(not(any(hax, eurydice)))]
+crate::impl_kem_trait!(
+    MlKem512,
+    MlKem512PublicKey,
+    MlKem512PrivateKey,
+    MlKem512Ciphertext
+);
+
+// Provide the (packed) PQCP APIs
+#[cfg(feature = "pqcp")]
+crate::pqcp::pqcp_api!(
+    "use libcrux_iot_ml_kem::mlkem512::pqcp::*;",
+    MlKem512,
+    " 512 "
+);
+
 /// An ML-KEM 512 Ciphertext
 pub type MlKem512Ciphertext = MlKemCiphertext<CPA_PKE_CIPHERTEXT_SIZE>;
 /// An ML-KEM 512 Private key
@@ -409,6 +428,16 @@ pub mod rand {
 pub(crate) mod kyber {
     use super::*;
 
+    /// The Kyber 512 algorithms
+    pub struct Kyber512;
+
+    crate::impl_kem_trait!(
+        Kyber512,
+        MlKem512PublicKey,
+        MlKem512PrivateKey,
+        MlKem512Ciphertext
+    );
+
     /// Generate Kyber 512 Key Pair
     ///
     /// The input is a byte array of size

libcrux-iot/ml-kem/src/mlkem768.rs

@@ -38,6 +38,25 @@ const PRF_OUTPUT_SIZE2: usize = ETA2_RANDOMNESS_SIZE * RANK;
 
 const IMPLICIT_REJECTION_HASH_INPUT_SIZE: usize = SHARED_SECRET_SIZE + CPA_PKE_CIPHERTEXT_SIZE;
 
+/// The ML-KEM 768 algorithms
+pub struct MlKem768;
+
+#[cfg(not(any(hax, eurydice)))]
+crate::impl_kem_trait!(
+    MlKem768,
+    MlKem768PublicKey,
+    MlKem768PrivateKey,
+    MlKem768Ciphertext
+);
+
+// Provide the (packed) PQCP APIs
+#[cfg(feature = "pqcp")]
+crate::pqcp::pqcp_api!(
+    "use libcrux_iot_ml_kem::mlkem768::pqcp::*;",
+    MlKem768,
+    " 768 "
+);
+
 /// An ML-KEM 768 Ciphertext
 pub type MlKem768Ciphertext = MlKemCiphertext<CPA_PKE_CIPHERTEXT_SIZE>;
 /// An ML-KEM 768 Private key
@@ -401,6 +420,16 @@ pub mod rand {
 pub(crate) mod kyber {
     use super::*;
 
+    /// The Kyber 768 algorithms
+    pub struct Kyber768;
+
+    crate::impl_kem_trait!(
+        Kyber768,
+        MlKem768PublicKey,
+        MlKem768PrivateKey,
+        MlKem768Ciphertext
+    );
+
     /// Generate Kyber 768 Key Pair
     ///
     /// Generate a Kyber key pair. The input is a byte array of size