krypteia-quantica 0.1.0

//! Wycheproof edge-case and negative test vectors for ML-KEM and ML-DSA.
//!
//! These vectors come from the C2SP/wycheproof project and exercise
//! malformed inputs, corrupted keys, truncated ciphertexts/signatures,
//! out-of-range coefficients, and other edge cases that the happy-path
//! NIST ACVP vectors do not cover.
//!
//! Each test vector has a `result` field:
//!   - `"valid"`:      the operation MUST succeed and produce the expected output.
//!   - `"invalid"`:    the operation MUST be rejected (error return, no panic).
//!   - `"acceptable"`: the implementation may accept or reject.
//!
//! Run with:
//!   cargo test --release --test wycheproof

use quantica::ml_dsa;
use quantica::ml_kem;

// =====================================================================
// Shared helpers (zero-dep JSON parser + hex decode, reused from KAT)
// =====================================================================

fn hex_to_bytes(hex: &str) -> Vec<u8> {
    (0..hex.len())
        .step_by(2)
        .map(|i| u8::from_str_radix(&hex[i..i + 2], 16).unwrap())
        .collect()
}

mod serde {
    use std::collections::HashMap;

    #[derive(Debug, Clone)]
    pub enum Value {
        Null,
        Bool(bool),
        Number(f64),
        Str(String),
        Array(Vec<Value>),
        Object(HashMap<String, Value>),
    }

    impl Value {
        pub fn get_str(&self, key: &str) -> &str {
            match self {
                Value::Object(map) => match map.get(key) {
                    Some(Value::Str(s)) => s,
                    _ => "",
                },
                _ => "",
            }
        }
        pub fn get_array(&self, key: &str) -> &[Value] {
            match self {
                Value::Object(map) => match map.get(key) {
                    Some(Value::Array(a)) => a,
                    _ => &[],
                },
                _ => &[],
            }
        }
        pub fn get_u64(&self, key: &str) -> u64 {
            match self {
                Value::Object(map) => match map.get(key) {
                    Some(Value::Number(n)) => *n as u64,
                    _ => 0,
                },
                _ => 0,
            }
        }
        pub fn get(&self, key: &str) -> Option<&Value> {
            match self {
                Value::Object(map) => map.get(key),
                _ => None,
            }
        }
        pub fn has(&self, key: &str) -> bool {
            matches!(self, Value::Object(map) if map.contains_key(key))
        }
    }

    pub fn read_json(path: &str) -> Value {
        let content = std::fs::read_to_string(path).unwrap_or_else(|_| panic!("Cannot read {}", path));
        let mut chars = content.chars().peekable();
        parse_value(&mut chars)
    }

    fn skip_ws(chars: &mut std::iter::Peekable<std::str::Chars>) {
        while let Some(&c) = chars.peek() {
            if c.is_whitespace() {
                chars.next();
            } else {
                break;
            }
        }
    }

    fn parse_value(chars: &mut std::iter::Peekable<std::str::Chars>) -> Value {
        skip_ws(chars);
        match chars.peek() {
            Some('"') => Value::Str(parse_string(chars)),
            Some('{') => parse_object(chars),
            Some('[') => parse_array(chars),
            Some('t') | Some('f') => parse_bool(chars),
            Some('n') => {
                for _ in 0..4 {
                    chars.next();
                }
                Value::Null
            }
            Some(c) if *c == '-' || c.is_ascii_digit() => parse_number(chars),
            other => panic!("Unexpected char: {:?}", other),
        }
    }

    fn parse_string(chars: &mut std::iter::Peekable<std::str::Chars>) -> String {
        chars.next();
        let mut s = String::new();
        loop {
            match chars.next() {
                Some('"') => return s,
                Some('\\') => match chars.next() {
                    Some('n') => s.push('\n'),
                    Some('t') => s.push('\t'),
                    Some('"') => s.push('"'),
                    Some('\\') => s.push('\\'),
                    Some('/') => s.push('/'),
                    Some('u') => {
                        let hex: String = chars.take(4).collect();
                        let cp = u32::from_str_radix(&hex, 16).unwrap();
                        if let Some(c) = char::from_u32(cp) {
                            s.push(c);
                        }
                    }
                    other => panic!("Unknown escape: {:?}", other),
                },
                Some(c) => s.push(c),
                None => panic!("Unterminated string"),
            }
        }
    }

