crypto/mlkem: improved semi-expanded wycheproof tests

Takes an upstream update to the Wycheproof test vectors. This includes
an update to the ML-KEM semi-expanded form decapsulation test schema
that adds the expected K. We use that to tighten up
TestMLKEMSemiExpandedDecapsWycheproof to take advantage.

Change-Id: I67abcc9c1c6728ab154742f9a466e18be22eef89
Reviewed-on: https://go-review.googlesource.com/c/go/+/792940
Reviewed-by: David Chase <drchase@google.com>
Reviewed-by: Dmitri Shuralyov <dmitshur@google.com>
Reviewed-by: Filippo Valsorda <filippo@golang.org>
Auto-Submit: Daniel McCarney <daniel@binaryparadox.net>
LUCI-TryBot-Result: golang-scoped@luci-project-accounts.iam.gserviceaccount.com <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
diff --git a/src/crypto/internal/cryptotest/wycheproof/_schema/go.mod b/src/crypto/internal/cryptotest/wycheproof/_schema/go.mod
index 1253c1c..d7c20ce 100644
--- a/src/crypto/internal/cryptotest/wycheproof/_schema/go.mod
+++ b/src/crypto/internal/cryptotest/wycheproof/_schema/go.mod
@@ -4,7 +4,7 @@
 
 require (
 	github.com/atombender/go-jsonschema v0.23.1
-	github.com/c2sp/wycheproof v0.0.0-20260606153636-6d7cccd0fcb1
+	github.com/c2sp/wycheproof v0.0.0-20260625212325-ee7b4f7e6119
 )
 
 require (
diff --git a/src/crypto/internal/cryptotest/wycheproof/_schema/go.sum b/src/crypto/internal/cryptotest/wycheproof/_schema/go.sum
index bc5fe4e..5ad2f74 100644
--- a/src/crypto/internal/cryptotest/wycheproof/_schema/go.sum
+++ b/src/crypto/internal/cryptotest/wycheproof/_schema/go.sum
@@ -1,7 +1,7 @@
 dario.cat/mergo v1.0.2 h1:85+piFYR1tMbRrLcDwR18y4UKJ3aH1Tbzi24VRW1TK8=
 dario.cat/mergo v1.0.2/go.mod h1:E/hbnu0NxMFBjpMIE34DRGLWqDy0g5FuKDhCb31ngxA=
-github.com/c2sp/wycheproof v0.0.0-20260606153636-6d7cccd0fcb1 h1:veaGluF2y7kZnV/owZzY3FuOZWZQdzqnIcJsbJpFV7A=
-github.com/c2sp/wycheproof v0.0.0-20260606153636-6d7cccd0fcb1/go.mod h1:/5JsOpi3fKmnnig5BQ2mz49TNtnrD308VfTVHcbr5x4=
+github.com/c2sp/wycheproof v0.0.0-20260625212325-ee7b4f7e6119 h1:jOCu6erj86P93+eFn8Qxp6XLwftn1nW0C9RUx7xZwZg=
+github.com/c2sp/wycheproof v0.0.0-20260625212325-ee7b4f7e6119/go.mod h1:/5JsOpi3fKmnnig5BQ2mz49TNtnrD308VfTVHcbr5x4=
 github.com/davecgh/go-spew v0.0.0-20161028175848-04cdfd42973b/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
diff --git a/src/crypto/internal/cryptotest/wycheproof/schema.go b/src/crypto/internal/cryptotest/wycheproof/schema.go
index 50c655f..06939c8 100644
--- a/src/crypto/internal/cryptotest/wycheproof/schema.go
+++ b/src/crypto/internal/cryptotest/wycheproof/schema.go
@@ -4980,6 +4980,10 @@
 }
 
 type MLKEMDecapsTestGroupTestsElem struct {
+	// If present, the shared key the implementation MUST return on a successful
+	// Decapsulate call.
+	K *string `json:"K,omitempty,omitzero"`
+
 	// An input ciphertext
 	C string `json:"c"`
 
@@ -4989,6 +4993,9 @@
 	// The full decapsulation key
 	Dk string `json:"dk"`
 
+	// The encapsulation key bytes of dk.
+	Ek string `json:"ek"`
+
 	// A list of flags
 	Flags []string `json:"flags"`
 
@@ -5011,6 +5018,9 @@
 	if _, ok := raw["dk"]; raw != nil && !ok {
 		return fmt.Errorf("field dk in MLKEMDecapsTestGroupTestsElem: required")
 	}
+	if _, ok := raw["ek"]; raw != nil && !ok {
+		return fmt.Errorf("field ek in MLKEMDecapsTestGroupTestsElem: required")
+	}
 	if _, ok := raw["flags"]; raw != nil && !ok {
 		return fmt.Errorf("field flags in MLKEMDecapsTestGroupTestsElem: required")
 	}
diff --git a/src/crypto/internal/cryptotest/wycheproof/schemaversion.go b/src/crypto/internal/cryptotest/wycheproof/schemaversion.go
index 642266d..9a837a2 100644
--- a/src/crypto/internal/cryptotest/wycheproof/schemaversion.go
+++ b/src/crypto/internal/cryptotest/wycheproof/schemaversion.go
@@ -8,4 +8,4 @@
 
 // wycheproofVersion is the github.com/c2sp/wycheproof module version that
 // schema.go was generated against.
-const wycheproofVersion = "v0.0.0-20260606153636-6d7cccd0fcb1"
+const wycheproofVersion = "v0.0.0-20260625212325-ee7b4f7e6119"
diff --git a/src/crypto/mlkem/mlkem_wycheproof_test.go b/src/crypto/mlkem/mlkem_wycheproof_test.go
index 1188437..c516252 100644
--- a/src/crypto/mlkem/mlkem_wycheproof_test.go
+++ b/src/crypto/mlkem/mlkem_wycheproof_test.go
@@ -356,6 +356,9 @@
 		if !bytes.Equal(mlkem.TestingOnlyExpandedBytes768(dk), dkBytes) {
 			t.Errorf("expanded decapsulation key roundtrip mismatch")
 		}
+		if want := wycheproof.MustDecodeHex(tv.Ek); !bytes.Equal(dk.EncapsulationKey().Bytes(), want) {
+			t.Errorf("encapsulation key mismatch:\n got: %x\nwant: %x", dk.EncapsulationKey().Bytes(), want)
+		}
 		k, err := dk.Decapsulate(ciphertext)
 		if err != nil {
 			if shouldPass {
@@ -367,16 +370,11 @@
 			t.Errorf("Decapsulate unexpectedly succeeded")
 			return
 		}
-		if len(k) != SharedKeySize {
-			t.Errorf("shared key has wrong length: got %d, want %d", len(k), SharedKeySize)
+		if tv.K == nil {
+			t.Fatalf("Decapsulate succeeded but test vector has no expected K")
 		}
-		kFresh, cFresh := dk.EncapsulationKey().Encapsulate()
-		kRT, err := dk.Decapsulate(cFresh)
-		if err != nil {
-			t.Fatalf("Decapsulate of fresh ciphertext: %v", err)
-		}
-		if !bytes.Equal(kFresh, kRT) {
-			t.Errorf("encaps/decaps roundtrip key mismatch")
+		if want := wycheproof.MustDecodeHex(*tv.K); !bytes.Equal(k, want) {
+			t.Errorf("shared key mismatch:\n got: %x\nwant: %x", k, want)
 		}
 
 	case wycheproof.MLKEMDecapsTestGroupParameterSetMLKEM1024:
@@ -390,6 +388,9 @@
 		if !bytes.Equal(mlkem.TestingOnlyExpandedBytes1024(dk), dkBytes) {
 			t.Errorf("expanded decapsulation key roundtrip mismatch")
 		}
+		if want := wycheproof.MustDecodeHex(tv.Ek); !bytes.Equal(dk.EncapsulationKey().Bytes(), want) {
+			t.Errorf("encapsulation key mismatch:\n got: %x\nwant: %x", dk.EncapsulationKey().Bytes(), want)
+		}
 		k, err := dk.Decapsulate(ciphertext)
 		if err != nil {
 			if shouldPass {
@@ -401,16 +402,11 @@
 			t.Errorf("Decapsulate unexpectedly succeeded")
 			return
 		}
-		if len(k) != SharedKeySize {
-			t.Errorf("shared key has wrong length: got %d, want %d", len(k), SharedKeySize)
+		if tv.K == nil {
+			t.Fatalf("Decapsulate succeeded but test vector has no expected K")
 		}
-		kFresh, cFresh := dk.EncapsulationKey().Encapsulate()
-		kRT, err := dk.Decapsulate(cFresh)
-		if err != nil {
-			t.Fatalf("Decapsulate of fresh ciphertext: %v", err)
-		}
-		if !bytes.Equal(kFresh, kRT) {
-			t.Errorf("encaps/decaps roundtrip key mismatch")
+		if want := wycheproof.MustDecodeHex(*tv.K); !bytes.Equal(k, want) {
+			t.Errorf("shared key mismatch:\n got: %x\nwant: %x", k, want)
 		}
 
 	default: