Support custom role claim (#1492)

akrambek · web-flow · commit 09b88c8dba74 · 2025-06-11T11:02:26.000-07:00
diff --git a/runtime/guard-jwt/src/main/java/io/aklivity/zilla/runtime/guard/jwt/config/JwtOptionsConfig.java b/runtime/guard-jwt/src/main/java/io/aklivity/zilla/runtime/guard/jwt/config/JwtOptionsConfig.java
@@ -27,6 +27,7 @@ public class JwtOptionsConfig extends OptionsConfig
 {
     public final String issuer;
     public final String audience;
+    public final String roles;
     public final List<JwtKeyConfig> keys;
     public final Optional<Duration> challenge;
     public final String identity;
@@ -46,13 +47,15 @@ public static <T> JwtOptionsConfigBuilder<T> builder(
     JwtOptionsConfig(
         String issuer,
         String audience,
+        String roles,
         List<JwtKeyConfig> keys,
         Duration challenge,
         String identity,
         String keysURL)
     {
         this.issuer = issuer;
         this.audience = audience;
+        this.roles = roles;
         this.keys = keys;
         this.challenge = ofNullable(challenge);
         this.identity = identity;
diff --git a/runtime/guard-jwt/src/main/java/io/aklivity/zilla/runtime/guard/jwt/config/JwtOptionsConfigBuilder.java b/runtime/guard-jwt/src/main/java/io/aklivity/zilla/runtime/guard/jwt/config/JwtOptionsConfigBuilder.java
@@ -24,10 +24,13 @@
 
 public class JwtOptionsConfigBuilder<T> extends ConfigBuilder<T, JwtOptionsConfigBuilder<T>>
 {
+    public static final String ROLES_DEFAULT = "scope";
+
     private final Function<OptionsConfig, T> mapper;
 
     private String issuer;
     private String audience;
+    private String roles;
     private List<JwtKeyConfig> keys;
     private Duration challenge;
     private String identity;
@@ -60,6 +63,13 @@ public JwtOptionsConfigBuilder<T> audience(
         return this;
     }
 
+    public JwtOptionsConfigBuilder<T> roles(
+        String roles)
+    {
+        this.roles = roles;
+        return this;
+    }
+
     public JwtOptionsConfigBuilder<T> challenge(
         Duration challenge)
     {
@@ -107,6 +117,8 @@ public JwtOptionsConfigBuilder<T> keysURL(
     @Override
     public T build()
     {
-        return mapper.apply(new JwtOptionsConfig(issuer, audience, keys, challenge, identity, keysURL));
+        String roles = this.roles != null ? this.roles : ROLES_DEFAULT;
+
+        return mapper.apply(new JwtOptionsConfig(issuer, audience, roles, keys, challenge, identity, keysURL));
     }
 }
diff --git a/runtime/guard-jwt/src/main/java/io/aklivity/zilla/runtime/guard/jwt/internal/JwtGuardHandler.java b/runtime/guard-jwt/src/main/java/io/aklivity/zilla/runtime/guard/jwt/internal/JwtGuardHandler.java
@@ -53,10 +53,14 @@
 
 public class JwtGuardHandler implements GuardHandler
 {
+    private static final String SPLIT_VALUE_PATTERN = "\\s+";
+    private static final String SPLIT_PATH_PATTERN = "\\.";
+
     private final JsonWebSignature signature = new JsonWebSignature();
 
     private final String issuer;
     private final String audience;
+    private final String roles;
     private final Duration challenge;
     private final String identity;
     private final Map<String, JsonWebKey> keys;
@@ -72,6 +76,7 @@ public JwtGuardHandler(
     {
         this.issuer = options.issuer;
         this.audience = options.audience;
+        this.roles = options.roles;
         this.challenge = options.challenge.orElse(null);
         this.identity = options.identity;
 
@@ -179,11 +184,15 @@ public long reauthorize(
                 break authorize;
             }
 
-            List<String> roles = Optional.ofNullable(claims.getClaimValue("scope"))
-                .map(s -> s.toString().intern())
-                .map(s -> s.split("\\s+"))
-                .map(Arrays::asList)
-                .orElse(null);
+            Object rolesValue = claimValue(claims, this.roles);
+            @SuppressWarnings("unchecked")
+            List<String> roles = (rolesValue instanceof List)
+                    ? (List<String>) rolesValue
+                    : Optional.ofNullable(rolesValue)
+                        .map(Object::toString)
+                        .map(s -> s.split(SPLIT_VALUE_PATTERN))
+                        .map(Arrays::asList)
+                        .orElse(null);
 
             JwtSessionStore sessionStore = supplySessionStore(contextId);
             session = sessionStore.supplySession(identity, roles);
@@ -407,6 +416,33 @@ private void unshareIfNecessary()
         }
     }
 
+    private static Object claimValue(
+        Object node,
+        String path)
+    {
+        Object current = node;
+        for (String part : path.split(SPLIT_PATH_PATTERN))
+        {
+            if (current == null)
+            {
+                break;
+            }
+            if (current instanceof JwtClaims)
+            {
+                current = ((JwtClaims) current).getClaimValue(part);
+            }
+            else if (current instanceof Map)
+            {
+                current = ((Map<?, ?>) current).get(part);
+            }
+            else
+            {
+                current = null;
+            }
+        }
+        return current;
+    }
+
     private static String readKeys(
         Path keysPath)
     {
diff --git a/runtime/guard-jwt/src/main/java/io/aklivity/zilla/runtime/guard/jwt/internal/config/JwtOptionsConfigAdapter.java b/runtime/guard-jwt/src/main/java/io/aklivity/zilla/runtime/guard/jwt/internal/config/JwtOptionsConfigAdapter.java
@@ -14,6 +14,7 @@
  */
 package io.aklivity.zilla.runtime.guard.jwt.internal.config;
 
+import static io.aklivity.zilla.runtime.guard.jwt.config.JwtOptionsConfigBuilder.ROLES_DEFAULT;
 import static java.util.Collections.emptyList;
 
 import java.time.Duration;
@@ -38,6 +39,7 @@ public final class JwtOptionsConfigAdapter implements OptionsConfigAdapterSpi, J
 {
     private static final String ISSUER_NAME = "issuer";
     private static final String AUDIENCE_NAME = "audience";
+    private static final String ROLES = "roles";
     private static final String KEYS_NAME = "keys";
     private static final String CHALLENGE_NAME = "challenge";
     private static final String IDENTITY_NAME = "identity";
@@ -76,6 +78,11 @@ public JsonObject adaptToJson(
             object.add(AUDIENCE_NAME, jwtOptions.audience);
         }
 
+        if (jwtOptions.roles != null && !ROLES_DEFAULT.equals(jwtOptions.roles))
+        {
+            object.add(ROLES, jwtOptions.roles);
+        }
+
         if (jwtOptions.keys != null)
         {
             JsonArrayBuilder newKeys = Json.createArrayBuilder();
@@ -118,6 +125,11 @@ public OptionsConfig adaptFromJson(
             jwtOptions.audience(object.getString(AUDIENCE_NAME));
         }
 
+        if (object.containsKey(ROLES))
+        {
+            jwtOptions.roles(object.getString(ROLES));
+        }
+
         if (object.containsKey(KEYS_NAME))
         {
             JsonValue keysValue = object.getValue(String.format("/%s", KEYS_NAME));
diff --git a/runtime/guard-jwt/src/test/java/io/aklivity/zilla/runtime/guard/jwt/internal/JwtGuardHandlerTest.java b/runtime/guard-jwt/src/test/java/io/aklivity/zilla/runtime/guard/jwt/internal/JwtGuardHandlerTest.java
@@ -31,6 +31,8 @@
 import java.time.Clock;
 import java.time.Duration;
 import java.time.Instant;
+import java.util.List;
+import java.util.Map;
 
 import org.agrona.collections.MutableLong;
 import org.jose4j.jws.JsonWebSignature;
@@ -650,6 +652,41 @@ public void shouldDeauthorize() throws Exception
         guard.deauthorize(sessionId);
     }
 
+    @Test
+    public void shouldAuthorizeWithCustomRole() throws Exception
+    {
+        Duration challenge = ofSeconds(3L);
+        JwtOptionsConfig options = JwtOptionsConfig.builder()
+            .inject(identity())
+            .issuer("test issuer")
+            .audience("testAudience")
+            .roles("realm_access.roles")
+            .key(RFC7515_RS256_CONFIG)
+            .challenge(challenge)
+            .build();
+        JwtGuardHandler guard = new JwtGuardHandler(options, context, new MutableLong(1L)::getAndIncrement);
+
+        Instant now = Instant.now();
+
+        JwtClaims claims = new JwtClaims();
+        claims.setClaim("iss", "test issuer");
+        claims.setClaim("aud", "testAudience");
+        claims.setClaim("sub", "testSubject");
+        claims.setClaim("exp", now.getEpochSecond() + 10L);
+        claims.setClaim("realm_access",
+            Map.of("roles", List.of("default-roles-backend", "offline_access", "uma_authorization")));
+        String token = sign(claims.toJson(), "test", RFC7515_RS256, "RS256");
+
+        long sessionId = guard.reauthorize(0L, 0L, 101L, token);
+
+        assertThat(sessionId, not(equalTo(0L)));
+        assertThat(guard.identity(sessionId), equalTo("testSubject"));
+        assertThat(guard.expiresAt(sessionId), equalTo(ofSeconds(now.getEpochSecond() + 10L).toMillis()));
+        assertThat(guard.expiringAt(sessionId), equalTo(ofSeconds(now.getEpochSecond() + 10L).minus(challenge).toMillis()));
+        assertTrue(guard.verify(sessionId, asList("default-roles-backend", "offline_access", "uma_authorization")));
+        assertFalse(guard.verify(sessionId, asList("admin")));
+    }
+
     static String sign(
         String payload,
         String kid,
diff --git a/specs/guard-jwt.spec/src/main/scripts/io/aklivity/zilla/specs/guard/jwt/config/guard.yaml b/specs/guard-jwt.spec/src/main/scripts/io/aklivity/zilla/specs/guard/jwt/config/guard.yaml
@@ -21,6 +21,7 @@ guards:
     options:
       issuer: https://auth.example.com
       audience: https://api.example.com
+      roles: "realm_access.roles"
       keys:
         - kty: EC
           crv: P-256
diff --git a/specs/guard-jwt.spec/src/main/scripts/io/aklivity/zilla/specs/guard/jwt/schema/jwt.schema.patch.json b/specs/guard-jwt.spec/src/main/scripts/io/aklivity/zilla/specs/guard/jwt/schema/jwt.schema.patch.json
@@ -41,6 +41,12 @@
                                 "title": "Audience",
                                 "type": "string"
                             },
+                            "roles":
+                            {
+                                "title": "Roles",
+                                "type": "string",
+                                "default": "scope"
+                            },
                             "keys":
                             {
                                 "title": "Keys",

Original file line number	Diff line number	Diff line change
`@@ -24,10 +24,13 @@`
`24`	`24`
`25`	`25`	`public class JwtOptionsConfigBuilder<T> extends ConfigBuilder<T, JwtOptionsConfigBuilder<T>>`
`26`	`26`	`{`
	`27`	`+ public static final String ROLES_DEFAULT = "scope";`
	`28`	`+`
`27`	`29`	`private final Function<OptionsConfig, T> mapper;`
`28`	`30`
`29`	`31`	`private String issuer;`
`30`	`32`	`private String audience;`
	`33`	`+ private String roles;`
`31`	`34`	`private List<JwtKeyConfig> keys;`
`32`	`35`	`private Duration challenge;`
`33`	`36`	`private String identity;`
`@@ -60,6 +63,13 @@ public JwtOptionsConfigBuilder<T> audience(`
`60`	`63`	`return this;`
`61`	`64`	`}`
`62`	`65`
	`66`	`+ public JwtOptionsConfigBuilder<T> roles(`
	`67`	`+ String roles)`
	`68`	`+ {`
	`69`	`+ this.roles = roles;`
	`70`	`+ return this;`
	`71`	`+ }`
	`72`	`+`
`63`	`73`	`public JwtOptionsConfigBuilder<T> challenge(`
`64`	`74`	`Duration challenge)`
`65`	`75`	`{`
`@@ -107,6 +117,8 @@ public JwtOptionsConfigBuilder<T> keysURL(`
`107`	`117`	`@Override`
`108`	`118`	`public T build()`
`109`	`119`	`{`
`110`		`- return mapper.apply(new JwtOptionsConfig(issuer, audience, keys, challenge, identity, keysURL));`
	`120`	`+ String roles = this.roles != null ? this.roles : ROLES_DEFAULT;`
	`121`	`+`
	`122`	`+ return mapper.apply(new JwtOptionsConfig(issuer, audience, roles, keys, challenge, identity, keysURL));`
`111`	`123`	`}`
`112`	`124`	`}`