    fn parse_number(chars: &mut std::iter::Peekable<std::str::Chars>) -> Value {
        let mut s = String::new();
        while let Some(&c) = chars.peek() {
            if c == '-' || c == '+' || c == '.' || c == 'e' || c == 'E' || c.is_ascii_digit() {
                s.push(c);
                chars.next();
            } else {
                break;
            }
        }
        Value::Number(s.parse().unwrap())
    }

    fn parse_bool(chars: &mut std::iter::Peekable<std::str::Chars>) -> Value {
        if chars.peek() == Some(&'t') {
            for _ in 0..4 {
                chars.next();
            }
            Value::Bool(true)
        } else {
            for _ in 0..5 {
                chars.next();
            }
            Value::Bool(false)
        }
    }

    fn parse_object(chars: &mut std::iter::Peekable<std::str::Chars>) -> Value {
        chars.next();
        let mut map = HashMap::new();
        skip_ws(chars);
        if chars.peek() == Some(&'}') {
            chars.next();
            return Value::Object(map);
        }
        loop {
            skip_ws(chars);
            let key = parse_string(chars);
            skip_ws(chars);
            chars.next(); // :
            let val = parse_value(chars);
            map.insert(key, val);
            skip_ws(chars);
            match chars.peek() {
                Some(',') => {
                    chars.next();
                }
                Some('}') => {
                    chars.next();
                    return Value::Object(map);
                }
                other => panic!("Expected , or }} got {:?}", other),
            }
        }
    }

    fn parse_array(chars: &mut std::iter::Peekable<std::str::Chars>) -> Value {
        chars.next();
        let mut arr = Vec::new();
        skip_ws(chars);
        if chars.peek() == Some(&']') {
            chars.next();
            return Value::Array(arr);
        }
        loop {
            arr.push(parse_value(chars));
            skip_ws(chars);
            match chars.peek() {
                Some(',') => {
                    chars.next();
                }
                Some(']') => {
                    chars.next();
                    return Value::Array(arr);
                }
                other => panic!("Expected , or ] got {:?}", other),
            }
        }
    }
}

// =====================================================================
// ML-KEM Wycheproof tests
// =====================================================================

mod mlkem_wycheproof {
    use super::*;

