From b31b71e83c1036fab220b2ef0ff23e6296f11eb1 Mon Sep 17 00:00:00 2001
From: Zike Yang <ar@armail.top>
Date: Tue, 27 Apr 2021 00:53:52 +0800
Subject: [PATCH] Support array type claims in JWT (#10375)

---
 .../AuthenticationProviderToken.java          | 11 ++++-
 .../AuthenticationProviderTokenTest.java      | 49 +++++++++++++++++++
 2 files changed, 59 insertions(+), 1 deletion(-)

diff --git a/pulsar-broker-common/src/main/java/org/apache/pulsar/broker/authentication/AuthenticationProviderToken.java b/pulsar-broker-common/src/main/java/org/apache/pulsar/broker/authentication/AuthenticationProviderToken.java
index 5b2db049519a2..f2eaf7e337fee 100644
--- a/pulsar-broker-common/src/main/java/org/apache/pulsar/broker/authentication/AuthenticationProviderToken.java
+++ b/pulsar-broker-common/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/pulsar-broker-common/src/test/java/org/apache/pulsar/broker/authentication/AuthenticationProviderTokenTest.java b/pulsar-broker-common/src/test/java/org/apache/pulsar/broker/authentication/AuthenticationProviderTokenTest.java
index fe4d6a7ac09bb..cdc5e210d265f 100644
--- a/pulsar-broker-common/src/test/java/org/apache/pulsar/broker/authentication/AuthenticationProviderTokenTest.java
+++ b/pulsar-broker-common/src/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)