Handle roles in nested claim for JWT auth backends

CHANGELOG.md

@@ -9,6 +9,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Github workflow for changelog verification ([#5318](https://github.com/opensearch-project/security/pull/5318))
 - Register cluster settings listener for `plugins.security.cache.ttl_minutes` ([#5324](https://github.com/opensearch-project/security/pull/5324))
 - Add flush cache endpoint for individual user ([#5337](https://github.com/opensearch-project/security/pull/5337))
+- Handle roles in nested claim for JWT auth backends ([#5355](https://github.com/opensearch-project/security/pull/5355))
 
 ### Changed
 - Use extendedPlugins in integrationTest framework for sample resource plugin testing ([#5322](https://github.com/opensearch-project/security/pull/5322))

src/integrationTest/java/org/opensearch/security/http/JwtAuthenticationNestedClaimsTests.java

@@ -0,0 +1,127 @@
+/*
+* Copyright OpenSearch Contributors
+* 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.
+*
+*/
+package org.opensearch.security.http;
+
+import java.security.KeyPair;
+import java.util.Arrays;
+import java.util.Base64;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import com.carrotsearch.randomizedtesting.annotations.ThreadLeakScope;
+import org.apache.hc.core5.http.Header;
+import org.junit.ClassRule;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import org.opensearch.test.framework.JwtConfigBuilder;
+import org.opensearch.test.framework.TestSecurityConfig;
+import org.opensearch.test.framework.cluster.ClusterManager;
+import org.opensearch.test.framework.cluster.LocalCluster;
+import org.opensearch.test.framework.cluster.TestRestClient;
+import org.opensearch.test.framework.cluster.TestRestClient.HttpResponse;
+import org.opensearch.test.framework.log.LogsRule;
+
+import io.jsonwebtoken.SignatureAlgorithm;
+import io.jsonwebtoken.security.Keys;
+
+import static java.nio.charset.StandardCharsets.US_ASCII;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.containsInAnyOrder;
+import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.hasSize;
+import static org.opensearch.security.http.JwtAuthenticationTests.POINTER_BACKEND_ROLES;
+import static org.opensearch.security.http.JwtAuthenticationTests.POINTER_USERNAME;
+import static org.opensearch.test.framework.TestSecurityConfig.AuthcDomain.BASIC_AUTH_DOMAIN_ORDER;
+
+@RunWith(com.carrotsearch.randomizedtesting.RandomizedRunner.class)
+@ThreadLeakScope(ThreadLeakScope.Scope.NONE)
+public class JwtAuthenticationNestedClaimsTests {
+
+    public static final String CLAIM_USERNAME = "preferred-username";
+    public static final List<String> CLAIM_ROLES = List.of("attributes", "roles");
+
+    public static final String USER_SUPERHERO = "superhero";
+    private static final KeyPair KEY_PAIR1 = Keys.keyPairFor(SignatureAlgorithm.RS256);
+    private static final String PUBLIC_KEY1 = new String(Base64.getEncoder().encode(KEY_PAIR1.getPublic().getEncoded()), US_ASCII);
+    private static final String JWT_AUTH_HEADER = "jwt-auth";
+
+    private static final JwtAuthorizationHeaderFactory tokenFactory1 = new JwtAuthorizationHeaderFactory(
+        KEY_PAIR1.getPrivate(),
+        CLAIM_USERNAME,
+        CLAIM_ROLES,
+        JWT_AUTH_HEADER
+    );
+    public static final TestSecurityConfig.AuthcDomain JWT_AUTH_DOMAIN = new TestSecurityConfig.AuthcDomain(
+        "jwt",
+        BASIC_AUTH_DOMAIN_ORDER - 1
+    ).jwtHttpAuthenticator(
+        new JwtConfigBuilder().jwtHeader(JWT_AUTH_HEADER).signingKey(List.of(PUBLIC_KEY1)).subjectKey(CLAIM_USERNAME).rolesKey(CLAIM_ROLES)
+    ).backend("noop");
+
+    @ClassRule
+    public static final LocalCluster cluster = new LocalCluster.Builder().clusterManager(ClusterManager.SINGLENODE)
+        .anonymousAuth(false)
+        .authc(JWT_AUTH_DOMAIN)
+        .build();
+
+    @Rule
+    public LogsRule logsRule = new LogsRule("org.opensearch.security.auth.http.jwt.HTTPJwtAuthenticator");
+
+    // TODO write tests for scenarios where roles are in nested claim. i.e. rolesKey: ['attributes', 'roles']
+    @Test
+    public void shouldAuthenticateWithNestedRolesClaim() {
+        // Create nested claims structure
+        Map<String, Object> attributes = new HashMap<>();
+        List<String> rolesClaim = Arrays.asList("all_access", "securitymanager");
+        attributes.put("roles", rolesClaim);
+
+        Map<String, Object> nestedClaims = new HashMap<>();
+        nestedClaims.put("attributes", attributes);
+
+        // Generate token with nested claims
+        Header header = tokenFactory1.generateValidTokenWithCustomClaims(USER_SUPERHERO, null, nestedClaims);
+
+        try (TestRestClient client = cluster.getRestClient(header)) {
+            HttpResponse response = client.getAuthInfo();
+
+            response.assertStatusCode(200);
+            String username = response.getTextFromJsonBody(POINTER_USERNAME);
+            assertThat(username, equalTo(USER_SUPERHERO));
+            List<String> roles = response.getTextArrayFromJsonBody(POINTER_BACKEND_ROLES);
+            assertThat(roles, hasSize(2));
+            assertThat(roles, containsInAnyOrder("all_access", "securitymanager"));
+        }
+    }
+
+    @Test
+    public void shouldHandleMissingNestedRolesClaim() {
+        // Create invalid nested claims structure
+        Map<String, Object> attributes = new HashMap<>();
+        attributes.put("wrong", "missing"); // Invalid format - should be a list
+
+        Map<String, Object> nestedClaims = new HashMap<>();
+        nestedClaims.put("attributes", attributes);
+
+        Header header = tokenFactory1.generateValidTokenWithCustomClaims(USER_SUPERHERO, null, nestedClaims);
+
+        try (TestRestClient client = cluster.getRestClient(header)) {
+            HttpResponse response = client.getAuthInfo();
+
+            response.assertStatusCode(200);
+            String username = response.getTextFromJsonBody(POINTER_USERNAME);
+            assertThat(username, equalTo(USER_SUPERHERO));
+            List<String> roles = response.getTextArrayFromJsonBody(POINTER_BACKEND_ROLES);
+            assertThat(roles, hasSize(0));
+        }
+    }
+}

src/integrationTest/java/org/opensearch/security/http/JwtAuthenticationTests.java

@@ -69,7 +69,7 @@
 public class JwtAuthenticationTests {
 
     public static final String CLAIM_USERNAME = "preferred-username";
-    public static final String CLAIM_ROLES = "backend-user-roles";
+    public static final List<String> CLAIM_ROLES = List.of("backend-user-roles");
 
     public static final String USER_SUPERHERO = "superhero";
     public static final String USERNAME_ROOT = "root";
@@ -305,5 +305,4 @@ public void secondKeypairShouldAuthenticateWithJwtToken_positiveWithAnotherUsern
             assertThat(username, equalTo(USERNAME_ROOT));
         }
     }
-
 }

src/integrationTest/java/org/opensearch/security/http/JwtAuthenticationWithUrlParamTests.java

@@ -51,7 +51,7 @@
 public class JwtAuthenticationWithUrlParamTests {
 
     public static final String CLAIM_USERNAME = "preferred-username";
-    public static final String CLAIM_ROLES = "backend-user-roles";
+    public static final List<String> CLAIM_ROLES = List.of("backend-user-roles");
     public static final String POINTER_USERNAME = "/user_name";
 
     private static final KeyPair KEY_PAIR = Keys.keyPairFor(SignatureAlgorithm.RS256);

src/integrationTest/java/org/opensearch/security/http/JwtAuthorizationHeaderFactory.java

@@ -10,11 +10,10 @@
 package org.opensearch.security.http;
 
 import java.security.PrivateKey;
-import java.util.Arrays;
 import java.util.Date;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
-import java.util.stream.Collectors;
 
 import com.google.common.collect.ImmutableMap;
 import org.apache.commons.lang3.StringUtils;
@@ -33,11 +32,11 @@ class JwtAuthorizationHeaderFactory {
 
     private final String usernameClaimName;
 
-    private final String rolesClaimName;
+    private final List<String> rolesClaimName;
 
     private final String headerName;
 
-    public JwtAuthorizationHeaderFactory(PrivateKey privateKey, String usernameClaimName, String rolesClaimName, String headerName) {
+    public JwtAuthorizationHeaderFactory(PrivateKey privateKey, String usernameClaimName, List<String> rolesClaimName, String headerName) {
         this.privateKey = requireNonNull(privateKey, "Private key is required");
         this.usernameClaimName = requireNonNull(usernameClaimName, "Username claim name is required");
         this.rolesClaimName = requireNonNull(rolesClaimName, "Roles claim name is required.");
@@ -64,8 +63,28 @@ private Map<String, Object> customClaimsMap(String username, String[] roles) {
         if (StringUtils.isNoneEmpty(username)) {
             builder.put(usernameClaimName, username);
         }
-        if ((roles != null) && (roles.length > 0)) {
-            builder.put(rolesClaimName, Arrays.stream(roles).collect(Collectors.joining(",")));
+        if (roles != null && roles.length > 0) {
+            if (rolesClaimName.size() == 1) {
+                // Simple case - no nesting
+                builder.put(rolesClaimName.get(0), String.join(",", roles));
+            } else {
+                // Handle nested claims
+                Map<String, Object> nestedMap = new HashMap<>();
+                Map<String, Object> currentMap = nestedMap;
+
+                // Build the nested structure
+                for (int i = 0; i < rolesClaimName.size() - 1; i++) {
+                    Map<String, Object> nextMap = new HashMap<>();
+                    currentMap.put(rolesClaimName.get(i), nextMap);
+                    currentMap = nextMap;
+                }
+
+                // Add the roles array at the deepest level
+                currentMap.put(rolesClaimName.get(rolesClaimName.size() - 1), String.join(",", roles));
+
+                // Add the entire nested structure to the builder
+                builder.putAll(nestedMap);
+            }
         }
         return builder.build();
     }

src/integrationTest/java/org/opensearch/test/framework/JwtConfigBuilder.java

@@ -22,7 +22,7 @@ public class JwtConfigBuilder {
     private String jwtUrlParameter;
     private List<String> signingKeys;
     private String subjectKey;
-    private String rolesKey;
+    private List<String> rolesKey;
 
     public JwtConfigBuilder jwtHeader(String jwtHeader) {
         this.jwtHeader = jwtHeader;
@@ -45,6 +45,11 @@ public JwtConfigBuilder subjectKey(String subjectKey) {
     }
 
     public JwtConfigBuilder rolesKey(String rolesKey) {
+        this.rolesKey = List.of(rolesKey);
+        return this;
+    }
+
+    public JwtConfigBuilder rolesKey(List<String> rolesKey) {
         this.rolesKey = rolesKey;
         return this;
     }
@@ -64,7 +69,7 @@ public Map<String, Object> build() {
         if (isNoneBlank(subjectKey)) {
             builder.put("subject_key", subjectKey);
         }
-        if (isNoneBlank(rolesKey)) {
+        if (rolesKey != null && !rolesKey.isEmpty()) {
             builder.put("roles_key", rolesKey);
         }
         return builder.build();

src/main/java/org/opensearch/security/auth/http/jwt/AbstractHTTPJwtAuthenticator.java

@@ -61,7 +61,7 @@
     private final boolean isDefaultAuthHeader;
     private final String jwtUrlParameter;
     private final String subjectKey;
-    private final String rolesKey;
+    private final List<String> rolesKey;
     private final List<String> requiredAudience;
     private final String requiredIssuer;
 
@@ -72,7 +72,7 @@
         jwtUrlParameter = settings.get("jwt_url_parameter");
         jwtHeaderName = settings.get("jwt_header", AUTHORIZATION);
         isDefaultAuthHeader = AUTHORIZATION.equalsIgnoreCase(jwtHeaderName);
-        rolesKey = settings.get("roles_key");
+        rolesKey = settings.getAsList("roles_key");
         subjectKey = settings.get("subject_key");
         clockSkewToleranceSeconds = settings.getAsInt("jwt_clock_skew_tolerance_seconds", DEFAULT_CLOCK_SKEW_TOLERANCE_SECONDS);
         requiredAudience = settings.getAsList("required_audience");
@@ -219,7 +219,21 @@
             return new String[0];
         }
 
-        Object rolesObject = claims.getClaim(rolesKey);
+        Object rolesObject = null;
+        Map<String, Object> claimsMap = claims.getClaims();
+        for (int i = 0; i < rolesKey.size(); i++) {
+            if (i == rolesKey.size() - 1) {
+                rolesObject = claimsMap.get(rolesKey.get(i));
+            } else if (claimsMap.get(rolesKey.get(i)) instanceof Map) {
+                claimsMap = (Map<String, Object>) claimsMap.get(rolesKey.get(i));
+            } else {
+                log.warn(
+                    "Failed to get roles from JWT claims with roles_key '{}'. Check if this key is correct and available in the JWT payload.",
+                    rolesKey
+                );
+                return new String[0];
+            }
+        }
 
         if (rolesObject == null) {
             log.warn(