Merge branch 'main' into retention-policy-app

conf/openmetadata.yaml

@@ -180,6 +180,7 @@ authenticationConfiguration:
   # This will only be valid when provider type specified is customOidc
   providerName: ${CUSTOM_OIDC_AUTHENTICATION_PROVIDER_NAME:-""}
   publicKeyUrls: ${AUTHENTICATION_PUBLIC_KEYS:-[http://localhost:8585/api/v1/system/config/jwks]}
+  tokenValidationAlgorithm: ${AUTHENTICATION_TOKEN_VALIDATION_ALGORITHM:-"RS256"}
   authority: ${AUTHENTICATION_AUTHORITY:-https://accounts.google.com}
   clientId: ${AUTHENTICATION_CLIENT_ID:-""}
   callbackUrl: ${AUTHENTICATION_CALLBACK_URL:-""}

openmetadata-service/src/main/java/org/openmetadata/service/OpenMetadataApplication.java

@@ -190,7 +190,10 @@ public void run(OpenMetadataApplicationConfig catalogConfig, Environment environ
     EntityMaskerFactory.createEntityMasker();
 
     // Instantiate JWT Token Generator
-    JWTTokenGenerator.getInstance().init(catalogConfig.getJwtTokenConfiguration());
+    JWTTokenGenerator.getInstance()
+        .init(
+            catalogConfig.getAuthenticationConfiguration().getTokenValidationAlgorithm(),
+            catalogConfig.getJwtTokenConfiguration());
 
     // Set the Database type for choosing correct queries from annotations
     jdbi.getConfig(SqlObjects.class)

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

@@ -22,6 +22,7 @@
 import static org.openmetadata.service.security.SecurityUtil.validatePrincipalClaimsMapping;
 import static org.openmetadata.service.security.jwt.JWTTokenGenerator.ROLES_CLAIM;
 import static org.openmetadata.service.security.jwt.JWTTokenGenerator.TOKEN_TYPE;
+import static org.openmetadata.service.security.jwt.JWTTokenGenerator.getAlgorithm;
 
 import com.auth0.jwk.Jwk;
 import com.auth0.jwk.JwkProvider;
@@ -71,6 +72,7 @@ public class JwtFilter implements ContainerRequestFilter {
   private boolean enforcePrincipalDomain;
   private AuthProvider providerType;
   private boolean useRolesFromProvider = false;
+  private AuthenticationConfiguration.TokenValidationAlgorithm tokenValidationAlgorithm;
 
   private static final List<String> DEFAULT_PUBLIC_KEY_URLS =
       Arrays.asList(
@@ -123,6 +125,7 @@ public JwtFilter(
     this.principalDomain = authorizerConfiguration.getPrincipalDomain();
     this.enforcePrincipalDomain = authorizerConfiguration.getEnforcePrincipalDomain();
     this.useRolesFromProvider = authorizerConfiguration.getUseRolesFromProvider();
+    this.tokenValidationAlgorithm = authenticationConfiguration.getTokenValidationAlgorithm();
   }
 
   @VisibleForTesting
@@ -224,7 +227,8 @@ public Map<String, Claim> validateJwtAndGetClaims(String token) {
 
     // Validate JWT with public key
     Jwk jwk = jwkProvider.get(jwt.getKeyId());
-    Algorithm algorithm = Algorithm.RSA256((RSAPublicKey) jwk.getPublicKey(), null);
+    Algorithm algorithm =
+        getAlgorithm(tokenValidationAlgorithm, (RSAPublicKey) jwk.getPublicKey(), null);
     try {
       algorithm.verify(jwt);
     } catch (RuntimeException runtimeException) {

openmetadata-service/src/main/java/org/openmetadata/service/security/jwt/JWTTokenGenerator.java

@@ -37,6 +37,7 @@
 import java.util.Set;
 import lombok.Getter;
 import lombok.extern.slf4j.Slf4j;
+import org.openmetadata.schema.api.security.AuthenticationConfiguration;
 import org.openmetadata.schema.api.security.jwt.JWTTokenConfiguration;
 import org.openmetadata.schema.auth.JWTAuthMechanism;
 import org.openmetadata.schema.auth.JWTTokenExpiry;
@@ -56,6 +57,7 @@ public class JWTTokenGenerator {
   @Getter private RSAPublicKey publicKey;
   private String issuer;
   private String kid;
+  private AuthenticationConfiguration.TokenValidationAlgorithm tokenValidationAlgorithm;
 
   private JWTTokenGenerator() {
     /* Private constructor for singleton */
@@ -66,7 +68,9 @@ public static JWTTokenGenerator getInstance() {
   }
 
   /** Expected to be initialized only once during application start */
-  public void init(JWTTokenConfiguration jwtTokenConfiguration) {
+  public void init(
+      AuthenticationConfiguration.TokenValidationAlgorithm algorithm,
+      JWTTokenConfiguration jwtTokenConfiguration) {
     try {
       if (jwtTokenConfiguration.getRsaprivateKeyFilePath() != null
           && !jwtTokenConfiguration.getRsaprivateKeyFilePath().isEmpty()
@@ -84,6 +88,7 @@ public void init(JWTTokenConfiguration jwtTokenConfiguration) {
         publicKey = (RSAPublicKey) kf.generatePublic(spec);
         issuer = jwtTokenConfiguration.getJwtissuer();
         kid = jwtTokenConfiguration.getKeyId();
+        tokenValidationAlgorithm = algorithm;
       }
     } catch (Exception ex) {
       LOG.error("Failed to initialize JWTTokenGenerator ", ex);
@@ -141,7 +146,7 @@ public JWTAuthMechanism getJwtAuthMechanism(
         }
       }
       JWTAuthMechanism jwtAuthMechanism = new JWTAuthMechanism().withJWTTokenExpiry(expiry);
-      Algorithm algorithm = Algorithm.RSA256(null, privateKey);
+      Algorithm algorithm = getAlgorithm(tokenValidationAlgorithm, null, privateKey);
       String token =
           JWT.create()
               .withIssuer(issuer)
@@ -214,4 +219,15 @@ public Date getTokenExpiryFromJWT(String token) {
 
     return jwt.getExpiresAt();
   }
+
+  public static Algorithm getAlgorithm(
+      AuthenticationConfiguration.TokenValidationAlgorithm algorithm,
+      RSAPublicKey publicKey,
+      RSAPrivateKey privateKey) {
+    return switch (algorithm) {
+      case RS_256 -> Algorithm.RSA256(publicKey, privateKey);
+      case RS_384 -> Algorithm.RSA384(publicKey, privateKey);
+      case RS_512 -> Algorithm.RSA512(publicKey, privateKey);
+    };
+  }
 }

openmetadata-service/src/test/java/org/openmetadata/service/security/JWTTokenGeneratorTest.java

@@ -15,6 +15,7 @@
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.TestInstance;
+import org.openmetadata.schema.api.security.AuthenticationConfiguration;
 import org.openmetadata.schema.api.security.jwt.JWTTokenConfiguration;
 import org.openmetadata.schema.auth.JWTAuthMechanism;
 import org.openmetadata.schema.auth.JWTTokenExpiry;
@@ -38,7 +39,8 @@ public void setup() {
     jwtTokenConfiguration.setRsaprivateKeyFilePath(rsaPrivateKeyPath);
     jwtTokenConfiguration.setRsapublicKeyFilePath(rsaPublicKeyPath);
     jwtTokenGenerator = JWTTokenGenerator.getInstance();
-    jwtTokenGenerator.init(jwtTokenConfiguration);
+    jwtTokenGenerator.init(
+        AuthenticationConfiguration.TokenValidationAlgorithm.RS_256, jwtTokenConfiguration);
   }
 
   @Test

openmetadata-spec/src/main/resources/json/schema/configuration/authenticationConfiguration.json

@@ -46,6 +46,12 @@
         "type": "string"
       }
     },
+    "tokenValidationAlgorithm": {
+      "description": "Token Validation Algorithm to use.",
+      "type": "string",
+      "enum": ["RS256", "RS384", "RS512"],
+      "default": "RS256"
+    },
     "authority": {
       "description": "Authentication Authority",
       "type": "string"

openmetadata-ui/src/main/resources/ui/src/generated/configuration/authenticationConfiguration.ts

@@ -10,9 +10,7 @@
  *  See the License for the specific language governing permissions and
  *  limitations under the License.
  */
-
-
- /**
+/**
  * This schema defines the Authentication Configuration.
  */
 export interface AuthenticationConfiguration {
@@ -69,6 +67,10 @@ export interface AuthenticationConfiguration {
      * Saml Configuration that is applicable only when the provider is Saml
      */
     samlConfiguration?: SamlSSOClientConfig;
+    /**
+     * Token Validation Algorithm to use.
+     */
+    tokenValidationAlgorithm?: TokenValidationAlgorithm;
 }
 
 /**
@@ -492,3 +494,12 @@ export interface SP {
      */
     spX509Certificate?: string;
 }
+
+/**
+ * Token Validation Algorithm to use.
+ */
+export enum TokenValidationAlgorithm {
+    Rs256 = "RS256",
+    Rs384 = "RS384",
+    Rs512 = "RS512",
+}