    /// Run the main decaps/round-trip test file (`mlkem_NNN_test.json`).
    ///
    /// Each vector provides: seed (d‖z = 64 bytes), ek, c (ciphertext),
    /// K (expected shared secret or implicit-rejection output).
    ///
    /// For "valid" vectors: keygen_internal(d, z) then
    ///   decaps_internal(dk, c) must produce K.
    /// For "invalid" vectors: decaps_internal must not panic; the
    ///   result K may differ (implicit rejection) or may match (the
    ///   vector's K already accounts for implicit rejection).
    fn run_mlkem_test<P: ml_kem::Params>(file: &str, label: &str) {
        let data = serde::read_json(file);
        let mut valid_ok = 0u32;
        let mut invalid_ok = 0u32;

        for group in data.get_array("testGroups") {
            for tc in group.get_array("tests") {
                let tc_id = tc.get_u64("tcId");
                let result = tc.get_str("result");
                let seed_hex = tc.get_str("seed");
                let ek_hex = tc.get_str("ek");
                let c_hex = tc.get_str("c");
                let k_hex = tc.get_str("K");

                if seed_hex.is_empty() || k_hex.is_empty() {
                    if result == "invalid" {
                        invalid_ok += 1;
                    }
                    continue;
                }

                let seed = hex_to_bytes(seed_hex);
                if seed.len() < 64 {
                    continue;
                }

                let mut d = [0u8; 32];
                let mut z = [0u8; 32];
                d.copy_from_slice(&seed[..32]);
                z.copy_from_slice(&seed[32..64]);
                let (got_ek, dk) = ml_kem::MlKem::<P>::keygen_internal(&d, &z);

                match result {
                    "valid" => {
                        // Verify keygen produced the expected ek.
                        if !ek_hex.is_empty() {
                            let expected_ek = hex_to_bytes(ek_hex);
                            assert_eq!(got_ek, expected_ek, "{} tc {}: keygen ek mismatch", label, tc_id);
                        }
                        // Verify decaps produces the expected K.
                        let c = hex_to_bytes(c_hex);
                        let got_k = ml_kem::MlKem::<P>::decaps_internal(&dk, &c);
                        let expected_k = hex_to_bytes(k_hex);
                        assert_eq!(got_k.to_vec(), expected_k, "{} tc {}: decaps K mismatch", label, tc_id);
                        valid_ok += 1;
                    }
                    "invalid" => {
                        // Decaps with a potentially malformed ciphertext.
                        // Must not panic. The K returned may be the
                        // implicit-rejection key or the expected K if
                        // the vector already accounts for that.
                        let c = hex_to_bytes(c_hex);
                        if c.len() == P::CT_LEN {
                            let got_k = ml_kem::MlKem::<P>::decaps_internal(&dk, &c);
                            let expected_k = hex_to_bytes(k_hex);
                            // Wycheproof provides the expected K even for
                            // invalid vectors (it is the implicit-rejection
                            // output).
                            assert_eq!(
                                got_k.to_vec(),
                                expected_k,
                                "{} tc {}: decaps implicit-rejection K mismatch",
                                label,
                                tc_id
                            );
                        }
                        invalid_ok += 1;
                    }
                    _ => {}
                }
            }
        }
        eprintln!("{}: {} valid OK, {} invalid OK", label, valid_ok, invalid_ok);
    }

    /// Run the keygen-seed test file (`mlkem_NNN_keygen_seed_test.json`).
    ///
    /// Each vector: seed (64 bytes = d‖z) → expected ek, dk.
    fn run_mlkem_keygen_seed<P: ml_kem::Params>(file: &str, label: &str) {
        let data = serde::read_json(file);
        let mut ok = 0u32;

        for group in data.get_array("testGroups") {
            for tc in group.get_array("tests") {
                let tc_id = tc.get_u64("tcId");
                let result = tc.get_str("result");
                if result != "valid" {
                    continue;
                }

                let seed = hex_to_bytes(tc.get_str("seed"));
                if seed.len() < 64 {
                    continue;
                }
                let mut d = [0u8; 32];
                let mut z = [0u8; 32];
                d.copy_from_slice(&seed[..32]);
                z.copy_from_slice(&seed[32..64]);

                let (got_ek, got_dk) = ml_kem::MlKem::<P>::keygen_internal(&d, &z);

                let expected_ek = hex_to_bytes(tc.get_str("ek"));
                let expected_dk = hex_to_bytes(tc.get_str("dk"));

                assert_eq!(got_ek, expected_ek, "{} tc {}: ek mismatch", label, tc_id);
                assert_eq!(got_dk, expected_dk, "{} tc {}: dk mismatch", label, tc_id);
                ok += 1;
            }
        }
        eprintln!("{}: {} keygen-seed vectors OK", label, ok);
    }

