From 219df7087c1823fdfa4340ce66198d722a4e5648 Mon Sep 17 00:00:00 2001
From: Jason O'Donnell <2160810+jasonodonnell@users.noreply.github.com>
Date: Wed, 16 Mar 2022 21:48:10 -0400
Subject: [PATCH] identity/token: fix duplicate keys in well-known (#14543)

* identity/token: fix duplicate kids in well-known

* Remove unused check

* changelog

* use map-based approach to dedup key IDs

* improve changelog description

* move jwks closer to usage; specify capacity

Co-authored-by: Austin Gebauer <agebauer@hashicorp.com>
---
 changelog/14543.txt               |  3 ++
 vault/identity_store_oidc.go      | 31 +++++++++++++-------
 vault/identity_store_oidc_test.go | 48 +++++++++++++++++++++++++++++++
 3 files changed, 71 insertions(+), 11 deletions(-)
 create mode 100644 changelog/14543.txt

diff --git a/changelog/14543.txt b/changelog/14543.txt
new file mode 100644
index 000000000..649d1e72c
--- /dev/null
+++ b/changelog/14543.txt
@@ -0,0 +1,3 @@
+```release-note:bug
+identity/token: Fixes a bug where duplicate public keys could appear in the .well-known JWKS
+```
diff --git a/vault/identity_store_oidc.go b/vault/identity_store_oidc.go
index a019ed558..f30edf598 100644
--- a/vault/identity_store_oidc.go
+++ b/vault/identity_store_oidc.go
@@ -621,9 +621,11 @@ func (i *IdentityStore) keyIDsByName(ctx context.Context, s logical.Storage, nam
 	if err := entry.DecodeJSON(&key); err != nil {
 		return keyIDs, err
 	}
+
 	for _, k := range key.KeyRing {
 		keyIDs = append(keyIDs, k.KeyID)
 	}
+
 	return keyIDs, nil
 }
 
@@ -1660,16 +1662,14 @@ func (i *IdentityStore) generatePublicJWKS(ctx context.Context, s logical.Storag
 		return nil, err
 	}
 
-	jwks := &jose.JSONWebKeySet{
-		Keys: make([]jose.JSONWebKey, 0),
-	}
-
 	// only return keys that are associated with a role
 	roleNames, err := s.List(ctx, roleConfigPath)
 	if err != nil {
 		return nil, err
 	}
 
+	// collect and deduplicate the key IDs for all roles
+	keyIDs := make(map[string]struct{})
 	for _, roleName := range roleNames {
 		role, err := i.getOIDCRole(ctx, s, roleName)
 		if err != nil {
@@ -1679,20 +1679,29 @@ func (i *IdentityStore) generatePublicJWKS(ctx context.Context, s logical.Storag
 			continue
 		}
 
-		keyIDs, err := i.keyIDsByName(ctx, s, role.Key)
+		roleKeyIDs, err := i.keyIDsByName(ctx, s, role.Key)
 		if err != nil {
 			return nil, err
 		}
 
-		for _, keyID := range keyIDs {
-			key, err := loadOIDCPublicKey(ctx, s, keyID)
-			if err != nil {
-				return nil, err
-			}
-			jwks.Keys = append(jwks.Keys, *key)
+		for _, keyID := range roleKeyIDs {
+			keyIDs[keyID] = struct{}{}
 		}
 	}
 
+	jwks := &jose.JSONWebKeySet{
+		Keys: make([]jose.JSONWebKey, 0, len(keyIDs)),
+	}
+
+	// load the JSON web key for each key ID
+	for keyID := range keyIDs {
+		key, err := loadOIDCPublicKey(ctx, s, keyID)
+		if err != nil {
+			return nil, err
+		}
+		jwks.Keys = append(jwks.Keys, *key)
+	}
+
 	if err := i.oidcCache.SetDefault(ns, "jwks", jwks); err != nil {
 		return nil, err
 	}
diff --git a/vault/identity_store_oidc_test.go b/vault/identity_store_oidc_test.go
index 8d23d63b3..2f8d0e3c6 100644
--- a/vault/identity_store_oidc_test.go
+++ b/vault/identity_store_oidc_test.go
@@ -759,6 +759,54 @@ func TestOIDC_PublicKeys(t *testing.T) {
 	assertPublicKeyCount(t, ctx, storage, c, 2)
 }
 
+// TestOIDC_PublicKeys tests that public keys are updated by
+// key creation, rotation, and deletion
+func TestOIDC_SharedPublicKeysByRoles(t *testing.T) {
+	c, _, _ := TestCoreUnsealed(t)
+	ctx := namespace.RootContext(nil)
+	storage := &logical.InmemStorage{}
+
+	// Create a test key "test-key"
+	c.identityStore.HandleRequest(ctx, &logical.Request{
+		Path:      "oidc/key/test-key",
+		Operation: logical.CreateOperation,
+		Storage:   storage,
+	})
+
+	// Create a test role "test-role"
+	c.identityStore.HandleRequest(ctx, &logical.Request{
+		Path:      "oidc/role/test-role",
+		Operation: logical.CreateOperation,
+		Data: map[string]interface{}{
+			"key": "test-key",
+		},
+		Storage: storage,
+	})
+
+	// Create a test role "test-role2"
+	c.identityStore.HandleRequest(ctx, &logical.Request{
+		Path:      "oidc/role/test-role2",
+		Operation: logical.CreateOperation,
+		Data: map[string]interface{}{
+			"key": "test-key",
+		},
+		Storage: storage,
+	})
+
+	// Create a test role "test-role3"
+	c.identityStore.HandleRequest(ctx, &logical.Request{
+		Path:      "oidc/role/test-role3",
+		Operation: logical.CreateOperation,
+		Data: map[string]interface{}{
+			"key": "test-key",
+		},
+		Storage: storage,
+	})
+
+	// .well-known/keys should contain 2 public keys
+	assertPublicKeyCount(t, ctx, storage, c, 2)
+}
+
 // TestOIDC_SignIDToken tests acquiring a signed token and verifying the public portion
 // of the signing key
 func TestOIDC_SignIDToken(t *testing.T) {