From 0559e8b99cc9ff41565be05bcb619d52ade73122 Mon Sep 17 00:00:00 2001
From: Stanley Shen <sshen37@ucsc.edu>
Date: Sat, 26 Sep 2026 19:08:49 -0700
Subject: [PATCH] hpke: set the scheme on hybrid KEM keys

DeriveKeyPair and the two unmarshalers left the scheme field of
hybridKEMPrivKey and hybridKEMPubKey nil, so Scheme() returned a nil
kem.Scheme and calling any method on it panicked.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
---
 hpke/hybridkem.go | 20 ++++++++++-------
 hpke/kem_test.go  | 55 +++++++++++++++++++++++++++++++++++++++--------
 2 files changed, 58 insertions(+), 17 deletions(-)

diff --git a/hpke/hybridkem.go b/hpke/hybridkem.go
index ed99d2741..9a3166366 100644
--- a/hpke/hybridkem.go
+++ b/hpke/hybridkem.go
@@ -190,12 +190,14 @@ func (h hybridKEM) DeriveKeyPair(seed []byte) (kem.PublicKey, kem.PrivateKey) {
 	pubB, privB := h.kemB.DeriveKeyPair(seedB)
 
 	privKey := &hybridKEMPrivKey{
-		privA: privA,
-		privB: privB,
+		scheme: h,
+		privA:  privA,
+		privB:  privB,
 	}
 	pubKey := &hybridKEMPubKey{
-		pubA: pubA,
-		pubB: pubB,
+		scheme: h,
+		pubA:   pubA,
+		pubB:   pubB,
 	}
 
 	return pubKey, privKey
@@ -225,8 +227,9 @@ func (h hybridKEM) UnmarshalBinaryPrivateKey(data []byte) (kem.PrivateKey, error
 	}
 
 	return &hybridKEMPrivKey{
-		privA: skA,
-		privB: skB,
+		scheme: h,
+		privA:  skA,
+		privB:  skB,
 	}, nil
 }
 
@@ -244,7 +247,8 @@ func (h hybridKEM) UnmarshalBinaryPublicKey(data []byte) (kem.PublicKey, error)
 	}
 
 	return &hybridKEMPubKey{
-		pubA: pkA,
-		pubB: pkB,
+		scheme: h,
+		pubA:   pkA,
+		pubB:   pkB,
 	}, nil
 }
diff --git a/hpke/kem_test.go b/hpke/kem_test.go
index 40cdced19..976298f90 100644
--- a/hpke/kem_test.go
+++ b/hpke/kem_test.go
@@ -6,27 +6,42 @@ import (
 
 	"github.com/cloudflare/circl/hpke"
 	"github.com/cloudflare/circl/internal/test"
+	"github.com/cloudflare/circl/kem"
 )
 
+var kemIDs = []hpke.KEM{
+	hpke.KEM_P256_HKDF_SHA256,
+	hpke.KEM_P384_HKDF_SHA384,
+	hpke.KEM_P521_HKDF_SHA512,
+	hpke.KEM_X25519_HKDF_SHA256,
+	hpke.KEM_X448_HKDF_SHA512,
+	hpke.KEM_X25519_KYBER768_DRAFT00,
+}
+
 func TestKemKeysExactLength(t *testing.T) {
-	for _, kemID := range []hpke.KEM{
-		hpke.KEM_P256_HKDF_SHA256,
-		hpke.KEM_P384_HKDF_SHA384,
-		hpke.KEM_P521_HKDF_SHA512,
-		hpke.KEM_X25519_HKDF_SHA256,
-		hpke.KEM_X448_HKDF_SHA512,
-		hpke.KEM_X25519_KYBER768_DRAFT00,
-	} {
+	for _, kemID := range kemIDs {
 		checkExactLengthUnmarshal(t, kemID)
 	}
 }
 
-func checkExactLengthUnmarshal(t *testing.T, kemID hpke.KEM) {
+func TestKemKeysScheme(t *testing.T) {
+	for _, kemID := range kemIDs {
+		checkKeysScheme(t, kemID)
+	}
+}
+
+func genKeyPair(t *testing.T, kemID hpke.KEM) (kem.Scheme, kem.PublicKey, kem.PrivateKey) {
+	t.Helper()
 	scheme := kemID.Scheme()
 	pk, sk, err := scheme.GenerateKeyPair()
 	if err != nil {
 		t.Fatal(err)
 	}
+	return scheme, pk, sk
+}
+
+func checkExactLengthUnmarshal(t *testing.T, kemID hpke.KEM) {
+	scheme, pk, sk := genKeyPair(t, kemID)
 	skBytes, err := sk.MarshalBinary()
 	test.CheckNoErr(t, err, "marshal private key")
 	pkBytes, err := pk.MarshalBinary()
@@ -58,3 +73,25 @@ func checkExactLengthUnmarshal(t *testing.T, kemID hpke.KEM) {
 		test.CheckOk(pk.Equal(gotPk), "public keys must match", t)
 	})
 }
+
+func checkKeysScheme(t *testing.T, kemID hpke.KEM) {
+	scheme, pk, sk := genKeyPair(t, kemID)
+
+	t.Run(scheme.Name(), func(t *testing.T) {
+		test.CheckOk(sk.Scheme().Name() == scheme.Name(), "private key scheme must match", t)
+		test.CheckOk(pk.Scheme().Name() == scheme.Name(), "public key scheme must match", t)
+		test.CheckOk(sk.Public().Scheme().Name() == scheme.Name(), "derived public key scheme must match", t)
+
+		skBytes, err := sk.MarshalBinary()
+		test.CheckNoErr(t, err, "marshal private key")
+		gotSk, err := scheme.UnmarshalBinaryPrivateKey(skBytes)
+		test.CheckNoErr(t, err, "unmarshal private key")
+		test.CheckOk(gotSk.Scheme().Name() == scheme.Name(), "unmarshaled private key scheme must match", t)
+
+		pkBytes, err := pk.MarshalBinary()
+		test.CheckNoErr(t, err, "marshal public key")
+		gotPk, err := scheme.UnmarshalBinaryPublicKey(pkBytes)
+		test.CheckNoErr(t, err, "unmarshal public key")
+		test.CheckOk(gotPk.Scheme().Name() == scheme.Name(), "unmarshaled public key scheme must match", t)
+	})
+}