From 1ca3d82c09065d5936382edfe39e907850145652 Mon Sep 17 00:00:00 2001
From: Tony Arcieri <bascule@gmail.com>
Date: Thu, 29 Jan 2026 10:52:13 -0700
Subject: [PATCH] ml-kem: add initial Wycheproof test vectors

Uses a similar strategy to `ml-dsa` (and `rsa`), cloning the Wycheproof
test vectors as a submodule and parsing the JSON files directly.

Adds the following as an initial integration:
- `mlkem_512_keygen_seed_test`
- `mlkem_768_keygen_seed_test`
- `mlkem_1024_keygen_seed_test`
---
 .github/workflows/ml-kem.yml |   4 ++
 .gitmodules                  |   3 +
 ml-kem/tests/wycheproof.rs   | 106 +++++++++++++++++++++++++++++++++++
 thirdparty/wycheproof        |   1 +
 4 files changed, 114 insertions(+)
 create mode 100644 .gitmodules
 create mode 100644 ml-kem/tests/wycheproof.rs
 create mode 160000 thirdparty/wycheproof

diff --git a/.github/workflows/ml-kem.yml b/.github/workflows/ml-kem.yml
index a71f229..2deeead 100644
--- a/.github/workflows/ml-kem.yml
+++ b/.github/workflows/ml-kem.yml
@@ -69,6 +69,8 @@ jobs:
           - stable
     steps:
       - uses: actions/checkout@v6
+        with:
+          submodules: recursive
       - uses: RustCrypto/actions/cargo-cache@master
       - uses: dtolnay/rust-toolchain@master
         with:
@@ -91,6 +93,8 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v6
+        with:
+          submodules: recursive
       - uses: dtolnay/rust-toolchain@master
         with:
           toolchain: ${{ matrix.rust }}
diff --git a/.gitmodules b/.gitmodules
new file mode 100644
index 0000000..0b040ef
--- /dev/null
+++ b/.gitmodules
@@ -0,0 +1,3 @@
+[submodule "wycheproof"]
+	path = thirdparty/wycheproof
+	url = https://github.com/C2SP/wycheproof.git
diff --git a/ml-kem/tests/wycheproof.rs b/ml-kem/tests/wycheproof.rs
new file mode 100644
index 0000000..22ca823
--- /dev/null
+++ b/ml-kem/tests/wycheproof.rs
@@ -0,0 +1,106 @@
+//! Test against the Wycheproof test vectors.
+
+use ml_kem::{EncodedSizeUser, KemCore, MlKem512, MlKem768, MlKem1024, kem::KeyExport};
+use serde::Deserialize;
+use std::fs::File;
+
+#[derive(Deserialize, Debug)]
+struct TestFile {
+    algorithm: String,
+    schema: String,
+    #[serde(rename(deserialize = "testGroups"))]
+    groups: Vec<TestGroup>,
+}
+
+#[derive(Deserialize, Debug)]
+struct TestGroup {
+    #[allow(dead_code)]
+    #[serde(rename(deserialize = "type"))]
+    type_: String,
+    #[serde(default, rename(deserialize = "parameterSet"))]
+    parameter_set: String,
+    source: Source,
+    tests: Vec<Test>,
+}
+
+#[derive(Deserialize, Debug)]
+struct Source {
+    name: String,
+    version: String,
+}
+
+#[derive(Deserialize, Debug)]
+struct Test {
+    #[serde(rename(deserialize = "tcId"))]
+    id: usize,
+    comment: String,
+    #[serde(with = "hex::serde")]
+    seed: Vec<u8>,
+    #[serde(default, with = "hex::serde")]
+    ek: Vec<u8>,
+    #[serde(with = "hex::serde")]
+    dk: Vec<u8>,
+    result: ExpectedResult,
+}
+
+#[derive(Copy, Clone, Deserialize, Debug, PartialEq)]
+#[serde(rename_all = "lowercase")]
+enum ExpectedResult {
+    Valid,
+    Invalid,
+    Acceptable,
+}
+
+macro_rules! load_json_file {
+    ($json_file:expr) => {{
+        let path = format!("../thirdparty/wycheproof/testvectors_v1/{}", $json_file);
+        let data_file = File::open(&path)
+            .expect("failed to open data file (try running `git submodule update --init`)");
+
+        println!("Loading file: {path}");
+
+        let tests: TestFile = serde_json::from_reader(data_file).expect("invalid test JSON");
+        println!("{} ({})", tests.algorithm, tests.schema);
+        tests
+    }};
+}
+
+macro_rules! mlkem_keygen_seed_test {
+    ($name:ident, $json_file:expr, $kem:ident) => {
+        #[test]
+        fn $name() {
+            let tests = load_json_file!($json_file);
+
+            for group in tests.groups {
+                println!(
+                    "Parameter set: {} ({} v{})\n",
+                    &group.parameter_set, &group.source.name, &group.source.version
+                );
+
+                for test in &group.tests {
+                    println!("Test #{}: {} ({:?})", test.id, &test.comment, &test.result);
+
+                    let (dk, ek) = $kem::from_seed(test.seed.as_slice().try_into().unwrap());
+                    assert_eq!(test.dk.as_slice(), dk.to_encoded_bytes().as_slice());
+                    assert_eq!(test.ek.as_slice(), ek.to_bytes().as_slice());
+                }
+            }
+        }
+    };
+}
+
+mlkem_keygen_seed_test!(
+    mlkem_512_keygen_seed_test,
+    "mlkem_512_keygen_seed_test.json",
+    MlKem512
+);
+mlkem_keygen_seed_test!(
+    mlkem_768_keygen_seed_test,
+    "mlkem_768_keygen_seed_test.json",
+    MlKem768
+);
+mlkem_keygen_seed_test!(
+    mlkem_1024_keygen_seed_test,
+    "mlkem_1024_keygen_seed_test.json",
+    MlKem1024
+);
diff --git a/thirdparty/wycheproof b/thirdparty/wycheproof
new file mode 160000
index 0000000..ad6ac18
--- /dev/null
+++ b/thirdparty/wycheproof
@@ -0,0 +1 @@
+Subproject commit ad6ac186a990529963a1750e2aaa120f5f289cbc