Add config for using multiple JWKs providers (#3738)

catalog-rest-service/src/main/java/org/openmetadata/catalog/security/AuthenticationConfiguration.java

@@ -19,7 +19,14 @@ import lombok.Setter;
 
 public class AuthenticationConfiguration {
   @Getter @Setter private String provider;
-  @Getter @Setter private String publicKey;
+
+  /** @deprecated Use publicKeyUrls */
+  @Deprecated(since = "0.9.1", forRemoval = true)
+  @Getter
+  @Setter
+  private String publicKey;
+
+  @Getter @Setter private List<String> publicKeyUrls;
   @Getter @Setter private String authority;
   @Getter @Setter private String clientId;
   @Getter @Setter private String callbackUrl;
@@ -31,9 +38,8 @@ public class AuthenticationConfiguration {
         + "provider='"
         + provider
         + '\''
-        + ", publicKey='"
-        + publicKey
-        + '\''
+        + ", publicKeyUrls="
+        + publicKeyUrls
         + ", authority='"
         + authority
         + '\''

catalog-rest-service/src/main/java/org/openmetadata/catalog/security/JwtFilter.java

@@ -14,13 +14,16 @@
 package org.openmetadata.catalog.security;
 
 import com.auth0.jwk.Jwk;
-import com.auth0.jwk.UrlJwkProvider;
+import com.auth0.jwk.JwkProvider;
 import com.auth0.jwt.JWT;
 import com.auth0.jwt.algorithms.Algorithm;
+import com.auth0.jwt.exceptions.JWTDecodeException;
 import com.auth0.jwt.interfaces.DecodedJWT;
 import com.fasterxml.jackson.databind.node.TextNode;
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.collect.ImmutableList;
 import io.dropwizard.util.Strings;
-import java.net.URI;
+import java.net.URL;
 import java.security.interfaces.RSAPublicKey;
 import java.util.Calendar;
 import java.util.List;
@@ -40,15 +43,28 @@ import org.openmetadata.catalog.security.auth.CatalogSecurityContext;
 public class JwtFilter implements ContainerRequestFilter {
   public static final String AUTHORIZATION_HEADER = "Authorization";
   public static final String TOKEN_PREFIX = "Bearer";
-  private String publicKeyUri;
+
   private List<String> jwtPrincipalClaims;
+  private JwkProvider jwkProvider;
 
   @SuppressWarnings("unused")
   private JwtFilter() {}
 
+  @SneakyThrows
   public JwtFilter(AuthenticationConfiguration authenticationConfiguration) {
-    this.publicKeyUri = authenticationConfiguration.getPublicKey();
     this.jwtPrincipalClaims = authenticationConfiguration.getJwtPrincipalClaims();
+
+    ImmutableList.Builder<URL> publicKeyUrlsBuilder = ImmutableList.builder();
+    for (String publicKeyUrlStr : authenticationConfiguration.getPublicKeyUrls()) {
+      publicKeyUrlsBuilder.add(new URL(publicKeyUrlStr));
+    }
+    this.jwkProvider = new MultiUrlJwkProvider(publicKeyUrlsBuilder.build());
+  }
+
+  @VisibleForTesting
+  JwtFilter(JwkProvider jwkProvider, List<String> jwtPrincipalClaims) {
+    this.jwkProvider = jwkProvider;
+    this.jwtPrincipalClaims = jwtPrincipalClaims;
   }
 
   @SneakyThrows
@@ -65,7 +81,12 @@ public class JwtFilter implements ContainerRequestFilter {
     LOG.debug("Token from header:{}", tokenFromHeader);
 
     // Decode JWT Token
-    DecodedJWT jwt = JWT.decode(tokenFromHeader);
+    DecodedJWT jwt;
+    try {
+      jwt = JWT.decode(tokenFromHeader);
+    } catch (JWTDecodeException e) {
+      throw new AuthenticationException("Invalid token", e);
+    }
 
     // Check if expired
     if (jwt.getExpiresAt().before(Calendar.getInstance(TimeZone.getTimeZone("UTC")).getTime())) {
@@ -73,9 +94,7 @@ public class JwtFilter implements ContainerRequestFilter {
     }
 
     // Validate JWT with public key
-    final URI uri = new URI(publicKeyUri).normalize();
-    UrlJwkProvider urlJwkProvider = new UrlJwkProvider(uri.toURL());
-    Jwk jwk = urlJwkProvider.get(jwt.getKeyId());
+    Jwk jwk = jwkProvider.get(jwt.getKeyId());
     Algorithm algorithm = Algorithm.RSA256((RSAPublicKey) jwk.getPublicKey(), null);
     try {
       algorithm.verify(jwt);

catalog-rest-service/src/main/java/org/openmetadata/catalog/security/MultiUrlJwkProvider.java

@@ -0,0 +1,44 @@
+/*
+ *  Copyright 2021 Collate
+ *  Licensed under the Apache License, Version 2.0 (the "License");
+ *  you may not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *  http://www.apache.org/licenses/LICENSE-2.0
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+
+package org.openmetadata.catalog.security;
+
+import com.auth0.jwk.Jwk;
+import com.auth0.jwk.JwkException;
+import com.auth0.jwk.JwkProvider;
+import com.auth0.jwk.SigningKeyNotFoundException;
+import com.auth0.jwk.UrlJwkProvider;
+import java.net.URL;
+import java.util.List;
+import java.util.stream.Collectors;
+
+final class MultiUrlJwkProvider implements JwkProvider {
+  private final List<UrlJwkProvider> urlJwkProviders;
+
+  public MultiUrlJwkProvider(List<URL> publicKeyUris) {
+    this.urlJwkProviders = publicKeyUris.stream().map(UrlJwkProvider::new).collect(Collectors.toUnmodifiableList());
+  }
+
+  @Override
+  public Jwk get(String keyId) throws JwkException {
+    JwkException lastException = new SigningKeyNotFoundException("No key found in with kid " + keyId, null);
+    for (UrlJwkProvider jwkProvider : urlJwkProviders) {
+      try {
+        return jwkProvider.get(keyId);
+      } catch (JwkException e) {
+        lastException.addSuppressed(e);
+      }
+    }
+    throw lastException;
+  }
+}

catalog-rest-service/src/test/java/org/openmetadata/catalog/security/JwtFilterTest.java

@@ -0,0 +1,201 @@
+/*
+ *  Copyright 2021 Collate
+ *  Licensed under the Apache License, Version 2.0 (the "License");
+ *  you may not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *  http://www.apache.org/licenses/LICENSE-2.0
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+
+package org.openmetadata.catalog.security;
+
+import static java.lang.String.format;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import com.auth0.jwk.Jwk;
+import com.auth0.jwk.JwkProvider;
+import com.auth0.jwt.JWT;
+import com.auth0.jwt.algorithms.Algorithm;
+import java.net.URI;
+import java.security.KeyPair;
+import java.security.KeyPairGenerator;
+import java.security.NoSuchAlgorithmException;
+import java.security.interfaces.RSAPrivateKey;
+import java.security.interfaces.RSAPublicKey;
+import java.time.Instant;
+import java.time.temporal.ChronoUnit;
+import java.util.Date;
+import java.util.List;
+import java.util.Locale;
+import java.util.Map;
+import javax.ws.rs.container.ContainerRequestContext;
+import javax.ws.rs.core.MultivaluedHashMap;
+import javax.ws.rs.core.SecurityContext;
+import javax.ws.rs.core.UriInfo;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+import org.mockito.ArgumentCaptor;
+
+class JwtFilterTest {
+
+  private static JwtFilter jwtFilter;
+
+  private static Algorithm algorithm;
+  private static UriInfo mockRequestURIInfo;
+
+  @BeforeAll
+  static void before() throws Exception {
+    // Create a RSA256 algorithm wth random public/private key pair
+    KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance("RSA");
+    keyPairGenerator.initialize(512);
+    KeyPair keyPair = keyPairGenerator.generateKeyPair();
+    algorithm = Algorithm.RSA256((RSAPublicKey) keyPair.getPublic(), (RSAPrivateKey) keyPair.getPrivate());
+
+    // Mock a JwkProvider that has a single JWK containing the public key from the algorithm above
+    // This is used to verify the JWT
+    Jwk mockJwk = mock(Jwk.class);
+    when(mockJwk.getPublicKey()).thenReturn(keyPair.getPublic());
+    JwkProvider jwkProvider = mock(JwkProvider.class);
+    when(jwkProvider.get(algorithm.getSigningKeyId())).thenReturn(mockJwk);
+
+    // This is needed by JwtFilter for some metadata, not very important
+    URI uri = URI.create("POST:http://localhost:8080/login");
+    mockRequestURIInfo = mock(UriInfo.class);
+    when(mockRequestURIInfo.getPath()).thenReturn("/login");
+    when(mockRequestURIInfo.getRequestUri()).thenReturn(uri);
+
+    List<String> principalClaims = List.of("sub", "email");
+    jwtFilter = new JwtFilter(jwkProvider, principalClaims);
+  }
+
+  @Test
+  void testSuccessfulFilter() {
+    String jwt =
+        JWT.create()
+            .withExpiresAt(Date.from(Instant.now().plus(1, ChronoUnit.DAYS)))
+            .withClaim("sub", "sam")
+            .sign(algorithm);
+
+    ContainerRequestContext context = createRequestContextWithJwt(jwt);
+
+    jwtFilter.filter(context);
+
+    ArgumentCaptor<SecurityContext> securityContextArgument = ArgumentCaptor.forClass(SecurityContext.class);
+    verify(context, times(1)).setSecurityContext(securityContextArgument.capture());
+
+    assertEquals("sam", securityContextArgument.getValue().getUserPrincipal().getName());
+  }
+
+  @Test
+  void testFilterWithEmailClaim() {
+    String jwt =
+        JWT.create()
+            .withExpiresAt(Date.from(Instant.now().plus(1, ChronoUnit.DAYS)))
+            .withClaim("email", "sam@gmail.com")
+            .sign(algorithm);
+
+    ContainerRequestContext context = createRequestContextWithJwt(jwt);
+
+    jwtFilter.filter(context);
+
+    ArgumentCaptor<SecurityContext> securityContextArgument = ArgumentCaptor.forClass(SecurityContext.class);
+    verify(context, times(1)).setSecurityContext(securityContextArgument.capture());
+
+    assertEquals("sam", securityContextArgument.getValue().getUserPrincipal().getName());
+  }
+
+  @Test
+  void testMissingToken() {
+    MultivaluedHashMap<String, String> headers = new MultivaluedHashMap<>();
+    ContainerRequestContext context = mock(ContainerRequestContext.class);
+    when(context.getUriInfo()).thenReturn(mockRequestURIInfo);
+    when(context.getHeaders()).thenReturn(headers);
+
+    Exception exception = assertThrows(AuthenticationException.class, () -> jwtFilter.filter(context));
+    assertTrue(exception.getMessage().toLowerCase(Locale.ROOT).contains("token not present"));
+  }
+
+  @Test
+  void testInvalidToken() {
+    ContainerRequestContext context = createRequestContextWithJwt("invalid-token");
+
+    Exception exception = assertThrows(AuthenticationException.class, () -> jwtFilter.filter(context));
+    assertTrue(exception.getMessage().toLowerCase(Locale.ROOT).contains("invalid token"));
+  }
+
+  @Test
+  void testExpiredToken() {
+    String jwt =
+        JWT.create()
+            .withExpiresAt(Date.from(Instant.now().minus(1, ChronoUnit.DAYS)))
+            .withClaim("sub", "sam")
+            .sign(algorithm);
+
+    ContainerRequestContext context = createRequestContextWithJwt(jwt);
+
+    Exception exception = assertThrows(AuthenticationException.class, () -> jwtFilter.filter(context));
+    assertTrue(exception.getMessage().toLowerCase(Locale.ROOT).contains("expired"));
+  }
+
+  @Test
+  void testNoClaimsInToken() {
+    String jwt =
+        JWT.create()
+            .withExpiresAt(Date.from(Instant.now().plus(1, ChronoUnit.DAYS)))
+            .withClaim("emailAddress", "sam@gmail.com")
+            .sign(algorithm);
+
+    ContainerRequestContext context = createRequestContextWithJwt(jwt);
+
+    Exception exception = assertThrows(AuthenticationException.class, () -> jwtFilter.filter(context));
+    assertTrue(exception.getMessage().toLowerCase(Locale.ROOT).contains("claim"));
+  }
+
+  @Test
+  void testInvalidSignatureJwt() throws NoSuchAlgorithmException {
+    KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance("RSA");
+    keyPairGenerator.initialize(512);
+    KeyPair keyPair = keyPairGenerator.generateKeyPair();
+    Algorithm secondaryAlgorithm =
+        Algorithm.RSA256((RSAPublicKey) keyPair.getPublic(), (RSAPrivateKey) keyPair.getPrivate());
+
+    String jwt =
+        JWT.create()
+            .withExpiresAt(Date.from(Instant.now().plus(1, ChronoUnit.DAYS)))
+            .withClaim("sub", "sam")
+            .sign(secondaryAlgorithm);
+
+    ContainerRequestContext context = createRequestContextWithJwt(jwt);
+
+    Exception exception = assertThrows(AuthenticationException.class, () -> jwtFilter.filter(context));
+    assertTrue(exception.getMessage().toLowerCase(Locale.ROOT).contains("invalid token"));
+  }
+
+  /**
+   * Creates the ContainerRequestsContext that is passed to the filter. This object can be quite complex, but the
+   * JwtFilter cares only about the Authorization header and request URI.
+   *
+   * @param jwt JWT in string format to be added to headers
+   * @return Mocked ContainerRequestContext with an Authorization header and request URI info
+   */
+  private static ContainerRequestContext createRequestContextWithJwt(String jwt) {
+    MultivaluedHashMap<String, String> headers =
+        new MultivaluedHashMap<>(Map.of(JwtFilter.AUTHORIZATION_HEADER, format("%s %s", JwtFilter.TOKEN_PREFIX, jwt)));
+
+    ContainerRequestContext context = mock(ContainerRequestContext.class);
+    when(context.getUriInfo()).thenReturn(mockRequestURIInfo);
+    when(context.getHeaders()).thenReturn(headers);
+
+    return context;
+  }
+}

deploy/deploy-on-bare-metal/enable-security/auth0-sso/auth0-config.md

@@ -9,7 +9,8 @@
 ```
 authenticationConfiguration:
   provider: "auth0"
-  publicKey: "https://parth-panchal.us.auth0.com/.well-known/jwks.json"
+  publicKeyUrls:
+    - "https://parth-panchal.us.auth0.com/.well-known/jwks.json"
   authority: "https://parth-panchal.us.auth0.com/"
   clientId: "{Client ID}"
   callbackUrl: "http://localhost:8585/callback"

deploy/deploy-on-bare-metal/enable-security/google-sso/google-config.md

@@ -7,7 +7,8 @@ Once the `client id` and `client secret` are generated, add `client id` as the v
 ```
 authenticationConfiguration:
   provider: "google"
-  publicKey: "https://www.googleapis.com/oauth2/v3/certs"
+  publicKeyUrls:
+    - "https://www.googleapis.com/oauth2/v3/certs"
   authority: "https://accounts.google.com"
   clientId: "{client id}"
   callbackUrl: "http://localhost:8585/callback"

deploy/deploy-on-bare-metal/enable-security/okta-sso/okta-config.md

@@ -7,7 +7,8 @@ Once the **Client Id**, and **Issuer URL** are generated, add those details in `
 ```yaml
 authenticationConfiguration:
   provider: "okta"
-  publicKey: "{ISSUER_URL}/v1/keys"
+  publicKeyUrls:
+    - "{ISSUER_URL}/v1/keys"
   authority: "{ISSUER_URL}"
   clientId: "{CLIENT_ID - SPA APP}"
   callbackUrl: "http://localhost:8585/callback"

deploy/deploy-on-bare-metal/enable-security/okta-sso/okta-server-creds.md

@@ -42,7 +42,8 @@ description: >-
     * **Refresh Token** - For the refresh token behavior, it is recommended to select the option to 'Rotate token after every use'.
     * **Implicit (hybrid)** - Select the options to allow ID Token and Access Token with implicit grant type.
   * Enter the **Sign-in redirect URIs**
-    * [http://localhost:8585/signin
      \
+    * [http://localhost:8585/signin
+      \
       http://localhost:8585](http://localhost:8585/signinhttp://localhost:8585)
   * Enter the **Sign-out redirect URIs**
   * Enter the **Base URIs**
@@ -147,7 +148,8 @@ authorizerConfiguration:
 
 authenticationConfiguration:
   provider: "okta"
-  publicKey: "{ISSUER_URL}/v1/keys"
+  publicKeyUrls:
+    - "{ISSUER_URL}/v1/keys"
   authority: "{ISSUER_URL}"
   clientId: "{CLIENT_ID - SPA APP}"
   callbackUrl: "http://localhost:8585/callback"

deploy/deploy-on-kubernetes/enable-security/auth0-sso/auth0-config.md

@@ -9,7 +9,8 @@
 ```
 authenticationConfiguration:
   provider: "auth0"
-  publicKey: "https://parth-panchal.us.auth0.com/.well-known/jwks.json"
+  publicKeyUrls:
+    - "https://parth-panchal.us.auth0.com/.well-known/jwks.json"
   authority: "https://parth-panchal.us.auth0.com/"
   clientId: "{Client ID}"
   callbackUrl: "http://localhost:8585/callback"

deploy/deploy-on-kubernetes/enable-security/google-sso/google-config.md

@@ -7,7 +7,8 @@ Once the `client id` and `client secret` are generated, add `client id` as the v
 ```
 authenticationConfiguration:
   provider: "google"
-  publicKey: "https://www.googleapis.com/oauth2/v3/certs"
+  publicKeyUrls:
+    - "https://www.googleapis.com/oauth2/v3/certs"
   authority: "https://accounts.google.com"
   clientId: "{client id}"
   callbackUrl: "http://localhost:8585/callback"

deploy/deploy-on-kubernetes/enable-security/okta-sso/okta-config.md

@@ -7,7 +7,8 @@ Once the **Client Id**, and **Issuer URL** are generated, add those details in `
 ```yaml
 authenticationConfiguration:
   provider: "okta"
-  publicKey: "{ISSUER_URL}/v1/keys"
+  publicKeyUrls:
+    - "{ISSUER_URL}/v1/keys"
   authority: "{ISSUER_URL}"
   clientId: "{CLIENT_ID - SPA APP}"
   callbackUrl: "http://localhost:8585/callback"

deploy/deploy-on-kubernetes/enable-security/okta-sso/okta-server-creds.md

@@ -42,7 +42,8 @@ description: >-
     * **Refresh Token** - For the refresh token behavior, it is recommended to select the option to 'Rotate token after every use'.
     * **Implicit (hybrid)** - Select the options to allow ID Token and Access Token with implicit grant type.
   * Enter the **Sign-in redirect URIs**
-    * [http://localhost:8585/signin
      \
+    * [http://localhost:8585/signin
+      \
       http://localhost:8585](http://localhost:8585/signinhttp://localhost:8585)
   * Enter the **Sign-out redirect URIs**
   * Enter the **Base URIs**
@@ -147,7 +148,8 @@ authorizerConfiguration:
 
 authenticationConfiguration:
   provider: "okta"
-  publicKey: "{ISSUER_URL}/v1/keys"
+  publicKeyUrls:
+    - "{ISSUER_URL}/v1/keys"
   authority: "{ISSUER_URL}"
   clientId: "{CLIENT_ID - SPA APP}"
   callbackUrl: "http://localhost:8585/callback"

docs/install/README.md

@@ -49,7 +49,8 @@ Enter following information in ***/conf/openmetadata-security.yaml*** file:
 authorizerConfiguration:
   className: <authorizer_classname>
   containerRequestFilter: <JWT-filter>
-  publicKeyUri: <sign-on_provider_public-key>
+  publicKeyUrls:
+    - <sign-on_provider_public-key>
   clientAuthorizer:
     authority: <sign-on_issuer-url>
     client_id: <sign-on_client_id>

docs/install/enable-security/google-sso/google-config.md

@@ -7,7 +7,8 @@
 ```
 authenticationConfiguration:
   provider: "google"
-  publicKey: "https://www.googleapis.com/oauth2/v3/certs"
+  publicKeyUrls:
+    - "https://www.googleapis.com/oauth2/v3/certs"
   authority: "https://accounts.google.com"
   clientId: "{client id}"
   callbackUrl: "http://localhost:8585/callback"

docs/install/enable-security/okta-sso/okta-config.md

@@ -7,7 +7,8 @@
 ```
 authenticationConfiguration:
   provider: "okta"
-  publicKey: "https://{okta_domain}/oauth2/default/v1/keys"
+  publicKeyUrls:
+    - "https://{okta_domain}/oauth2/default/v1/keys"
   authority: "{okta_domain}"
   clientId: "{Client Secret}"
   callbackUrl: "http://localhost:8585/callback"