Role mapping support in pac4j extension

docs/development/extensions-core/druid-pac4j.md

@@ -55,6 +55,7 @@ druid.auth.authenticator.jwt.type=jwt
 |`druid.auth.pac4j.oidc.discoveryURI`|discovery URI for fetching OP metadata [see this](http://openid.net/specs/openid-connect-discovery-1_0.html).|none|Yes|
 |`druid.auth.pac4j.oidc.oidcClaim`|[claim](https://openid.net/specs/openid-connect-core-1_0.html#Claims) that will be extracted from the ID Token after validation.|name|No|
 |`druid.auth.pac4j.oidc.scope`| scope is used by an application during authentication to authorize access to a user's details.|`openid profile email`|No|
+|`druid.auth.pac4j.oidc.roleClaimPath`| Dot-separated path to the claim containing user roles|none|No|
 
 :::info
 Users must set a strong passphrase to ensure that an attacker is not able to guess it simply by brute force.

extensions-core/druid-basic-security/src/main/java/org/apache/druid/security/basic/authorization/LDAPRoleProvider.java

@@ -106,14 +106,24 @@ public Set<String> getRoles(String authorizerPrefix, AuthenticationResult authen
       }
     }
 
+    Set<String> claims = RoleProviderUtil.claimValuesFromCtx(authenticationResult.getContext());
+
     // Get the roles assigned to LDAP user from the metastore.
-    // This allow us to authorize LDAP users regardless of whether they belong to any groups or not in LDAP.  
-    BasicAuthorizerUser user = userMap.get(authenticationResult.getIdentity());
-    if (user != null) {
-      roleNames.addAll(user.getRoles());
+    // This allow us to authorize LDAP users regardless of whether they belong to any groups or not in LDAP.
+    if (claims != null) {
+      return RoleProviderUtil.getRolesByClaimValue(
+          authorizerPrefix,
+          claims,
+          roleNames,
+          cacheManager
+      );
+    } else {
+      return RoleProviderUtil.getRolesByIdentity(
+          userMap,
+          authenticationResult.getIdentity(),
+          roleNames
+      );
     }
-
-    return roleNames;
   }
 
   @Override

extensions-core/druid-basic-security/src/main/java/org/apache/druid/security/basic/authorization/MetadataStoreRoleProvider.java

@@ -57,11 +57,22 @@ public Set<String> getRoles(String authorizerPrefix, AuthenticationResult authen
       throw new IAE("Could not load userMap for authorizer [%s]", authorizerPrefix);
     }
 
-    BasicAuthorizerUser user = userMap.get(authenticationResult.getIdentity());
-    if (user != null) {
-      roleNames.addAll(user.getRoles());
+    Set<String> claims = RoleProviderUtil.claimValuesFromCtx(authenticationResult.getContext());
+
+    if (claims != null) {
+      return RoleProviderUtil.getRolesByClaimValue(
+          authorizerPrefix,
+          claims,
+          roleNames,
+          cacheManager
+      );
+    } else {
+      return RoleProviderUtil.getRolesByIdentity(
+          userMap,
+          authenticationResult.getIdentity(),
+          roleNames
+      );
     }
-    return roleNames;
   }
 
   @Override

extensions-core/druid-basic-security/src/main/java/org/apache/druid/security/basic/authorization/RoleProviderUtil.java

@@ -0,0 +1,90 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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.apache.druid.security.basic.authorization;
+
+import org.apache.druid.security.basic.authorization.db.cache.BasicAuthorizerCacheManager;
+import org.apache.druid.security.basic.authorization.entity.BasicAuthorizerRole;
+import org.apache.druid.security.basic.authorization.entity.BasicAuthorizerUser;
+
+import javax.annotation.Nullable;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+
+public class RoleProviderUtil
+{
+  public static final String ROLE_CLAIM_CONTEXT_KEY = "druidRoles";
+
+  public static Set<String> getRolesByIdentity(
+      Map<String, BasicAuthorizerUser> userMap,
+      String identity,
+      Set<String> roleNames
+  )
+  {
+    BasicAuthorizerUser user = userMap.get(identity);
+    if (user != null) {
+      roleNames.addAll(user.getRoles());
+    }
+    return roleNames;
+  }
+
+  public static Set<String> getRolesByClaimValue(
+      String authorizerPrefix,
+      Set<String> claimValue,
+      Set<String> roleNames,
+      BasicAuthorizerCacheManager cacheManager
+  )
+  {
+    Map<String, BasicAuthorizerRole> roleMap = cacheManager.getRoleMap(authorizerPrefix);
+
+    if (roleMap == null) {
+      return Set.of();
+    }
+
+    roleMap.keySet()
+           .stream()
+           .filter(claimValue::contains)
+           .forEach(roleNames::add);
+
+    return roleNames;
+  }
+
+  @Nullable
+  protected static Set<String> claimValuesFromCtx(Map<String, Object> ctx)
+  {
+    Object value = (ctx == null) ? null : ctx.get(RoleProviderUtil.ROLE_CLAIM_CONTEXT_KEY);
+    if (!(value instanceof Set)) {
+      return null;
+    }
+    Set<?> rawClaimValues = (Set<?>) value;
+
+    Set<String> result = new HashSet<>();
+    for (Object claimValue : rawClaimValues) {
+      if (!(claimValue instanceof String)) {
+        return null;
+      }
+      String str = ((String) claimValue).trim();
+      if (!str.isEmpty()) {
+        result.add(str);
+      }
+    }
+    return result.isEmpty() ? null : result;
+  }
+}

extensions-core/druid-basic-security/src/test/java/org/apache/druid/security/authorization/MetadataStoreRoleProviderGetRolesTest.java

@@ -0,0 +1,87 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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.apache.druid.security.authorization;
+
+import org.apache.druid.security.basic.authorization.MetadataStoreRoleProvider;
+import org.apache.druid.security.basic.authorization.RoleProviderUtil;
+import org.apache.druid.security.basic.authorization.db.cache.BasicAuthorizerCacheManager;
+import org.apache.druid.security.basic.authorization.entity.BasicAuthorizerRole;
+import org.apache.druid.security.basic.authorization.entity.BasicAuthorizerUser;
+import org.apache.druid.server.security.AuthenticationResult;
+import org.junit.Test;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Set;
+
+import static org.junit.Assert.assertEquals;
+
+public class MetadataStoreRoleProviderGetRolesTest
+{
+
+  @Test
+  public void returnsRolesByClaimValuesWhenPresent()
+  {
+    Map<String, BasicAuthorizerRole> roles = new HashMap<>();
+    roles.put("admin", null);
+    roles.put("viewer", null);
+
+    Set<String> viewerRole = Set.of("viewer");
+
+    BasicAuthorizerUser user = new BasicAuthorizerUser("alice", viewerRole);
+
+    Map<String, BasicAuthorizerUser> users = Map.of("alice", user);
+
+    BasicAuthorizerCacheManager cache = new StubCacheManager(users, roles);
+    MetadataStoreRoleProvider provider = new MetadataStoreRoleProvider(cache);
+
+    Set<String> claims = Set.of("admin", "extraneous");
+
+    Map<String, Object> ctx = Map.of(RoleProviderUtil.ROLE_CLAIM_CONTEXT_KEY, claims);
+
+    AuthenticationResult ar = new AuthenticationResult("alice", "basic", "pac4j", ctx);
+
+    Set<String> out = provider.getRoles("basic", ar);
+    Set<String> expected = Set.of("admin");
+    assertEquals(expected, out);
+  }
+
+  @Test
+  public void fallsBackToIdentityWhenNoClaimContext()
+  {
+    Set<String> viewerRole = Set.of("viewer");
+    BasicAuthorizerUser user = new BasicAuthorizerUser("alice", viewerRole);
+
+    Map<String, BasicAuthorizerUser> users = Map.of("alice", user);
+
+    Map<String, BasicAuthorizerRole> roles = new HashMap<>();
+    roles.put("admin", null);
+
+    BasicAuthorizerCacheManager cache = new StubCacheManager(users, roles);
+    MetadataStoreRoleProvider provider = new MetadataStoreRoleProvider(cache);
+
+    AuthenticationResult ar = new AuthenticationResult("alice", "basic", "pac4j", Collections.emptyMap());
+
+    Set<String> out = provider.getRoles("basic", ar);
+    Set<String> expected = Set.of("viewer");
+    assertEquals(expected, out);
+  }
+}

extensions-core/druid-basic-security/src/test/java/org/apache/druid/security/authorization/RoleProviderUtilTest.java

@@ -0,0 +1,82 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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.apache.druid.security.authorization;
+
+import org.apache.druid.security.basic.authorization.RoleProviderUtil;
+import org.apache.druid.security.basic.authorization.db.cache.BasicAuthorizerCacheManager;
+import org.apache.druid.security.basic.authorization.entity.BasicAuthorizerRole;
+import org.apache.druid.security.basic.authorization.entity.BasicAuthorizerUser;
+import org.junit.Test;
+
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+
+public class RoleProviderUtilTest
+{
+
+  @Test
+  public void getRolesByIdentityAddsRolesWhenUserFound()
+  {
+    Set<String> roles = Set.of("r1", "r2");
+    BasicAuthorizerUser user = new BasicAuthorizerUser("id", roles);
+
+    Map<String, BasicAuthorizerUser> userMap = Map.of("id", user);
+
+    Set<String> out = RoleProviderUtil.getRolesByIdentity(userMap, "id", new HashSet<>());
+    assertEquals(roles, out);
+  }
+
+  @Test
+  public void getRolesByIdentityNoopWhenUserMissing()
+  {
+    Map<String, BasicAuthorizerUser> userMap = Map.of();
+    Set<String> out = RoleProviderUtil.getRolesByIdentity(userMap, "missing", new HashSet<>());
+    assertTrue(out.isEmpty());
+  }
+
+  @Test
+  public void getRolesByClaimValuesFiltersByRoleNames()
+  {
+    Map<String, BasicAuthorizerRole> roles = new HashMap<>();
+    roles.put("r1", null);
+    roles.put("r2", null);
+
+    BasicAuthorizerCacheManager cache = new StubCacheManager(Map.of(), roles);
+
+    Set<String> claims = Set.of("r2", "nope");
+    Set<String> out = RoleProviderUtil.getRolesByClaimValue("authz", claims, new HashSet<>(), cache);
+    assertEquals(Set.of("r2"), out);
+  }
+
+  @Test
+  public void getRolesByClaimValuesThrowsWhenRoleMapNull()
+  {
+    BasicAuthorizerCacheManager cache = new StubCacheManager(Map.of(), null);
+    assertTrue(RoleProviderUtil.getRolesByClaimValue("authz", Set.of("r2"),
+                                                     new HashSet<>(), cache
+    ).isEmpty());
+  }
+}