ml-kem: encapsulation keys are not validated

[karalabe]

I've found a discrepancy between this implementation of ML-KEM and the Go one. There's a repro below for a case which is rejected by Go but accepted by Rust. Seems the issue is that public keys loaded from bytes are not validated, but they could be invalid.

//! Reproduction case: ml-kem accepts invalid encapsulation keys
//!
//! ML-KEM encapsulation keys contain polynomial coefficients encoded as 12-bit
//! values. Per FIPS 203, each coefficient must be in the range [0, 3329) where
//! 3329 is the modulus q. The ml-kem crate does not validate this, accepting
//! arbitrary bytes as valid encapsulation keys.
//!
//! Go's crypto/mlkem implementation validates coefficients and rejects invalid keys.

use ml_kem::{EncodedSizeUser, MlKem768Params};
use ml_kem::kem::EncapsulationKey;

fn main() {
    // Create an invalid key: all bytes set to 0xFF
    // When decoded as 12-bit coefficients, this produces values of 0xFFF = 4095 > 3329
    let invalid_key = [0xFFu8; 1184];

    // ml-kem accepts this without error
    let ek = EncapsulationKey::<MlKem768Params>::from_bytes(&invalid_key.into());

    println!("ml-kem accepted invalid key with coefficients >= 3329");
    println!("Encapsulation key bytes: {:?}...", &ek.as_bytes()[..16]);

    // Demonstrate the issue: the first 3 bytes [0xFF, 0xFF, 0xFF] decode as:
    // - coeff1 = 0xFF | ((0xFF & 0x0F) << 8) = 0xFF | 0xF00 = 0xFFF = 4095
    // - coeff2 = (0xFF >> 4) | (0xFF << 4) = 0x0F | 0xFF0 = 0xFFF = 4095
    // Both exceed q=3329, making this an invalid ML-KEM key

    let bytes: [u8; 3] = [0xFF, 0xFF, 0xFF];
    let coeff1 = u16::from(bytes[0]) | ((u16::from(bytes[1]) & 0x0F) << 8);
    let coeff2 = (u16::from(bytes[1]) >> 4) | (u16::from(bytes[2]) << 4);

    println!("\nDecoded coefficients from [0xFF, 0xFF, 0xFF]:");
    println!("  coeff1 = {} (max valid: 3328)", coeff1);
    println!("  coeff2 = {} (max valid: 3328)", coeff2);
    println!("\nBoth exceed q=3329, but ml-kem accepts this key.");
}

ml-kem accepted invalid key with coefficients >= 3329
Encapsulation key bytes: [254, 226, 47, 254, 226, 47, 254, 226, 47, 254, 226, 47, 254, 226, 47, 254]...

Decoded coefficients from [0xFF, 0xFF, 0xFF]:
  coeff1 = 4095 (max valid: 3328)
  coeff2 = 4095 (max valid: 3328)

[karalabe]

I think this is the check Go does: https://github.com/golang/go/blob/master/src/crypto/internal/fips140/mlkem/field.go#L164-L181

[tarcieri]

I made decoding fallible in #178, but we still need to implement the same checks as Go.

Looking at the current implementation perhaps PkeParams::decode_u12 could also be made fallible and perform the check. That could also be potentially used to implement expanded private key validation if desired.

cc @bifurcation

[tarcieri]

Getting into it a bit deeper I came across byte_decode (which implements "Algorithm 5 ByteDecode_d(F)") and noticed this in particular:

https://github.com/RustCrypto/KEMs/blob/7a029b2/ml-kem/src/encode.rs#L60

I'm wondering if the check should be inserted there, or if there are valid use cases where it should perform a modular reduction. If so, it should probably be a constant-time one (unfortunately that seems like another case that isn't triggering clippy::integer_division_remainder_used) It seems like since it's used for parsing expanded decapsulation keys it would ideally be constant-time for those, but if we do get rid of that support and it's only used for encapsulation keys it wouldn't need to be.

I guess I'll have to take a look at what the spec says about that particular algorithm.

Edit: I do see a modular reduction specified in FIPS 203, so we can't simply remove it:

Edit again: I now notice the algorithm numbers in encode.rs seem off-by-one, possibly an artifact of the FIPS 203 draft -> final. The image above shows Algorithm 6, but the code comment shows Algorithm 5 (which is ByteEncode_d(F), currently doc'd as Algorithm 4)

[tarcieri]

Aha, I think #179 should do the trick