crypto/mlkem/mlkemtest: add derandomized Encapsulate768/1024

Fixes #73627

Change-Id: I6a6a69649927e9b1cdff910832084fdc04ff5bc2
Reviewed-on: https://go-review.googlesource.com/c/go/+/703795
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Auto-Submit: Filippo Valsorda <filippo@golang.org>
Reviewed-by: Mark Freeman <markfreeman@google.com>
Reviewed-by: Junyang Shao <shaojunyang@google.com>
Reviewed-by: Daniel McCarney <daniel@binaryparadox.net>

api/next/73627.txt

@@ -0,0 +1,2 @@
+pkg crypto/mlkem/mlkemtest, func Encapsulate1024(*mlkem.EncapsulationKey1024, []uint8) ([]uint8, []uint8, error) #73627
+pkg crypto/mlkem/mlkemtest, func Encapsulate768(*mlkem.EncapsulationKey768, []uint8) ([]uint8, []uint8, error) #73627

doc/next/6-stdlib/99-minor/crypto/mlkem/mlkemtest/73627.md

@@ -0,0 +1,3 @@
+The new [crypto/mlkem/mlkemtest] package exposes the [Encapsulate768] and
+[Encapsulate1024] functions which implement derandomized ML-KEM encapsulation,
+for use with known-answer tests.

src/crypto/mlkem/mlkem.go

@@ -108,6 +108,9 @@ func (ek *EncapsulationKey768) Bytes() []byte {
 // encapsulation key, drawing random bytes from the default crypto/rand source.
 //
 // The shared key must be kept secret.
+//
+// For testing, derandomized encapsulation is provided by the
+// [crypto/mlkem/mlkemtest] package.
 func (ek *EncapsulationKey768) Encapsulate() (sharedKey, ciphertext []byte) {
 	return ek.key.Encapsulate()
 }
@@ -187,6 +190,9 @@ func (ek *EncapsulationKey1024) Bytes() []byte {
 // encapsulation key, drawing random bytes from the default crypto/rand source.
 //
 // The shared key must be kept secret.
+//
+// For testing, derandomized encapsulation is provided by the
+// [crypto/mlkem/mlkemtest] package.
 func (ek *EncapsulationKey1024) Encapsulate() (sharedKey, ciphertext []byte) {
 	return ek.key.Encapsulate()
 }

src/crypto/mlkem/mlkem_test.go

@@ -2,12 +2,14 @@
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.
 
-package mlkem
+package mlkem_test
 
 import (
 	"bytes"
 	"crypto/internal/fips140/mlkem"
 	"crypto/internal/fips140/sha3"
+	. "crypto/mlkem"
+	"crypto/mlkem/mlkemtest"
 	"crypto/rand"
 	"encoding/hex"
 	"flag"
@@ -176,7 +178,7 @@ func TestAccumulated(t *testing.T) {
 	s := sha3.NewShake128()
 	o := sha3.NewShake128()
 	seed := make([]byte, SeedSize)
-	var msg [32]byte
+	msg := make([]byte, 32)
 	ct1 := make([]byte, CiphertextSize768)
 
 	for i := 0; i < n; i++ {
@@ -188,8 +190,11 @@ func TestAccumulated(t *testing.T) {
 		ek := dk.EncapsulationKey()
 		o.Write(ek.Bytes())
 
-		s.Read(msg[:])
-		k, ct := ek.key.EncapsulateInternal(&msg)
+		s.Read(msg)
+		k, ct, err := mlkemtest.Encapsulate768(ek, msg)
+		if err != nil {
+			t.Fatal(err)
+		}
 		o.Write(ct)
 		o.Write(k)
 
@@ -231,8 +236,6 @@ func BenchmarkKeyGen(b *testing.B) {
 func BenchmarkEncaps(b *testing.B) {
 	seed := make([]byte, SeedSize)
 	rand.Read(seed)
-	var m [32]byte
-	rand.Read(m[:])
 	dk, err := NewDecapsulationKey768(seed)
 	if err != nil {
 		b.Fatal(err)
@@ -244,7 +247,7 @@ func BenchmarkEncaps(b *testing.B) {
 		if err != nil {
 			b.Fatal(err)
 		}
-		K, c := ek.key.EncapsulateInternal(&m)
+		K, c := ek.Encapsulate()
 		sink ^= c[0] ^ K[0]
 	}
 }

src/crypto/mlkem/mlkemtest/mlkemtest.go

@@ -0,0 +1,46 @@
+// Copyright 2025 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// Package mlkemtest provides testing functions for the ML-KEM algorithm.
+package mlkemtest
+
+import (
+	fips140mlkem "crypto/internal/fips140/mlkem"
+	"crypto/mlkem"
+	"errors"
+)
+
+// Encapsulate768 implements derandomized ML-KEM-768 encapsulation
+// (ML-KEM.Encaps_internal from FIPS 203) using the provided encapsulation key
+// ek and 32 bytes of randomness.
+//
+// It must only be used for known-answer tests.
+func Encapsulate768(ek *mlkem.EncapsulationKey768, random []byte) (sharedKey, ciphertext []byte, err error) {
+	if len(random) != 32 {
+		return nil, nil, errors.New("mlkemtest: Encapsulate768: random must be 32 bytes")
+	}
+	k, err := fips140mlkem.NewEncapsulationKey768(ek.Bytes())
+	if err != nil {
+		return nil, nil, errors.New("mlkemtest: Encapsulate768: failed to reconstruct key: " + err.Error())
+	}
+	sharedKey, ciphertext = k.EncapsulateInternal((*[32]byte)(random))
+	return sharedKey, ciphertext, nil
+}
+
+// Encapsulate1024 implements derandomized ML-KEM-1024 encapsulation
+// (ML-KEM.Encaps_internal from FIPS 203) using the provided encapsulation key
+// ek and 32 bytes of randomness.
+//
+// It must only be used for known-answer tests.
+func Encapsulate1024(ek *mlkem.EncapsulationKey1024, random []byte) (sharedKey, ciphertext []byte, err error) {
+	if len(random) != 32 {
+		return nil, nil, errors.New("mlkemtest: Encapsulate1024: random must be 32 bytes")
+	}
+	k, err := fips140mlkem.NewEncapsulationKey1024(ek.Bytes())
+	if err != nil {
+		return nil, nil, errors.New("mlkemtest: Encapsulate1024: failed to reconstruct key: " + err.Error())
+	}
+	sharedKey, ciphertext = k.EncapsulateInternal((*[32]byte)(random))
+	return sharedKey, ciphertext, nil
+}

src/go/build/deps_test.go

@@ -730,6 +730,9 @@ var depsRules = `
 	testing
 	< internal/testhash;
 
+	CRYPTO-MATH
+	< crypto/mlkem/mlkemtest;
+
 	CRYPTO-MATH, testing, internal/testenv, internal/testhash, encoding/json
 	< crypto/internal/cryptotest;