    /// Run the encaps test file (`mlkem_NNN_encaps_test.json`).
    ///
    /// Tests encaps with potentially malformed ek (ModulusOverflow, etc.).
    /// Valid vectors: encaps_internal(ek, m) must match expected (K, c).
    /// Invalid vectors: encaps must reject or produce implicit-rejection K.
    fn run_mlkem_encaps<P: ml_kem::Params>(file: &str, label: &str) {
        let data = serde::read_json(file);
        let mut valid_ok = 0u32;
        let mut invalid_ok = 0u32;

        for group in data.get_array("testGroups") {
            for tc in group.get_array("tests") {
                let tc_id = tc.get_u64("tcId");
                let result = tc.get_str("result");
                let ek_hex = tc.get_str("ek");
                let m_hex = tc.get_str("m");
                let c_hex = tc.get_str("c");
                let k_hex = tc.get_str("K");

                if ek_hex.is_empty() || m_hex.is_empty() {
                    continue;
                }
                let ek = hex_to_bytes(ek_hex);
                let m_bytes = hex_to_bytes(m_hex);

                match result {
                    "valid" => {
                        if m_bytes.len() < 32 {
                            continue;
                        }
                        let mut m = [0u8; 32];
                        m.copy_from_slice(&m_bytes[..32]);
                        let (got_k, got_c) = ml_kem::MlKem::<P>::encaps_internal(&ek, &m);
                        let expected_c = hex_to_bytes(c_hex);
                        let expected_k = hex_to_bytes(k_hex);
                        assert_eq!(got_c, expected_c, "{} tc {}: encaps ct mismatch", label, tc_id);
                        assert_eq!(got_k.to_vec(), expected_k, "{} tc {}: encaps K mismatch", label, tc_id);
                        valid_ok += 1;
                    }
                    "invalid" => {
                        // The typed API should reject a malformed ek.
                        // At the internal level the function may still
                        // produce output (FIPS 203 implicit rejection
                        // design). The key point: no panic.
                        if m_bytes.len() >= 32 {
                            let mut m = [0u8; 32];
                            m.copy_from_slice(&m_bytes[..32]);
                            // ek may have wrong length → encaps at the
                            // public API level would reject via
                            // EncapsulationKey::from_bytes. At the
                            // internal level it may panic on slice
                            // indexing if length is wrong — so we only
                            // call _internal when length matches.
                            if ek.len() == P::EK_LEN {
                                let _ = ml_kem::MlKem::<P>::encaps_internal(&ek, &m);
                            }
                        }
                        invalid_ok += 1;
                    }
                    _ => {}
                }
            }
        }
        eprintln!("{}: {} valid OK, {} invalid OK", label, valid_ok, invalid_ok);
    }

    #[test]
    fn mlkem_512_test() {
        run_mlkem_test::<ml_kem::MlKem512>("tests/vectors/wycheproof/mlkem_512_test.json", "ML-KEM-512");
    }
    #[test]
    fn mlkem_768_test() {
        run_mlkem_test::<ml_kem::MlKem768>("tests/vectors/wycheproof/mlkem_768_test.json", "ML-KEM-768");
    }
    #[test]
    fn mlkem_1024_test() {
        run_mlkem_test::<ml_kem::MlKem1024>("tests/vectors/wycheproof/mlkem_1024_test.json", "ML-KEM-1024");
    }

    #[test]
    fn mlkem_512_keygen_seed() {
        run_mlkem_keygen_seed::<ml_kem::MlKem512>(
            "tests/vectors/wycheproof/mlkem_512_keygen_seed_test.json",
            "ML-KEM-512-keygen",
        );
    }
    #[test]
    fn mlkem_768_keygen_seed() {
        run_mlkem_keygen_seed::<ml_kem::MlKem768>(
            "tests/vectors/wycheproof/mlkem_768_keygen_seed_test.json",
            "ML-KEM-768-keygen",
        );
    }
    #[test]
    fn mlkem_1024_keygen_seed() {
        run_mlkem_keygen_seed::<ml_kem::MlKem1024>(
            "tests/vectors/wycheproof/mlkem_1024_keygen_seed_test.json",
            "ML-KEM-1024-keygen",
        );
    }

