feat: add support to define custom oidc scopes (#9)

* feat: add support to define oidc scopes

* Add custom scopes to oidcIdpConfig

* Fix custom scope render and Handle unsafe char

source/DruidOidcExtension/src/main/java/com/amazon/solutions/druid/oidc/OidcAuthenticator.java

@@ -41,11 +41,15 @@
 import javax.net.ssl.SSLSocketFactory;
 import javax.servlet.DispatcherType;
 import javax.servlet.Filter;
+import java.util.ArrayList;
 import java.util.EnumSet;
+import java.util.List;
 import java.util.Map;
 
 @JsonTypeName("oidc")
 public class OidcAuthenticator implements Authenticator {
+    public static final String DEFAULT_SCOPE = "openid";
+
     private final String name;
     private final String authorizerName;
     private final Supplier<Config> pac4jConfigSupplier;
@@ -125,6 +129,14 @@ private Config createPac4jConfig(OidcConfig oidcConfig) {
                 // "discoveryURI"
                 new CustomSSLResourceRetriever(oidcConfig.getReadTimeout().getMillis(), sslSocketFactory));
 
+        if (oidcConfig.getCustomScopes() != null) {
+            List<String> finalScopes = new ArrayList<>(oidcConfig.getCustomScopes());
+            if (!finalScopes.contains(DEFAULT_SCOPE)) {
+                finalScopes.add(DEFAULT_SCOPE);
+            }
+            oidcConf.setScope(String.join(" ", finalScopes));
+        }
+
         OidcClient oidcClient = new OidcClient(oidcConf);
         oidcClient.setUrlResolver(new DefaultUrlResolver(true));
         oidcClient.setCallbackUrlResolver(new NoParameterCallbackUrlResolver());

source/DruidOidcExtension/src/main/java/com/amazon/solutions/druid/oidc/OidcConfig.java

@@ -26,6 +26,8 @@
 import org.apache.druid.metadata.PasswordProvider;
 import org.joda.time.Duration;
 
+import java.util.List;
+
 public class OidcConfig {
     @JsonProperty
     private final String clientID;
@@ -39,6 +41,9 @@ public class OidcConfig {
     @JsonProperty
     private final String groupClaimName;
 
+    @JsonProperty
+    private final List<String> customScopes;
+
     @JsonProperty
     private final boolean enableCustomSslContext;
 
@@ -63,6 +68,7 @@ public OidcConfig(
             @JsonProperty("clientSecret") PasswordProvider clientSecret,
             @JsonProperty("discoveryURI") String discoveryURI,
             @JsonProperty("groupClaimName") String groupClaimName,
+            @JsonProperty("customScopes") List<String> customScopes,
             @JsonProperty("enableCustomSslContext") boolean enableCustomSslContext,
             @JsonProperty("cookiePassphrase") PasswordProvider cookiePassphrase,
             @JsonProperty("readTimeout") Duration readTimeout,
@@ -73,6 +79,7 @@ public OidcConfig(
         this.clientSecret = Preconditions.checkNotNull(clientSecret, "null clientSecret");
         this.discoveryURI = Preconditions.checkNotNull(discoveryURI, "null discoveryURI");
         this.groupClaimName = groupClaimName;
+        this.customScopes = customScopes;
         this.enableCustomSslContext = enableCustomSslContext;
         this.cookiePassphrase = Preconditions.checkNotNull(cookiePassphrase, "null cookiePassphrase");
         this.readTimeout = readTimeout == null ? Duration.millis(5000) : readTimeout;
@@ -101,6 +108,11 @@ public String getGroupClaimName() {
         return groupClaimName;
     }
 
+    @JsonProperty
+    public List<String> getCustomScopes() {
+        return customScopes;
+    }
+
     @JsonProperty
     public boolean isEnableCustomSslContext() {
         return enableCustomSslContext;

source/DruidOidcExtension/src/test/java/com/amazon/solutions/druid/oidc/OidcAuthenticatorTest.java

@@ -33,6 +33,8 @@
 
 import com.google.inject.Provider;
 
+import java.util.Arrays;
+
 public class OidcAuthenticatorTest {
     private OidcConfig config;
     private Provider<SSLContext> provider;
@@ -50,15 +52,31 @@ public void setup() {
         when(passwordProvider.getPassword()).thenReturn("secret");
         when(config.getDiscoveryURI()).thenReturn("http://localhost");
         when(config.getReadTimeout()).thenReturn(new Duration(10));
+    }
 
+    @Test
+    public void canInitialiseOidcFilter() {
         authenticator = new OidcAuthenticator("name", "authorizerName", config, provider);
+        Filter oidcFilter = authenticator.getFilter();
+
+        assertNotNull(oidcFilter);
     }
 
     @Test
-    public void canInitialiseOidcFilter() {
+    public void canInitialiseOidcFilterWithoutCustomScopes() {
+        when(config.getCustomScopes()).thenReturn(null);
+        authenticator = new OidcAuthenticator("name", "authorizerName", config, provider);
         Filter oidcFilter = authenticator.getFilter();
 
         assertNotNull(oidcFilter);
     }
 
+    @Test
+    public void canInitialiseOidcFilterWithCustomScopes() {
+        when(config.getCustomScopes()).thenReturn(Arrays.asList("groups", "druid"));
+        authenticator = new OidcAuthenticator("name", "authorizerName", config, provider);
+        Filter oidcFilter = authenticator.getFilter();
+
+        assertNotNull(oidcFilter);
+    }
 }

source/DruidOidcExtension/src/test/java/com/amazon/solutions/druid/oidc/OidcConfigTest.java

@@ -25,6 +25,8 @@
 
 import com.fasterxml.jackson.databind.ObjectMapper;
 
+import java.util.Arrays;
+
 public class OidcConfigTest {
     @Test
     public void canParseConfig() throws Exception {
@@ -36,6 +38,7 @@ public void canParseConfig() throws Exception {
                 + "  \"clientSecret\": \"testsecret\",\n"
                 + "  \"discoveryURI\": \"testdiscoveryuri\",\n"
                 + "  \"groupClaimName\": \"group\",\n"
+                + "  \"customScopes\": [\"groups\", \"druid\"],\n"
                 + "  \"enableCustomSslContext\": true,\n"
                 + "  \"cookiePassphrase\": \"testcookiePassphrase\",\n"
                 + "  \"readTimeout\": \"PT10S\",\n"
@@ -52,9 +55,9 @@ public void canParseConfig() throws Exception {
         Assert.assertEquals("testsecret", config.getClientSecret().getPassword());
         Assert.assertEquals("testdiscoveryuri", config.getDiscoveryURI());
         Assert.assertEquals("group", config.getGroupClaimName());
+        Assert.assertEquals(Arrays.asList("groups", "druid"), config.getCustomScopes());
         Assert.assertEquals(true, config.isEnableCustomSslContext());
         Assert.assertEquals("testcookiePassphrase", config.getCookiePassphrase().getPassword());
         Assert.assertEquals(10_000L, config.getReadTimeout().getMillis());
-
     }
 }

source/lib/config/user_data/common_user_data

@@ -125,6 +125,7 @@ $PYTHON $DRUID_HOME/scripts/druid/render_druid_config.py \
     --oidc-client-id {{OIDC_CLIENT_ID}} \
     --oidc-discovery-uri {{OIDC_DISCOVERY_URI}} \
     --oidc-group-claim-name {{OIDC_GROUP_CLAIM_NAME}} \
+    --oidc-custom-scopes {{OIDC_CUSTOM_SCOPES}} \
     --druid-base-url {{DRUID_BASE_URL}} \
     --solution-version {{SOLUTION_VERSION}}
 

source/lib/constructs/druidAutoScalingGroup.ts

@@ -229,6 +229,10 @@ export class DruidAutoScalingGroup extends Construct {
                 asgContext.clusterParams.oidcIdpConfig?.groupClaimName,
                 ''
             ),
+            OIDC_CUSTOM_SCOPES: utils.ifUndefined(
+                JSON.stringify(asgContext.clusterParams.oidcIdpConfig?.customScopes),
+                ''
+            ),
             DRUID_BASE_URL: props.baseUrl,
             GRACEFUL_TERMINATION_PARAM_NAME: this.gracefulTerminationParam.parameterName,
             ADMIN_USER_SECRET_NAME:

source/lib/constructs/druidEksBase.ts

@@ -426,6 +426,8 @@ export abstract class DruidEksBase extends Construct {
             oidc_discovery_uri: this.props.druidClusterParams.oidcIdpConfig?.discoveryURI,
             oidc_group_claim_name:
                 this.props.druidClusterParams.oidcIdpConfig?.groupClaimName,
+            oidc_custom_scopes:
+                this.props.druidClusterParams.oidcIdpConfig?.customScopes,
             alb_scheme: this.props.druidClusterParams.internetFacing
                 ? 'internet-facing'
                 : 'internal',

source/lib/k8s-manifests/druid-cluster-eks.yaml

@@ -137,6 +137,9 @@ spec:
     {{#oidc_group_claim_name}}
     druid.auth.oidc.groupClaimName={{{oidc_group_claim_name}}}
     {{/oidc_group_claim_name}}
+    {{#oidc_custom_scopes}}
+    druid.auth.oidc.customScopes={{{oidc_custom_scopes}}}
+    {{/oidc_custom_scopes}}
 
     druid.escalator.type=basic
     druid.escalator.internalClientUsername=druid_system

source/lib/uploads/config/_common/common.runtime.properties

@@ -161,6 +161,9 @@ druid.auth.authenticator.jwt.authorizerName=oidc
 {% if oidc_group_claim_name  %}
 druid.auth.oidc.groupClaimName={{oidc_group_claim_name}}
 {% endif %}
+{% if oidc_custom_scopes  %}
+druid.auth.oidc.customScopes={{oidc_custom_scopes | safe}}
+{% endif %}
 
 druid.escalator.type=basic
 druid.escalator.internalClientUsername={{internal_client_username}}

source/lib/uploads/scripts/druid/render_druid_config.py

@@ -62,6 +62,8 @@ def parse_args():
                         help='Zookeeper IPs'),
     parser.add_argument('--oidc-group-claim-name', dest='oidc_group_claim_name', type=str, required=False,
                         nargs='?', const='', default='', help='OIDC group claim name'),
+    parser.add_argument('--oidc-custom-scopes', dest='oidc_custom_scopes', type=str, required=False,
+                        nargs='?', const='', default='', help='OIDC custom scopes'),
     parser.add_argument('--druid-base-url', dest='druid_base_url', type=str, required=False,
                         nargs='?', const='', default='', help='Base url of the druid cluster'),
     parser.add_argument('--solution-version', dest='solution_version', type=str, required=False,
@@ -85,6 +87,7 @@ def render_config(
         oidc_client_id=None,
         oidc_discovery_uri=None,
         oidc_group_claim_name=None,
+        oidc_custom_scopes=None,
         solution_version=None,
 ):
 
@@ -111,6 +114,7 @@ def render_config(
             'oidc_client_id': oidc_client_id,
             'oidc_discovery_uri': oidc_discovery_uri,
             'oidc_group_claim_name': oidc_group_claim_name,
+            'oidc_custom_scopes': json.dumps([e for e in oidc_custom_scopes[1:-1].split(',')]) if oidc_custom_scopes else None,
             'emitter_config': emitter_config,
             'druid_base_url': druid_base_url,
             'solution_version': solution_version,
@@ -146,6 +150,7 @@ def main():
         oidc_client_id=args.oidc_client_id,
         oidc_discovery_uri=args.oidc_discovery_uri,
         oidc_group_claim_name=args.oidc_group_claim_name,
+        oidc_custom_scopes=args.oidc_custom_scopes,
         solution_version=args.solution_version)
 
 

source/lib/utils/types.ts

@@ -301,6 +301,7 @@ export interface OidcIdpConfig {
     clientSecretArn: string;
     discoveryURI: string;
     groupClaimName?: string;
+    customScopes?: string[];
     groupRoleMappings?: Record<string, string[]>;
 }