feat: ML-KEM change key combiner to latest draft

See: openpgp-pqc/draft-openpgp-pqc#161
ProtonMail · Dec 5, 2024 · b5dbf0e · b5dbf0e
1 parent 2e09cd7
commit b5dbf0e
Show file tree

Hide file tree

Showing 3 changed files with 696 additions and 175 deletions.
diff --git a/openpgp/mlkem_ecdh/mlkem_ecdh.go b/openpgp/mlkem_ecdh/mlkem_ecdh.go
@@ -19,6 +19,7 @@ import (
 const (
 	maxSessionKeyLength = 64
 	MlKemSeedLen        = 64
+	kdfContext          = "OpenPGPCompositeKDFv1"
 )
 
 type PublicKey struct {
@@ -138,16 +139,11 @@ func Decrypt(priv *PrivateKey, kEphemeral, ecEphemeral, ciphertext []byte) (msg
 	return keywrap.Unwrap(kek, ciphertext)
 }
 
-// buildKey implements the composite KDF 2a from
-// https://mailarchive.ietf.org/arch/msg/openpgp/NMTCy707LICtxIhP3Xt1U5C8MF0/
+// buildKey implements the composite KDF from
+// https://github.com/openpgp-pqc/draft-openpgp-pqc/pull/161
 func buildKey(pub *PublicKey, eccSecretPoint, eccEphemeral, eccPublicKey, mlkemKeyShare, mlkemEphemeral []byte, mlkemPublicKey kem.PublicKey) ([]byte, error) {
-	h := sha3.New256()
-
-	// SHA3 never returns error
-	_, _ = h.Write(eccSecretPoint)
-	_, _ = h.Write(eccEphemeral)
-	_, _ = h.Write(eccPublicKey)
-	eccKeyShare := h.Sum(nil)
+	/// Set the output `ecdhKeyShare` to `eccSecretPoint`
+	eccKeyShare := eccSecretPoint
 
 	serializedMlkemPublicKey, err := mlkemPublicKey.MarshalBinary()
 	if err != nil {
@@ -162,16 +158,17 @@ func buildKey(pub *PublicKey, eccSecretPoint, eccEphemeral, eccPublicKey, mlkemK
 	//   eccEphemeral    - the ECDH ciphertext encoded as an octet string
 	//   eccPublicKey    - The ECDH public key of the recipient as an octet string
 
-	// 2a. SHA3-256(mlkemKeyShare || eccKeyShare || eccEphemeral || eccPublicKey || Domain)
-	//       where Domain is "Domain" for LAMPS, and "mlkemEphemeral || mlkemPublicKey || algId" for OpenPGP
-	h.Reset()
+	// SHA3-256(mlkemKeyShare || eccKeyShare || eccEphemeral || eccPublicKey ||
+	//          mlkemEphemeral || mlkemPublicKey || algId || "OpenPGPCompositeKDFv1")
+	h := sha3.New256()
 	_, _ = h.Write(mlkemKeyShare)
 	_, _ = h.Write(eccKeyShare)
 	_, _ = h.Write(eccEphemeral)
 	_, _ = h.Write(eccPublicKey)
 	_, _ = h.Write(mlkemEphemeral)
 	_, _ = h.Write(serializedMlkemPublicKey)
 	_, _ = h.Write([]byte{pub.AlgId})
+	_, _ = h.Write([]byte(kdfContext))
 	return h.Sum(nil), nil
 }
 

diff --git a/openpgp/read_test.go b/openpgp/read_test.go
@@ -967,13 +967,12 @@ var pqcDraftVectors = map[string]struct {
 	fingerprints      []string
 	armoredMessages   []string
 }{
-	// Update with fresh vectors
-	//"v6_Ed25519_ML-KEM-768+X25519": {
-	//	v6Ed25519Mlkem768X25519PrivateTestVector,
-	//	v6Ed25519Mlkem768X25519PublicTestVector,
-	//	[]string{"52343242345254050219ceff286e9c8e479ec88757f95354388984a02d7d0b59", "263e34b69938e753dc67ca8ee37652795135e0e16e48887103c11d7307df40ed"},
-	//	[]string{v6Ed25519Mlkem768X25519PrivateMessageTestVector},
-	//},
+	"v6_Ed25519_ML-KEM-768+X25519": {
+		v6Ed25519Mlkem768X25519PrivateTestVector,
+		v6Ed25519Mlkem768X25519PublicTestVector,
+		[]string{"bf262b24177002ac8ae5dc6da47c056d22ab9906d47d07952b75c358021901ca", "48b94bce2f9771788f5feb74122d599989c400cc0f49108bc98e0ea7945e4838"},
+		[]string{v6Ed25519Mlkem768X25519PrivateMessageTestVector},
+	},
 }
 
 func TestPqcDraftVectors(t *testing.T) {
@@ -1040,7 +1039,7 @@ func TestPqcDraftVectors(t *testing.T) {
 						return
 					}
 
-					if string(contents) != "Testing\n" {
+					if string(contents) != "Testing\r\n" {
 						t.Fatalf("Decrypted message is wrong: %s", contents)
 					}
 				})