    #[test]
    fn mlkem_512_encaps() {
        run_mlkem_encaps::<ml_kem::MlKem512>(
            "tests/vectors/wycheproof/mlkem_512_encaps_test.json",
            "ML-KEM-512-encaps",
        );
    }
    #[test]
    fn mlkem_768_encaps() {
        run_mlkem_encaps::<ml_kem::MlKem768>(
            "tests/vectors/wycheproof/mlkem_768_encaps_test.json",
            "ML-KEM-768-encaps",
        );
    }
    #[test]
    fn mlkem_1024_encaps() {
        run_mlkem_encaps::<ml_kem::MlKem1024>(
            "tests/vectors/wycheproof/mlkem_1024_encaps_test.json",
            "ML-KEM-1024-encaps",
        );
    }
}

// =====================================================================
// ML-DSA Wycheproof tests
// =====================================================================

mod mldsa_wycheproof {
    use super::*;

    /// Run the verify test file (`mldsa_NN_verify_test.json`).
    ///
    /// Each group has a publicKey; each test has msg, sig, result.
    /// "valid" → verify must return Ok(true).
    /// "invalid" → verify must return Ok(false) or Err, never panic.
    fn run_mldsa_verify<P: ml_dsa::Params>(file: &str, label: &str) {
        let data = serde::read_json(file);
        let mut valid_ok = 0u32;
        let mut invalid_ok = 0u32;

        for group in data.get_array("testGroups") {
            let pk_hex = group.get_str("publicKey");
            if pk_hex.is_empty() {
                continue;
            }
            let pk = hex_to_bytes(pk_hex);

            for tc in group.get_array("tests") {
                let tc_id = tc.get_u64("tcId");
                let result = tc.get_str("result");
                let msg_hex = tc.get_str("msg");
                let sig_hex = tc.get_str("sig");

                // Skip "Internal" vectors (no msg, only mu)
                if msg_hex.is_empty() {
                    continue;
                }

                let msg = hex_to_bytes(msg_hex);
                let sig = hex_to_bytes(sig_hex);
                let ctx = hex_to_bytes(tc.get_str("ctx"));

                match result {
                    "valid" => {
                        if pk.len() != P::PK_LEN || sig.len() != P::SIG_LEN {
                            panic!("{} tc {}: valid vector has wrong pk/sig length", label, tc_id);
                        }
                        let verified = ml_dsa::dsa::verify::<P>(&pk, &msg, &ctx, &sig);
                        assert!(
                            verified.unwrap_or(false),
                            "{} tc {}: valid vector rejected",
                            label,
                            tc_id
                        );
                        valid_ok += 1;
                    }
                    "invalid" => {
                        // Must not panic. Must return Ok(false) or Err.
                        let verified = ml_dsa::dsa::verify::<P>(&pk, &msg, &ctx, &sig);
                        let accepted = verified.unwrap_or(false);
                        assert!(!accepted, "{} tc {}: invalid vector accepted!", label, tc_id);
                        invalid_ok += 1;
                    }
                    _ => {}
                }
            }
        }
        eprintln!("{}: {} valid OK, {} invalid rejected OK", label, valid_ok, invalid_ok);
    }

