From b6a28bda368e8e5d99a8559859ca9fcd503d7af9 Mon Sep 17 00:00:00 2001
From: yuhan6665 <1588741+yuhan6665@users.noreply.github.com>
Date: Sun, 4 May 2025 17:04:23 -0400
Subject: [PATCH] crypto/internal/mlkem768: add EncapsulationKey type
 Change-Id: I3feacb044caa15ac9bbfc11f5d90bebf8a505510 Reviewed-on:
 https://go-review.googlesource.com/c/go/+/621980 Auto-Submit: Filippo
 Valsorda <filippo@golang.org> Reviewed-by: Roland Shoemaker
 <roland@golang.org> LUCI-TryBot-Result: Go LUCI
 <golang-scoped@luci-project-accounts.iam.gserviceaccount.com> Reviewed-by:
 Daniel McCarney <daniel@binaryparadox.net> Reviewed-by: Russ Cox
 <rsc@golang.org>

---
 handshake_client.go | 4 ++--
 key_schedule.go     | 9 +++++----
 2 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/handshake_client.go b/handshake_client.go
index fddee7a..679f1a6 100644
--- a/handshake_client.go
+++ b/handshake_client.go
@@ -164,7 +164,7 @@ func (c *Conn) makeClientHello() (*clientHelloMsg, *keySharePrivateKeys, *echCon
 			if _, err := io.ReadFull(config.rand(), seed); err != nil {
 				return nil, nil, nil, err
 			}
-			keyShareKeys.kyber, err = mlkem768.NewKeyFromSeed(seed)
+			keyShareKeys.kyber, err = mlkem768.NewDecapsulationKey(seed)
 			if err != nil {
 				return nil, nil, nil, err
 			}
@@ -174,7 +174,7 @@ func (c *Conn) makeClientHello() (*clientHelloMsg, *keySharePrivateKeys, *echCon
 			// both, as allowed by draft-ietf-tls-hybrid-design-09, Section 3.2.
 			hello.keyShares = []keyShare{
 				{group: x25519Kyber768Draft00, data: append(keyShareKeys.ecdhe.PublicKey().Bytes(),
-					keyShareKeys.kyber.EncapsulationKey()...)},
+					keyShareKeys.kyber.EncapsulationKey().Bytes()...)},
 				{group: X25519, data: keyShareKeys.ecdhe.PublicKey().Bytes()},
 			}
 		} else {
diff --git a/key_schedule.go b/key_schedule.go
index 8f5f2cf..f05a991 100644
--- a/key_schedule.go
+++ b/key_schedule.go
@@ -64,19 +64,20 @@ func kyberDecapsulate(dk *mlkem768.DecapsulationKey, c []byte) ([]byte, error) {
 	if err != nil {
 		return nil, err
 	}
-	return kyberSharedSecret(K, c), nil
+	return kyberSharedSecret(c, K), nil
 }
 
 // kyberEncapsulate implements encapsulation according to Kyber Round 3.
 func kyberEncapsulate(ek []byte) (c, ss []byte, err error) {
-	c, ss, err = mlkem768.Encapsulate(ek)
+	k, err := mlkem768.NewEncapsulationKey(ek)
 	if err != nil {
 		return nil, nil, err
 	}
-	return c, kyberSharedSecret(ss, c), nil
+	c, ss = k.Encapsulate()
+	return c, kyberSharedSecret(c, ss), nil
 }
 
-func kyberSharedSecret(K, c []byte) []byte {
+func kyberSharedSecret(c, K []byte) []byte {
 	// Package mlkem768 implements ML-KEM, which compared to Kyber removed a
 	// final hashing step. Compute SHAKE-256(K || SHA3-256(c), 32) to match Kyber.
 	// See https://words.filippo.io/mlkem768/#bonus-track-using-a-ml-kem-implementation-as-kyber-v3.