cert: Check an ML-KEM-1024 public key with the FIPS 203 modulus check alone, without expanding it

cert/pq_suite.go

@@ -88,11 +88,25 @@ var mlkem1024 = pqKEM{
 		return dk.EncapsulationKey().Bytes(), nil
 	},
 	checkPublicKey: func(pub []byte) bool {
-		_, err := mlkem.NewEncapsulationKey1024(pub)
-		return err == nil
+		return len(pub) == mlkem.EncapsulationKeySize1024 && mlkemModulusOK(pub[:len(pub)-32])
 	},
 }
 
+// mlkemModulusOK is the modulus check of FIPS 203 section 7.2, the one crypto/mlkem runs on an encapsulation key: t,
+// the key less its trailing 32 byte seed, packs 12 bit coefficients that must all be below q. crypto/mlkem would also
+// expand the seed into the ~10KB matrix of an encapsulation key, which a certificate check has no use for.
+func mlkemModulusOK(t []byte) bool {
+	const q = 3329
+	for i := 0; i+3 <= len(t); i += 3 {
+		d1 := uint16(t[i]) | uint16(t[i+1]&0x0f)<<8
+		d2 := uint16(t[i+1]>>4) | uint16(t[i+2])<<4
+		if d1 >= q || d2 >= q {
+			return false
+		}
+	}
+	return true
+}
+
 var mldsa87 = pqDSA{
 	name:           "MLDSA87",
 	publicKeySize:  mldsa.MLDSA87PublicKeySize,

cert/pq_suite_test.go

@@ -1,6 +1,9 @@
 package cert
 
 import (
+	"bytes"
+	"crypto/mlkem"
+	"crypto/rand"
 	"net/netip"
 	"strings"
 	"testing"
@@ -167,3 +170,42 @@ func TestPQSuite_Mismatch(t *testing.T) {
 		})
 	}
 }
+
+func TestMLKEMModulusOK(t *testing.T) {
+	stdlib := func(pub []byte) bool {
+		_, err := mlkem.NewEncapsulationKey1024(pub)
+		return err == nil
+	}
+	setCoeff := func(pub []byte, i int, v uint16) {
+		b := pub[i/2*3:]
+		if i%2 == 0 {
+			b[0] = byte(v)
+			b[1] = b[1]&0xf0 | byte(v>>8)
+		} else {
+			b[1] = b[1]&0x0f | byte(v<<4)
+			b[2] = byte(v >> 4)
+		}
+	}
+
+	for range 8 {
+		dk, err := mlkem.GenerateKey1024()
+		require.NoError(t, err)
+		pub := dk.EncapsulationKey().Bytes()
+		require.True(t, mlkem1024.checkPublicKey(pub))
+		for _, i := range []int{0, 1, 2, 3, 511, 512, 1022, 1023} {
+			for _, v := range []uint16{0, 3328, 3329, 4095} {
+				p := bytes.Clone(pub)
+				setCoeff(p, i, v)
+				assert.Equal(t, stdlib(p), mlkem1024.checkPublicKey(p), "coefficient %d = %d", i, v)
+				assert.Equal(t, v < 3329, mlkem1024.checkPublicKey(p), "coefficient %d = %d", i, v)
+			}
+		}
+	}
+
+	pub := make([]byte, mlkem.EncapsulationKeySize1024)
+	for range 64 {
+		_, _ = rand.Read(pub)
+		assert.Equal(t, stdlib(pub), mlkem1024.checkPublicKey(pub))
+	}
+	assert.False(t, mlkem1024.checkPublicKey(pub[:len(pub)-1]))
+}