crypto/internal/fips140/mlkem: make CAST conditional

It taks north of 130µs on my machine, which is enough to be worth
shaving off at init time.

Change-Id: I6a6a696463de228bc3e7b9ca10c12ddbaabdf384
Reviewed-on: https://go-review.googlesource.com/c/go/+/707695
Auto-Submit: Filippo Valsorda <filippo@golang.org>
Reviewed-by: Daniel McCarney <daniel@binaryparadox.net>
Reviewed-by: Roland Shoemaker <roland@golang.org>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Carlos Amedee <carlos@golang.org>

4 files changed, 34 insertions, 10 deletions

diff --git a/src/crypto/internal/fips140/mlkem/cast.go b/src/crypto/internal/fips140/mlkem/cast.go
index a432d1fdab..ea089c1b76 100644
--- a/src/crypto/internal/fips140/mlkem/cast.go
+++ b/src/crypto/internal/fips140/mlkem/cast.go
@@ -9,9 +9,10 @@ import (
 	"crypto/internal/fips140"
 	_ "crypto/internal/fips140/check"
 	"errors"
+	"sync"
 )
 
-func init() {
+var fipsSelfTest = sync.OnceFunc(func() {
 	fips140.CAST("ML-KEM-768", func() error {
 		var d = &[32]byte{
 			0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08,
@@ -40,14 +41,12 @@ func init() {
 		dk := &DecapsulationKey768{}
 		kemKeyGen(dk, d, z)
 		ek := dk.EncapsulationKey()
-		Ke, c := ek.EncapsulateInternal(m)
-		Kd, err := dk.Decapsulate(c)
-		if err != nil {
-			return err
-		}
+		var cc [CiphertextSize768]byte
+		Ke, _ := kemEncaps(&cc, ek, m)
+		Kd := kemDecaps(dk, &cc)
 		if !bytes.Equal(Ke, K) || !bytes.Equal(Kd, K) {
 			return errors.New("unexpected result")
 		}
 		return nil
 	})
-}
+})
diff --git a/src/crypto/internal/fips140/mlkem/mlkem1024.go b/src/crypto/internal/fips140/mlkem/mlkem1024.go
index 1419cf20fa..edde161422 100644
--- a/src/crypto/internal/fips140/mlkem/mlkem1024.go
+++ b/src/crypto/internal/fips140/mlkem/mlkem1024.go
@@ -113,6 +113,7 @@ func GenerateKey1024() (*DecapsulationKey1024, error) {
 }
 
 func generateKey1024(dk *DecapsulationKey1024) (*DecapsulationKey1024, error) {
+	fipsSelfTest()
 	var d [32]byte
 	drbg.Read(d[:])
 	var z [32]byte
@@ -126,6 +127,7 @@ func generateKey1024(dk *DecapsulationKey1024) (*DecapsulationKey1024, error) {
 // GenerateKeyInternal1024 is a derandomized version of GenerateKey1024,
 // exclusively for use in tests.
 func GenerateKeyInternal1024(d, z *[32]byte) *DecapsulationKey1024 {
+	fipsSelfTest()
 	dk := &DecapsulationKey1024{}
 	kemKeyGen1024(dk, d, z)
 	return dk
@@ -278,6 +280,7 @@ func (ek *EncapsulationKey1024) Encapsulate() (sharedKey, ciphertext []byte) {
 }
 
 func (ek *EncapsulationKey1024) encapsulate(cc *[CiphertextSize1024]byte) (sharedKey, ciphertext []byte) {
+	fipsSelfTest()
 	var m [messageSize]byte
 	drbg.Read(m[:])
 	// Note that the modulus check (step 2 of the encapsulation key check from
@@ -289,6 +292,7 @@ func (ek *EncapsulationKey1024) encapsulate(cc *[CiphertextSize1024]byte) (share
 // EncapsulateInternal is a derandomized version of Encapsulate, exclusively for
 // use in tests.
 func (ek *EncapsulationKey1024) EncapsulateInternal(m *[32]byte) (sharedKey, ciphertext []byte) {
+	fipsSelfTest()
 	cc := &[CiphertextSize1024]byte{}
 	return kemEncaps1024(cc, ek, m)
 }
@@ -394,6 +398,7 @@ func pkeEncrypt1024(cc *[CiphertextSize1024]byte, ex *encryptionKey1024, m *[mes
 //
 // The shared key must be kept secret.
 func (dk *DecapsulationKey1024) Decapsulate(ciphertext []byte) (sharedKey []byte, err error) {
+	fipsSelfTest()
 	if len(ciphertext) != CiphertextSize1024 {
 		return nil, errors.New("mlkem: invalid ciphertext length")
 	}
diff --git a/src/crypto/internal/fips140/mlkem/mlkem768.go b/src/crypto/internal/fips140/mlkem/mlkem768.go
index 298660e4e9..088c2954de 100644
--- a/src/crypto/internal/fips140/mlkem/mlkem768.go
+++ b/src/crypto/internal/fips140/mlkem/mlkem768.go
@@ -172,6 +172,7 @@ func GenerateKey768() (*DecapsulationKey768, error) {
 }
 
 func generateKey(dk *DecapsulationKey768) (*DecapsulationKey768, error) {
+	fipsSelfTest()
 	var d [32]byte
 	drbg.Read(d[:])
 	var z [32]byte
@@ -185,6 +186,7 @@ func generateKey(dk *DecapsulationKey768) (*DecapsulationKey768, error) {
 // GenerateKeyInternal768 is a derandomized version of GenerateKey768,
 // exclusively for use in tests.
 func GenerateKeyInternal768(d, z *[32]byte) *DecapsulationKey768 {
+	fipsSelfTest()
 	dk := &DecapsulationKey768{}
 	kemKeyGen(dk, d, z)
 	return dk
@@ -337,6 +339,7 @@ func (ek *EncapsulationKey768) Encapsulate() (sharedKey, ciphertext []byte) {
 }
 
 func (ek *EncapsulationKey768) encapsulate(cc *[CiphertextSize768]byte) (sharedKey, ciphertext []byte) {
+	fipsSelfTest()
 	var m [messageSize]byte
 	drbg.Read(m[:])
 	// Note that the modulus check (step 2 of the encapsulation key check from
@@ -348,6 +351,7 @@ func (ek *EncapsulationKey768) encapsulate(cc *[CiphertextSize768]byte) (sharedK
 // EncapsulateInternal is a derandomized version of Encapsulate, exclusively for
 // use in tests.
 func (ek *EncapsulationKey768) EncapsulateInternal(m *[32]byte) (sharedKey, ciphertext []byte) {
+	fipsSelfTest()
 	cc := &[CiphertextSize768]byte{}
 	return kemEncaps(cc, ek, m)
 }
@@ -453,6 +457,7 @@ func pkeEncrypt(cc *[CiphertextSize768]byte, ex *encryptionKey, m *[messageSize]
 //
 // The shared key must be kept secret.
 func (dk *DecapsulationKey768) Decapsulate(ciphertext []byte) (sharedKey []byte, err error) {
+	fipsSelfTest()
 	if len(ciphertext) != CiphertextSize768 {
 		return nil, errors.New("mlkem: invalid ciphertext length")
 	}
diff --git a/src/crypto/internal/fips140test/cast_test.go b/src/crypto/internal/fips140test/cast_test.go
index b043a71f04..5bbc964b61 100644
--- a/src/crypto/internal/fips140test/cast_test.go
+++ b/src/crypto/internal/fips140test/cast_test.go
@@ -48,8 +48,8 @@ var allCASTs = []string{
 	"HKDF-SHA2-256",
 	"HMAC-SHA2-256",
 	"KAS-ECC-SSC P-256",
-	"ML-KEM PCT",
-	"ML-KEM PCT",
+	"ML-KEM PCT", // -768
+	"ML-KEM PCT", // -1024
 	"ML-KEM-768",
 	"PBKDF2",
 	"RSA sign and verify PCT",
@@ -104,29 +104,44 @@ func TestAllCASTs(t *testing.T) {
 
 // TestConditionals causes the conditional CASTs and PCTs to be invoked.
 func TestConditionals(t *testing.T) {
-	mlkem.GenerateKey768()
+	// ML-KEM PCT
+	kMLKEM, err := mlkem.GenerateKey768()
+	if err != nil {
+		t.Error(err)
+	} else {
+		// ML-KEM-768
+		kMLKEM.EncapsulationKey().Encapsulate()
+	}
+	// ECDH PCT
 	kDH, err := ecdh.GenerateKey(ecdh.P256(), rand.Reader)
 	if err != nil {
 		t.Error(err)
 	} else {
+		// KAS-ECC-SSC P-256
 		ecdh.ECDH(ecdh.P256(), kDH, kDH.PublicKey())
 	}
+	// ECDSA PCT
 	kDSA, err := ecdsa.GenerateKey(ecdsa.P256(), rand.Reader)
 	if err != nil {
 		t.Error(err)
 	} else {
+		// ECDSA P-256 SHA2-512 sign and verify
 		ecdsa.SignDeterministic(ecdsa.P256(), sha256.New, kDSA, make([]byte, 32))
 	}
+	// Ed25519 sign and verify PCT
 	k25519, err := ed25519.GenerateKey()
 	if err != nil {
 		t.Error(err)
 	} else {
+		// Ed25519 sign and verify
 		ed25519.Sign(k25519, make([]byte, 32))
 	}
+	// RSA sign and verify PCT
 	kRSA, err := rsa.GenerateKey(rand.Reader, 2048)
 	if err != nil {
 		t.Error(err)
 	} else {
+		// RSASSA-PKCS-v1.5 2048-bit sign and verify
 		rsa.SignPKCS1v15(kRSA, crypto.SHA256.String(), make([]byte, 32))
 	}
 	t.Log("completed successfully")