[FEATURE] usage of JWKS with JWT (w/o OpenID connect)

Signed-off-by: Sebastian Michalski <shekerama@gmail.com>

Author: sebastianmichalski

src/main/java/com/amazon/dlic/auth/http/jwt/AbstractHTTPJwtAuthenticator.java

@@ -55,6 +55,8 @@ public abstract class AbstractHTTPJwtAuthenticator implements HTTPAuthenticator
     private final String jwtUrlParameter;
     private final String subjectKey;
     private final String rolesKey;
+    private final String requiredAudience;
+    private final String requiredIssuer;
 
     public static final int DEFAULT_CLOCK_SKEW_TOLERANCE_SECONDS = 30;
     private final int clockSkewToleranceSeconds ;
@@ -66,6 +68,8 @@ public AbstractHTTPJwtAuthenticator(Settings settings, Path configPath) {
         rolesKey = settings.get("roles_key");
         subjectKey = settings.get("subject_key");
         clockSkewToleranceSeconds = settings.getAsInt("jwt_clock_skew_tolerance_seconds", DEFAULT_CLOCK_SKEW_TOLERANCE_SECONDS);
+        requiredAudience = settings.get("required_audience");
+        requiredIssuer = settings.get("required_issuer");
 
         try {
             this.keyProvider = this.initKeyProvider(settings, configPath);
@@ -233,4 +237,12 @@ public boolean reRequestAuthentication(RestChannel channel, AuthCredentials auth
         return true;
     }
 
+    public String getRequiredAudience() {
+        return requiredAudience;
+    }
+
+    public String getRequiredIssuer() {
+        return requiredIssuer;
+    }
+
 }

src/main/java/com/amazon/dlic/auth/http/jwt/keybyoidc/HTTPJwtKeyByOpenIdConnectAuthenticator.java

@@ -32,9 +32,15 @@ protected KeyProvider initKeyProvider(Settings settings, Path configPath) throws
 
 		int refreshRateLimitTimeWindowMs = settings.getAsInt("refresh_rate_limit_time_window_ms", 10000);
 		int refreshRateLimitCount = settings.getAsInt("refresh_rate_limit_count", 10);
-
-		KeySetRetriever keySetRetriever = new KeySetRetriever(settings.get("openid_connect_url"),
-				getSSLConfig(settings, configPath), settings.getAsBoolean("cache_jwks_endpoint", false));
+		var jwksUri = settings.get("jwks_uri");
+
+		KeySetRetriever keySetRetriever;
+		if(jwksUri != null && !jwksUri.isBlank()) {
+			keySetRetriever =
+				new KeySetRetriever(getSSLConfig(settings, configPath), settings.getAsBoolean("cache_jwks_endpoint", false), jwksUri);
+		} else {
+			keySetRetriever = new KeySetRetriever(settings.get("openid_connect_url"), getSSLConfig(settings, configPath), settings.getAsBoolean("cache_jwks_endpoint", false));
+		}
 
 		keySetRetriever.setRequestTimeoutMs(idpRequestTimeoutMs);
 

src/main/java/com/amazon/dlic/auth/http/jwt/keybyoidc/KeySetRetriever.java

@@ -54,15 +54,20 @@ public class KeySetRetriever implements KeySetProvider {
 	private int oidcCacheModuleResponses = 0;
 	private long oidcRequests = 0;
 	private long lastCacheStatusLog = 0;
+	private String jwksUri;
 
 	KeySetRetriever(String openIdConnectEndpoint, SSLConfig sslConfig, boolean useCacheForOidConnectEndpoint) {
 		this.openIdConnectEndpoint = openIdConnectEndpoint;
 		this.sslConfig = sslConfig;
 
-		if (useCacheForOidConnectEndpoint) {
-			cacheConfig = CacheConfig.custom().setMaxCacheEntries(10).setMaxObjectSize(1024L * 1024L).build();
-			oidcHttpCacheStorage = new BasicHttpCacheStorage(cacheConfig);
-		}
+		configureCache(useCacheForOidConnectEndpoint);
+	}
+
+	KeySetRetriever(SSLConfig sslConfig, boolean useCacheForOidConnectEndpoint, String jwksUri) {
+		this.jwksUri = jwksUri;
+		this.sslConfig = sslConfig;
+
+		configureCache(useCacheForOidConnectEndpoint);
 	}
 
 	public JsonWebKeys get() throws AuthenticatorUnavailableException {
@@ -101,6 +106,10 @@ public JsonWebKeys get() throws AuthenticatorUnavailableException {
 
 	String getJwksUri() throws AuthenticatorUnavailableException {
 
+		if (jwksUri != null && !jwksUri.isBlank()) {
+			return jwksUri;
+		}
+
 		try (CloseableHttpClient httpClient = createHttpClient(oidcHttpCacheStorage)) {
 
 			HttpGet httpGet = new HttpGet(openIdConnectEndpoint);
@@ -204,6 +213,13 @@ private CloseableHttpClient createHttpClient(HttpCacheStorage httpCacheStorage)
 		return builder.build();
 	}
 
+	private void configureCache(boolean useCacheForOidConnectEndpoint) {
+		if (useCacheForOidConnectEndpoint) {
+			cacheConfig = CacheConfig.custom().setMaxCacheEntries(10).setMaxObjectSize(1024L * 1024L).build();
+			oidcHttpCacheStorage = new BasicHttpCacheStorage(cacheConfig);
+		}
+	}
+
 	public int getOidcCacheHits() {
 		return oidcCacheHits;
 	}

src/test/java/com/amazon/dlic/auth/http/jwt/keybyoidc/HTTPJwtKeyByOpenIdConnectAuthenticatorTest.java

@@ -18,6 +18,7 @@
 import org.junit.BeforeClass;
 import org.junit.Test;
 
+import org.opensearch.OpenSearchSecurityException;
 import org.opensearch.common.settings.Settings;
 import org.opensearch.security.user.AuthCredentials;
 import org.opensearch.security.util.FakeRestRequest;
@@ -58,6 +59,44 @@ public void basicTest() {
 		Assert.assertEquals(3, creds.getAttributes().size());
 	}
 
+
+	@Test
+	public void jwksUriTest() {
+		String requiredIssuer = "requiredIssuer";
+		String requiredAudience = "requiredAudience";
+		Settings settings = Settings.builder()
+			.put("jwks_uri", mockIdpServer.getJwksUri())
+			.put("required_issuer", requiredIssuer)
+			.put("required_audience", requiredAudience)
+			.build();
+
+		HTTPJwtKeyByOpenIdConnectAuthenticator jwtAuth = new HTTPJwtKeyByOpenIdConnectAuthenticator(settings, null);
+
+		AuthCredentials creds = jwtAuth.extractCredentials(new FakeRestRequest(
+			ImmutableMap.of("Authorization", TestJwts.MC_COY_SIGNED_OCT_1), new HashMap<>()), null);
+
+		Assert.assertNotNull(creds);
+		Assert.assertEquals(TestJwts.MCCOY_SUBJECT, creds.getUsername());
+		Assert.assertEquals(TestJwts.TEST_AUDIENCE, creds.getAttributes().get("attr.jwt.aud"));
+		Assert.assertEquals(0, creds.getBackendRoles().size());
+		Assert.assertEquals(3, creds.getAttributes().size());
+		Assert.assertEquals(requiredAudience, jwtAuth.getRequiredAudience());
+		Assert.assertEquals(requiredIssuer, jwtAuth.getRequiredIssuer());
+	}
+
+	@Test
+	public void jwksUriMissingTest() {
+		var exception = Assert.assertThrows(Exception.class, () -> {
+			HTTPJwtKeyByOpenIdConnectAuthenticator jwtAuth = new HTTPJwtKeyByOpenIdConnectAuthenticator(Settings.builder().build(), null);
+			jwtAuth.extractCredentials(
+				new FakeRestRequest(ImmutableMap.of("Authorization", TestJwts.MC_COY_SIGNED_OCT_1), new HashMap<>()),
+				null);
+		});
+
+		Assert.assertEquals("Authentication backend failed", exception.getMessage());
+		Assert.assertEquals(OpenSearchSecurityException.class, exception.getClass());
+	}
+
 	@Test
 	public void testEscapeKid() {
 		Settings settings = Settings.builder().put("openid_connect_url", mockIdpServer.getDiscoverUri()).build();

src/test/java/com/amazon/dlic/auth/http/jwt/keybyoidc/MockIpdServer.java

@@ -118,6 +118,10 @@ public String getDiscoverUri() {
 		return uri + CTX_DISCOVER;
 	}
 
+	public String getJwksUri() {
+		return uri + CTX_KEYS;
+	}
+
 	public int getPort() {
 		return port;
 	}