Usage of JWKS with JWT (without using OpenID connect) (#5578)

Signed-off-by: Rishav Kumar <rishavaz@amazon.com>
Co-authored-by: Rishav Kumar <rishavaz@amazon.com>

Author: Rishav9852Kumar

CHANGELOG.md

@@ -14,6 +14,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - [Resource Sharing] Keep track of tenant for sharable resources by persisting user requested tenant with sharing info ([#5588](https://github.com/opensearch-project/security/pull/5588))
 - [SecurityPlugin Health Check] Add AuthZ initialization completion check in health check API [(#5626)](https://github.com/opensearch-project/security/pull/5626)
 - [Resource Sharing] Adds API to provide dashboards support for resource access management ([#5597](https://github.com/opensearch-project/security/pull/5597))
+- Direct JWKS (JSON Web Key Set) support in the JWT authentication backend ([#5578](https://github.com/opensearch-project/security/pull/5578))
 
 
 ### Bug Fixes

config/config.yml

@@ -130,6 +130,7 @@ config:
           type: jwt
           challenge: false
           config:
+            jwks_uri: 'https://your-jwks-endpoint.com/.well-known/jwks.json'
             signing_key: "base64 encoded HMAC key or public RSA/ECDSA pem key"
             jwt_header: "Authorization"
             jwt_url_parameter: null

src/main/java/org/opensearch/security/auth/http/jwt/keybyjwks/HTTPJwtKeyByJWKSAuthenticator.java

@@ -0,0 +1,158 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ *
+ * Modifications Copyright OpenSearch Contributors. See
+ * GitHub history for details.
+ */
+
+package org.opensearch.security.auth.http.jwt.keybyjwks;
+
+import java.nio.file.Path;
+import java.util.Collections;
+import java.util.Set;
+
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+
+import org.opensearch.OpenSearchSecurityException;
+import org.opensearch.common.settings.Settings;
+import org.opensearch.common.util.concurrent.ThreadContext;
+import org.opensearch.core.common.Strings;
+import org.opensearch.security.auth.http.jwt.AbstractHTTPJwtAuthenticator;
+import org.opensearch.security.auth.http.jwt.HTTPJwtAuthenticator;
+import org.opensearch.security.auth.http.jwt.keybyoidc.KeyProvider;
+import org.opensearch.security.auth.http.jwt.keybyoidc.KeySetRetriever;
+import org.opensearch.security.auth.http.jwt.keybyoidc.SelfRefreshingKeySet;
+import org.opensearch.security.filter.SecurityRequest;
+import org.opensearch.security.user.AuthCredentials;
+import org.opensearch.security.util.SettingsBasedSSLConfigurator;
+
+/**
+ * JWT authenticator that uses JWKS (JSON Web Key Set) endpoints for key retrieval.
+ *
+ * This authenticator extends AbstractHTTPJwtAuthenticator and provides JWKS-specific
+ * key provider initialization. It supports direct JWKS endpoint access with caching,
+ * SSL configuration, automatic key refresh, and enhanced security features to protect
+ * against malicious JWKS endpoints.
+ *
+ * Security Features:
+ * - Response size validation before parsing to prevent memory exhaustion
+ * - Hard key count limit after parsing to reject oversized JWKS
+ * - Configurable timeouts and rate limiting
+ *
+ * Configuration:
+ * - jwks_uri: Direct JWKS endpoint URL (required)
+ * - cache_jwks_endpoint: Enable/disable caching (default: true)
+ * - jwks_request_timeout_ms: Request timeout in milliseconds (default: 5000)
+ * - jwks_queued_thread_timeout_ms: Queued thread timeout (default: 2500)
+ * - refresh_rate_limit_time_window_ms: Rate limit window (default: 10000)
+ * - refresh_rate_limit_count: Max refreshes per window (default: 10)
+ * - max_jwks_keys: HARD LIMIT - Rejects JWKS if exceeded (default: 10)
+ * - max_jwks_response_size_bytes: Max HTTP response size (default: 1MB)
+ */
+public class HTTPJwtKeyByJWKSAuthenticator extends AbstractHTTPJwtAuthenticator {
+
+    private final static Logger log = LogManager.getLogger(HTTPJwtKeyByJWKSAuthenticator.class);
+
+    // Fallback to static JWT authenticator if jwks_uri is null
+    private final HTTPJwtAuthenticator staticJwtAuthenticator;
+    private final boolean useJwks;
+    private final String jwtUrlParameter;
+
+    public HTTPJwtKeyByJWKSAuthenticator(Settings settings, Path configPath) {
+        super(settings, configPath);
+
+        String jwksUri = settings.get("jwks_uri");
+        this.useJwks = !Strings.isNullOrEmpty(jwksUri);
+        this.jwtUrlParameter = settings.get("jwt_url_parameter");
+
+        // Initialize static JWT authenticator as fallback if jwks_uri is not configured
+        if (!useJwks) {
+            log.warn("jwks_uri is not configured, falling back to static JWT authentication");
+            this.staticJwtAuthenticator = new HTTPJwtAuthenticator(settings, configPath);
+        } else {
+            this.staticJwtAuthenticator = null;
+        }
+    }
+
+    @Override
+    protected KeyProvider initKeyProvider(Settings settings, Path configPath) throws Exception {
+        String jwksUri = settings.get("jwks_uri");
+
+        // If jwks_uri is not configured, return null (will use static JWT fallback)
+        if (jwksUri == null || jwksUri.isBlank()) {
+            log.warn("jwks_uri is not configured, will use static JWT authentication fallback");
+            return null;
+        }
+
+        log.debug("Initializing JWKS key provider with endpoint: {}", jwksUri);
+
+        // Initialize configuration parameters
+        int jwksRequestTimeoutMs = settings.getAsInt("jwks_request_timeout_ms", 5000);
+        int jwksQueuedThreadTimeoutMs = settings.getAsInt("jwks_queued_thread_timeout_ms", 2500);
+        int refreshRateLimitTimeWindowMs = settings.getAsInt("refresh_rate_limit_time_window_ms", 10000);
+        int refreshRateLimitCount = settings.getAsInt("refresh_rate_limit_count", 10);
+        boolean cacheJwksEndpoint = settings.getAsBoolean("cache_jwks_endpoint", true);
+        int maxJwksKeys = settings.getAsInt("max_jwks_keys", -1);
+
+        log.warn("Initializing JWKS key provider with endpoint: {} (max keys: {})", jwksUri, maxJwksKeys);
+
+        // Add security configuration parameters
+        long maxJwksResponseSizeBytes = settings.getAsLong("max_jwks_response_size_bytes", 1024L * 1024L); // 1MB default
+
+        // Create secure key set retriever with HARD LIMIT enforcement using maxJwksKeys
+        KeySetRetriever keySetRetriever = KeySetRetriever.createForJwksUri(
+            getSSLConfig(settings, configPath),
+            cacheJwksEndpoint,
+            jwksUri,
+            maxJwksResponseSizeBytes,
+            maxJwksKeys
+        );
+        keySetRetriever.setRequestTimeoutMs(jwksRequestTimeoutMs);
+
+        // Create self-refreshing key set with caching and rate limiting
+        SelfRefreshingKeySet selfRefreshingKeySet = new SelfRefreshingKeySet(keySetRetriever);
+        selfRefreshingKeySet.setRequestTimeoutMs(jwksRequestTimeoutMs);
+        selfRefreshingKeySet.setQueuedThreadTimeoutMs(jwksQueuedThreadTimeoutMs);
+        selfRefreshingKeySet.setRefreshRateLimitTimeWindowMs(refreshRateLimitTimeWindowMs);
+        selfRefreshingKeySet.setRefreshRateLimitCount(refreshRateLimitCount);
+
+        return selfRefreshingKeySet;
+    }
+
+    @Override
+    public AuthCredentials extractCredentials(final SecurityRequest request, final ThreadContext context)
+        throws OpenSearchSecurityException {
+
+        // If jwks_uri is not configured, delegate to static JWT authenticator
+        if (!useJwks && staticJwtAuthenticator != null) {
+            log.debug("Delegating to static JWT authenticator since jwks_uri is not configured");
+            return staticJwtAuthenticator.extractCredentials(request, context);
+        }
+
+        // Use the standard JWKS authentication flow
+        return super.extractCredentials(request, context);
+    }
+
+    private static SettingsBasedSSLConfigurator.SSLConfig getSSLConfig(Settings settings, Path configPath) throws Exception {
+        return new SettingsBasedSSLConfigurator(settings, configPath, "jwks").buildSSLConfig();
+    }
+
+    @Override
+    public String getType() {
+        return "jwt";
+    }
+
+    @Override
+    public Set<String> getSensitiveUrlParams() {
+        if (jwtUrlParameter != null) {
+            return Set.of(jwtUrlParameter);
+        }
+        return Collections.emptySet();
+    }
+
+}

src/main/java/org/opensearch/security/auth/http/jwt/keybyoidc/JwtVerifier.java

@@ -57,6 +57,8 @@ public SignedJWT getVerifiedJwtToken(String encodedJwt) throws BadCredentialsExc
             String kid = escapedKid;
             if (!Strings.isNullOrEmpty(kid)) {
                 kid = StringEscapeUtils.unescapeJava(escapedKid);
+            } else {
+                log.debug("JWT token is missing 'kid' (Key ID) claim in header. This may cause key selection issues.");
             }
             JWK key = keyProvider.getKey(kid);
 
@@ -65,8 +67,10 @@ public SignedJWT getVerifiedJwtToken(String encodedJwt) throws BadCredentialsExc
 
             if (!signatureValid && Strings.isNullOrEmpty(kid)) {
                 key = keyProvider.getKeyAfterRefresh(null);
-                signatureVerifier = getInitializedSignatureVerifier(key, jwt);
-                signatureValid = jwt.verify(signatureVerifier);
+                if (key != null) {
+                    signatureVerifier = getInitializedSignatureVerifier(key, jwt);
+                    signatureValid = jwt.verify(signatureVerifier);
+                }
             }
 
             if (!signatureValid) {

src/main/java/org/opensearch/security/auth/http/jwt/keybyoidc/KeySetRetriever.java

@@ -57,6 +57,11 @@ public class KeySetRetriever implements KeySetProvider {
     private long lastCacheStatusLog = 0;
     private String jwksUri;
 
+    // Security validation settings (optional, for JWKS endpoints)
+    private long maxResponseSizeBytes = -1; // -1 means no limit
+    private int maxKeyCount = -1; // -1 means no limit
+    private boolean enableSecurityValidation = false;
+
     KeySetRetriever(String openIdConnectEndpoint, SSLConfig sslConfig, boolean useCacheForOidConnectEndpoint) {
         this.openIdConnectEndpoint = openIdConnectEndpoint;
         this.sslConfig = sslConfig;
@@ -71,10 +76,41 @@ public class KeySetRetriever implements KeySetProvider {
         configureCache(useCacheForOidConnectEndpoint);
     }
 
+    /**
+     * Factory method to create a KeySetRetriever for JWKS endpoint access.
+     * This method provides a public API for creating KeySetRetriever instances
+     * with built-in security validation to protect against malicious JWKS endpoints.
+     *
+     * @param sslConfig SSL configuration for HTTPS connections
+     * @param useCacheForJwksEndpoint whether to enable caching for JWKS endpoint
+     *                                When true, JWKS responses will be cached to improve performance
+     *                                and reduce network calls to the JWKS endpoint.
+     * @param jwksUri the JWKS endpoint URI
+     * @param maxResponseSizeBytes maximum allowed HTTP response size in bytes
+     * @param maxKeyCount maximum number of keys allowed in JWKS
+     * @return a new KeySetRetriever instance with security validation enabled
+     */
+    public static KeySetRetriever createForJwksUri(
+        SSLConfig sslConfig,
+        boolean useCacheForJwksEndpoint,
+        String jwksUri,
+        long maxResponseSizeBytes,
+        int maxKeyCount
+    ) {
+        KeySetRetriever retriever = new KeySetRetriever(sslConfig, useCacheForJwksEndpoint, jwksUri);
+        retriever.enableSecurityValidation = true;
+        retriever.maxResponseSizeBytes = maxResponseSizeBytes;
+        retriever.maxKeyCount = maxKeyCount;
+        return retriever;
+    }
+
     public JWKSet get() throws AuthenticatorUnavailableException {
         String uri = getJwksUri();
 
-        try (CloseableHttpClient httpClient = createHttpClient(null)) {
+        // Use cache storage if it's configured
+        HttpCacheStorage cacheStorage = oidcHttpCacheStorage;
+
+        try (CloseableHttpClient httpClient = createHttpClient(cacheStorage)) {
 
             HttpGet httpGet = new HttpGet(uri);
 
@@ -85,7 +121,20 @@ public JWKSet get() throws AuthenticatorUnavailableException {
 
             httpGet.setConfig(requestConfig);
 
-            try (CloseableHttpResponse response = httpClient.execute(httpGet)) {
+            // Configure HTTP client to only accept JSON responses for JWKS endpoints
+            if (enableSecurityValidation) {
+                httpGet.setHeader("Accept", "application/json, application/jwk-set+json");
+            }
+
+            HttpCacheContext httpContext = null;
+            if (cacheStorage != null) {
+                httpContext = new HttpCacheContext();
+            }
+
+            try (CloseableHttpResponse response = httpClient.execute(httpGet, httpContext)) {
+                if (httpContext != null) {
+                    logCacheResponseStatus(httpContext, true);
+                }
                 if (response.getCode() < 200 || response.getCode() >= 300) {
                     throw new AuthenticatorUnavailableException("Error while getting " + uri + ": " + response.getReasonPhrase());
                 }
@@ -95,11 +144,41 @@ public JWKSet get() throws AuthenticatorUnavailableException {
                 if (httpEntity == null) {
                     throw new AuthenticatorUnavailableException("Error while getting " + uri + ": Empty response entity");
                 }
+
+                // Apply security validation if enabled (for JWKS endpoints)
+                if (enableSecurityValidation) {
+                    // Validate response size
+                    if (maxResponseSizeBytes > 0) {
+                        long contentLength = httpEntity.getContentLength();
+                        if (contentLength > maxResponseSizeBytes) {
+                            throw new AuthenticatorUnavailableException(
+                                String.format(
+                                    "JWKS response too large from %s: %d bytes (max: %d)",
+                                    uri,
+                                    contentLength,
+                                    maxResponseSizeBytes
+                                )
+                            );
+                        }
+                    }
+                }
+
+                // Load JWKS using Nimbus JOSE (handles JSON parsing and validation)
                 JWKSet keySet = JWKSet.load(httpEntity.getContent());
 
+                // Apply minimal additional validation only for direct JWKS endpoints
+                if (enableSecurityValidation) {
+                    // Simple key count validation - HARD LIMIT
+                    if (maxKeyCount > 0 && keySet.getKeys().size() > maxKeyCount) {
+                        throw new AuthenticatorUnavailableException(
+                            String.format("JWKS from %s contains %d keys, but max allowed is %d", uri, keySet.getKeys().size(), maxKeyCount)
+                        );
+                    }
+                }
+
                 return keySet;
             } catch (ParseException e) {
-                throw new RuntimeException(e);
+                throw new AuthenticatorUnavailableException("Error parsing JWKS from " + uri + ": " + e.getMessage(), e);
             }
         } catch (IOException e) {
             throw new AuthenticatorUnavailableException("Error while getting " + uri + ": " + e, e);
@@ -177,21 +256,43 @@ public void setRequestTimeoutMs(int httpTimeoutMs) {
     }
 
     private void logCacheResponseStatus(HttpCacheContext httpContext) {
+        logCacheResponseStatus(httpContext, false);
+    }
+
+    private void logCacheResponseStatus(HttpCacheContext httpContext, boolean isJwksRequest) {
         this.oidcRequests++;
 
-        switch (httpContext.getCacheResponseStatus()) {
-            case CACHE_HIT:
-                this.oidcCacheHits++;
-                break;
-            case CACHE_MODULE_RESPONSE:
-                this.oidcCacheModuleResponses++;
-                break;
-            case CACHE_MISS:
+        // Handle cache statistics based on the response status
+        // For OIDC discovery flow, only count the JWKS request (not the discovery request)
+        // For direct JWKS URI, count all requests
+        boolean shouldCountStats = (jwksUri != null) || isJwksRequest;
+
+        if (!shouldCountStats) {
+            log.debug("Skipping cache statistics for OIDC discovery request #{}", this.oidcRequests);
+            return;
+        }
+
+        if (httpContext.getCacheResponseStatus() == null) {
+            if (oidcHttpCacheStorage != null) {
                 this.oidcCacheMisses++;
-                break;
-            case VALIDATED:
-                this.oidcCacheHitsValidated++;
-                break;
+                log.debug("Null cache status - counting as cache miss. Total misses: {}", this.oidcCacheMisses);
+            }
+        } else {
+            switch (httpContext.getCacheResponseStatus()) {
+                case CACHE_HIT:
+                    this.oidcCacheHits++;
+                    break;
+                case CACHE_MODULE_RESPONSE:
+                    this.oidcCacheModuleResponses++;
+                    break;
+                case CACHE_MISS:
+                    this.oidcCacheMisses++;
+                    break;
+                case VALIDATED:
+                    this.oidcCacheHits++;
+                    this.oidcCacheHitsValidated++;
+                    break;
+            }
         }
 
         long now = System.currentTimeMillis();
@@ -208,7 +309,6 @@ private void logCacheResponseStatus(HttpCacheContext httpContext) {
             );
             lastCacheStatusLog = now;
         }
-
     }
 
     private CloseableHttpClient createHttpClient(HttpCacheStorage httpCacheStorage) {
@@ -255,4 +355,5 @@ public int getOidcCacheHitsValidated() {
     public int getOidcCacheModuleResponses() {
         return oidcCacheModuleResponses;
     }
+
 }