diff --git a/cert/pq_suite.go b/cert/pq_suite.go index 7d41f250..67b8c8ae 100644 --- a/cert/pq_suite.go +++ b/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, diff --git a/cert/pq_suite_test.go b/cert/pq_suite_test.go index 3880aa6b..0ebf5130 100644 --- a/cert/pq_suite_test.go +++ b/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])) +}