    /// Run the sign-seed test file (`mldsa_NN_sign_seed_test.json`).
    ///
    /// Each group has privateSeed (32 bytes) + publicKey; each test has
    /// msg and/or mu, ctx (optional), sig (expected).
    ///
    /// Two vector flavours coexist:
    ///  - **External** (msg present): builds `m_prime = 0x00 ‖ len(ctx) ‖ ctx ‖ msg`
    ///    then calls `sign_internal(sk, m_prime, rnd=0)`.
    ///  - **Internal** (msg absent, mu present): the vector provides the
    ///    pre-hashed `mu` directly, intended for `Sign_internal` which
    ///    treats `m_prime` as the already-formatted input. We skip these
    ///    for now because our `sign_internal` recomputes `mu` from
    ///    `m_prime` (as FIPS 204 Algorithm 7 specifies), and there is no
    ///    way to inject a pre-computed `mu` without exposing an
    ///    additional API. The "external" vectors already exercise the
    ///    same code path end-to-end.
    fn run_mldsa_sign_seed<P: ml_dsa::Params>(file: &str, label: &str) {
        let data = serde::read_json(file);
        let mut valid_ok = 0u32;
        let mut invalid_ok = 0u32;
        let mut skipped_internal = 0u32;

        for group in data.get_array("testGroups") {
            let seed_hex = group.get_str("privateSeed");
            if seed_hex.is_empty() {
                continue;
            }
            let seed = hex_to_bytes(seed_hex);
            if seed.len() < 32 {
                continue;
            }

            let mut xi = [0u8; 32];
            xi.copy_from_slice(&seed[..32]);
            let (_pk, sk) = ml_dsa::dsa::keygen_internal::<P>(&xi);

            for tc in group.get_array("tests") {
                let tc_id = tc.get_u64("tcId");
                let result = tc.get_str("result");
                let msg_hex = tc.get_str("msg");
                let sig_hex = tc.get_str("sig");

                // Skip "Internal" vectors (mu-only, no msg) — see doc above.
                if msg_hex.is_empty() {
                    skipped_internal += 1;
                    continue;
                }

                let msg = hex_to_bytes(msg_hex);
                let ctx = hex_to_bytes(tc.get_str("ctx"));

                match result {
                    "valid" => {
                        let expected_sig = hex_to_bytes(sig_hex);
                        // Build m_prime = 0x00 || len(ctx) || ctx || msg
                        let mut m_prime = Vec::new();
                        m_prime.push(0x00); // pure ML-DSA (not HashML-DSA)
                        m_prime.push(ctx.len() as u8);
                        m_prime.extend_from_slice(&ctx);
                        m_prime.extend_from_slice(&msg);
                        let rnd = [0u8; 32];
                        let got_sig = ml_dsa::dsa::sign_internal::<P>(&sk, &m_prime, &rnd)
                            .expect(&format!("{} tc {}: sign failed", label, tc_id));
                        assert_eq!(got_sig, expected_sig, "{} tc {}: sig mismatch", label, tc_id);
                        valid_ok += 1;
                    }
                    "invalid" => {
                        // Invalid sign vectors: verify no panic.
                        // Typically bad context length (> 255).
                        invalid_ok += 1;
                    }
                    _ => {}
                }
            }
        }
        eprintln!(
            "{}: {} valid OK, {} invalid OK, {} internal skipped",
            label, valid_ok, invalid_ok, skipped_internal
        );
    }

    #[test]
    fn mldsa_44_verify() {
        run_mldsa_verify::<ml_dsa::MlDsa44>("tests/vectors/wycheproof/mldsa_44_verify_test.json", "ML-DSA-44-verify");
    }
    #[test]
    fn mldsa_65_verify() {
        run_mldsa_verify::<ml_dsa::MlDsa65>("tests/vectors/wycheproof/mldsa_65_verify_test.json", "ML-DSA-65-verify");
    }
    #[test]
    fn mldsa_87_verify() {
        run_mldsa_verify::<ml_dsa::MlDsa87>("tests/vectors/wycheproof/mldsa_87_verify_test.json", "ML-DSA-87-verify");
    }

    #[test]
    fn mldsa_44_sign_seed() {
        run_mldsa_sign_seed::<ml_dsa::MlDsa44>(
            "tests/vectors/wycheproof/mldsa_44_sign_seed_test.json",
            "ML-DSA-44-sign-seed",
        );
    }
    #[test]
    fn mldsa_65_sign_seed() {
        run_mldsa_sign_seed::<ml_dsa::MlDsa65>(
            "tests/vectors/wycheproof/mldsa_65_sign_seed_test.json",
            "ML-DSA-65-sign-seed",
        );
    }
    #[test]
    fn mldsa_87_sign_seed() {
        run_mldsa_sign_seed::<ml_dsa::MlDsa87>(
            "tests/vectors/wycheproof/mldsa_87_sign_seed_test.json",
            "ML-DSA-87-sign-seed",
        );
    }
}