From 862d14117104ddba4c572272bc72a46825b43fbb Mon Sep 17 00:00:00 2001
From: Livio Spring <livio.a@gmail.com>
Date: Mon, 26 Aug 2024 11:26:13 +0200
Subject: [PATCH] fix: internal check of JWT access tokens (#8486)

# Which Problems Are Solved

When using a JWT access_token on a ZITADEL API, the token was considered
invalid If the `WebKey` feature flag is enabled.

# How the Problems Are Solved

- Merge the new and old web keys if flag is enabled (as already done for
the keys endpoint).

# Additional Changes

None

# Additional Context

relates to #8449
---
 .../eventstore/token_verifier.go              | 32 +++++++++++++------
 1 file changed, 22 insertions(+), 10 deletions(-)

diff --git a/internal/authz/repository/eventsourcing/eventstore/token_verifier.go b/internal/authz/repository/eventsourcing/eventstore/token_verifier.go
index c9286d61dd..1a50c141d6 100644
--- a/internal/authz/repository/eventsourcing/eventstore/token_verifier.go
+++ b/internal/authz/repository/eventsourcing/eventstore/token_verifier.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/base64"
 	"fmt"
+	"slices"
 	"strings"
 	"time"
 
@@ -327,28 +328,39 @@ type openIDKeySet struct {
 
 // VerifySignature implements the oidc.KeySet interface
 // providing an implementation for the keys retrieved directly from Queries
-func (o *openIDKeySet) VerifySignature(ctx context.Context, jws *jose.JSONWebSignature) ([]byte, error) {
-	keySet, err := o.Queries.ActivePublicKeys(ctx, time.Now())
+func (o *openIDKeySet) VerifySignature(ctx context.Context, jws *jose.JSONWebSignature) (payload []byte, err error) {
+	keySet := new(jose.JSONWebKeySet)
+	if authz.GetFeatures(ctx).WebKey {
+		keySet, err = o.Queries.GetWebKeySet(ctx)
+		if err != nil {
+			return nil, err
+		}
+	}
+	legacyKeySet, err := o.Queries.ActivePublicKeys(ctx, time.Now())
 	if err != nil {
 		return nil, fmt.Errorf("error fetching keys: %w", err)
 	}
+	appendPublicKeysToWebKeySet(keySet, legacyKeySet)
 	keyID, alg := oidc.GetKeyIDAndAlg(jws)
-	key, err := oidc.FindMatchingKey(keyID, oidc.KeyUseSignature, alg, jsonWebKeys(keySet.Keys)...)
+	key, err := oidc.FindMatchingKey(keyID, oidc.KeyUseSignature, alg, keySet.Keys...)
 	if err != nil {
 		return nil, fmt.Errorf("invalid signature: %w", err)
 	}
 	return jws.Verify(&key)
 }
 
-func jsonWebKeys(keys []query.PublicKey) []jose.JSONWebKey {
-	webKeys := make([]jose.JSONWebKey, len(keys))
-	for i, key := range keys {
-		webKeys[i] = jose.JSONWebKey{
+func appendPublicKeysToWebKeySet(keyset *jose.JSONWebKeySet, pubkeys *query.PublicKeys) {
+	if pubkeys == nil || len(pubkeys.Keys) == 0 {
+		return
+	}
+	keyset.Keys = slices.Grow(keyset.Keys, len(pubkeys.Keys))
+
+	for _, key := range pubkeys.Keys {
+		keyset.Keys = append(keyset.Keys, jose.JSONWebKey{
+			Key:       key.Key(),
 			KeyID:     key.ID(),
 			Algorithm: key.Algorithm(),
 			Use:       key.Use().String(),
-			Key:       key.Key(),
-		}
+		})
 	}
-	return webKeys
 }