Support array type claims in JWT (apache#10375)

Jhanay · Apr 26, 2021 · b31b71e · b31b71e
1 parent 24f35af
commit b31b71e
Show file tree

Hide file tree

Showing 2 changed files with 59 additions and 1 deletion.
diff --git a/...on/src/main/java/org/apache/pulsar/broker/authentication/AuthenticationProviderToken.java b/...on/src/main/java/org/apache/pulsar/broker/authentication/AuthenticationProviderToken.java
@@ -31,6 +31,7 @@
 
 import com.google.common.annotations.VisibleForTesting;
 import io.jsonwebtoken.ExpiredJwtException;
+import io.jsonwebtoken.RequiredTypeException;
 import io.prometheus.client.Counter;
 import io.prometheus.client.Histogram;
 import org.apache.commons.lang3.StringUtils;
@@ -233,7 +234,15 @@ private Jwt<?, Claims> authenticateToken(final String token) throws Authenticati
     }
 
     private String getPrincipal(Jwt<?, Claims> jwt) {
-        return jwt.getBody().get(roleClaim, String.class);
+        try {
+            return jwt.getBody().get(roleClaim, String.class);
+        } catch (RequiredTypeException requiredTypeException) {
+            List list = jwt.getBody().get(roleClaim, List.class);
+            if (list != null && !list.isEmpty() && list.get(0) instanceof String) {
+                return (String) list.get(0);
+            }
+            return null;
+        }
     }
 
     /**

diff --git a/...rc/test/java/org/apache/pulsar/broker/authentication/AuthenticationProviderTokenTest.java b/...rc/test/java/org/apache/pulsar/broker/authentication/AuthenticationProviderTokenTest.java
@@ -34,6 +34,7 @@
 import io.jsonwebtoken.io.Decoders;
 import io.jsonwebtoken.security.Keys;
 import java.security.Key;
+import java.util.Arrays;
 import java.util.List;
 import lombok.Cleanup;
 
@@ -750,6 +751,54 @@ public void testMultiTokenAudienceNotInclude() throws Exception {
         testTokenAudienceWithDifferentConfig(properties, audienceClaim, audiences);
     }
 
+    @Test
+    public void testArrayTypeRoleClaim() throws Exception {
+        String authRoleClaim = "customClaim";
+        String authRole = "my-test-role";
+
+        KeyPair keyPair = Keys.keyPairFor(SignatureAlgorithm.RS256);
+
+        String privateKeyStr = AuthTokenUtils.encodeKeyBase64(keyPair.getPrivate());
+        String publicKeyStr = AuthTokenUtils.encodeKeyBase64(keyPair.getPublic());
+
+        AuthenticationProviderToken provider = new AuthenticationProviderToken();
+
+        Properties properties = new Properties();
+        // Use public key for validation
+        properties.setProperty(AuthenticationProviderToken.CONF_TOKEN_PUBLIC_KEY, publicKeyStr);
+        // Set custom claim field
+        properties.setProperty(AuthenticationProviderToken.CONF_TOKEN_AUTH_CLAIM, authRoleClaim);
+
+        ServiceConfiguration conf = new ServiceConfiguration();
+        conf.setProperties(properties);
+        provider.initialize(conf);
+
+        // Use private key to generate token
+        PrivateKey privateKey = AuthTokenUtils.decodePrivateKey(Decoders.BASE64.decode(privateKeyStr), SignatureAlgorithm.RS256);
+        String token = Jwts.builder()
+                .setClaims(new HashMap<String, Object>() {{
+                    put(authRoleClaim, Arrays.asList(authRole, "other-role"));
+                }})
+                .signWith(privateKey)
+                .compact();
+
+        // Pulsar protocol auth
+        String role = provider.authenticate(new AuthenticationDataSource() {
+            @Override
+            public boolean hasDataFromCommand() {
+                return true;
+            }
+
+            @Override
+            public String getCommandData() {
+                return token;
+            }
+        });
+        assertEquals(role, authRole);
+
+        provider.close();
+    }
+
     private static String createTokenWithAudience(Key signingKey, String audienceClaim, List<String> audience) {
         JwtBuilder builder = Jwts.builder()
                 .setSubject(SUBJECT)