From 4e53223d07b17d496eba81cb74c79ef91c13489c Mon Sep 17 00:00:00 2001
From: Gerard Snaauw <33763579+gerardsn@users.noreply.github.com>
Date: Fri, 18 Oct 2024 09:02:37 +0200
Subject: [PATCH] always add kid to header (#3499)

---
 crypto/jwx.go      | 13 +++++++------
 crypto/jwx_test.go | 18 ++++++++++++++++++
 crypto/memory.go   |  4 +---
 3 files changed, 26 insertions(+), 9 deletions(-)

diff --git a/crypto/jwx.go b/crypto/jwx.go
index a0779d15e8..381f42a2a4 100644
--- a/crypto/jwx.go
+++ b/crypto/jwx.go
@@ -82,9 +82,7 @@ func (client *Crypto) SignJWS(ctx context.Context, payload []byte, headers map[s
 		return "", err
 	}
 
-	if _, ok := headers["jwk"]; !ok {
-		headers["kid"] = kid
-	}
+	headers["kid"] = kid
 	return SignJWS(ctx, payload, headers, privateKey, detached)
 }
 
@@ -247,10 +245,13 @@ func SignJWS(ctx context.Context, payload []byte, protectedHeaders map[string]in
 			return "", fmt.Errorf("unable to set header %s: %w", key, err)
 		}
 	}
-	// The JWX library is fine with creating a JWK for a private key (including the private exponents), so
-	// we want to make sure the `jwk` header (if present) does not (accidentally) contain a private key.
-	// That would lead to the node leaking its private key material in the resulting JWS which would be very, very bad.
 	if headers.JWK() != nil {
+		// 'kid' has been logged, use 'jwk' to sign
+		_ = headers.Remove(jwk.KeyIDKey)
+
+		// The JWX library is fine with creating a JWK for a private key (including the private exponents), so
+		// we want to make sure the `jwk` header (if present) does not (accidentally) contain a private key.
+		// That would lead to the node leaking its private key material in the resulting JWS which would be very, very bad.
 		var jwkAsPrivateKey crypto.Signer
 		if err := headers.JWK().Raw(&jwkAsPrivateKey); err == nil {
 			// `err != nil` is good in this case, because that means the key is not assignable to crypto.Signer,
diff --git a/crypto/jwx_test.go b/crypto/jwx_test.go
index a87a343f7e..4aea26947f 100644
--- a/crypto/jwx_test.go
+++ b/crypto/jwx_test.go
@@ -281,6 +281,24 @@ func TestCrypto_SignJWS(t *testing.T) {
 		require.NoError(t, err)
 		auditLogs.AssertContains(t, ModuleName, "SignJWS", audit.TestActor, "Signing a JWS with key: kid")
 	})
+	t.Run("writes audit log for jwk", func(t *testing.T) {
+		auditLogs := audit.CaptureAuditLogs(t)
+		key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
+		require.NoError(t, err)
+		publicKeyAsJWK, _ := jwk.FromRaw(key.Public())
+		hdrs := map[string]interface{}{
+			"kid": kid,
+			"jwk": publicKeyAsJWK,
+		}
+
+		signature, err := SignJWS(audit.TestContext(), []byte{1, 2, 3}, hdrs, key, false)
+
+		require.NoError(t, err)
+		auditLogs.AssertContains(t, ModuleName, "SignJWS", audit.TestActor, "Signing a JWS with key: kid")
+		// kid is not in headers
+		msg, err := jws.Parse([]byte(signature))
+		assert.Empty(t, msg.Signatures()[0].ProtectedHeaders().KeyID())
+	})
 
 	t.Run("returns error for not found", func(t *testing.T) {
 		payload, _ := json.Marshal(map[string]interface{}{"iss": "nuts"})
diff --git a/crypto/memory.go b/crypto/memory.go
index 31b9bdad55..56d93ab951 100644
--- a/crypto/memory.go
+++ b/crypto/memory.go
@@ -66,9 +66,7 @@ func (m MemoryJWTSigner) SignJWS(ctx context.Context, payload []byte, headers ma
 	if kid != m.Key.KeyID() {
 		return "", ErrPrivateKeyNotFound
 	}
-	if _, ok := headers["jwk"]; !ok {
-		headers["kid"] = kid
-	}
+	headers["kid"] = kid
 	return SignJWS(ctx, payload, headers, signer, detached)
 }