fix: add support for multiple jwks keys per endpoint (#753)

CHANGELOG.md

@@ -11,6 +11,7 @@
 - fix: CORS: allow all request headers [#741](https://github.com/hypermodeinc/modus/pull/741)
 - fix: accept long base64 strings [#742](https://github.com/hypermodeinc/modus/pull/742)
 - fix: improve dgraph auth header passing [#752](https://github.com/hypermodeinc/modus/pull/752)
+- fix: jwks endpoint should use key ID if available [#730](https://github.com/hypermodeinc/modus/pull/753)
 
 ## 2025-01-24 - Runtime 0.17.1
 

runtime/middleware/jwt.go

@@ -205,34 +205,37 @@ func jwksEndpointsJsonToKeys(ctx context.Context, jwksEndpointsJson string) (map
 		return nil, err
 	}
 	keys := make(map[string]any)
-	for key, value := range jwksEndpoints {
+	for endpointKey, value := range jwksEndpoints {
 		jwks, err := jwk.Fetch(ctx, value)
 		if err != nil {
 			return nil, err
 		}
 
-		jwkKey, exists := jwks.Key(0)
-		if !exists {
-			return nil, errors.New("No keys found in JWKS for key: " + key)
-		}
+		for it := jwks.Keys(ctx); it.Next(ctx); {
+			jwkKey := it.Pair().Value.(jwk.Key)
+			var rawKey any
+			if err := jwkKey.Raw(&rawKey); err != nil {
+				return nil, err
+			}
 
-		var rawKey any
-		err = jwkKey.Raw(&rawKey)
-		if err != nil {
-			return nil, err
-		}
+			// Marshal the raw key into DER-encoded PKIX format
+			derBytes, err := x509.MarshalPKIXPublicKey(rawKey)
+			if err != nil {
+				return nil, err
+			}
 
-		// Marshal the raw key into DER-encoded PKIX format
-		derBytes, err := x509.MarshalPKIXPublicKey(rawKey)
-		if err != nil {
-			return nil, err
-		}
+			pubKey, err := x509.ParsePKIXPublicKey(derBytes)
+			if err != nil {
+				return nil, err
+			}
 
-		pubKey, err := x509.ParsePKIXPublicKey(derBytes)
-		if err != nil {
-			return nil, err
+			// Use a combination of endpoint key and key ID (if available) as the map key
+			keyID := endpointKey
+			if kid, exists := jwkKey.Get("kid"); exists {
+				keyID = endpointKey + "_" + kid.(string)
+			}
+			keys[keyID] = pubKey
 		}
-		keys[key] = pubKey
 	}
 	return keys, nil
 }