Require token has role claim

pulsar-broker-auth-oidc/src/main/java/org/apache/pulsar/broker/authentication/oidc/AuthenticationProviderOpenID.java

@@ -36,6 +36,7 @@
 import com.auth0.jwt.exceptions.TokenExpiredException;
 import com.auth0.jwt.interfaces.Claim;
 import com.auth0.jwt.interfaces.DecodedJWT;
+import com.auth0.jwt.interfaces.Verification;
 import io.kubernetes.client.openapi.ApiClient;
 import io.kubernetes.client.util.Config;
 import io.netty.handler.ssl.SslContext;
@@ -115,7 +116,8 @@ public class AuthenticationProviderOpenID implements AuthenticationProvider {
 
     private long acceptedTimeLeewaySeconds;
     private FallbackDiscoveryMode fallbackDiscoveryMode;
-    private String roleClaim;
+    private String roleClaim = ROLE_CLAIM_DEFAULT;
+    private boolean isRoleClaimNotSubject;
 
     static final String ALLOWED_TOKEN_ISSUERS = "openIDAllowedTokenIssuers";
     static final String ISSUER_TRUST_CERTS_FILE_PATH = "openIDTokenIssuerTrustCertsFilePath";
@@ -145,6 +147,7 @@ public class AuthenticationProviderOpenID implements AuthenticationProvider {
     public void initialize(ServiceConfiguration config) throws IOException {
         this.allowedAudiences = validateAllowedAudiences(getConfigValueAsSet(config, ALLOWED_AUDIENCES));
         this.roleClaim = getConfigValueAsString(config, ROLE_CLAIM, ROLE_CLAIM_DEFAULT);
+        this.isRoleClaimNotSubject = !ROLE_CLAIM_DEFAULT.equals(roleClaim);
         this.acceptedTimeLeewaySeconds = getConfigValueAsInt(config, ACCEPTED_TIME_LEEWAY_SECONDS,
                 ACCEPTED_TIME_LEEWAY_SECONDS_DEFAULT);
         boolean requireHttps = getConfigValueAsBoolean(config, REQUIRE_HTTPS, REQUIRE_HTTPS_DEFAULT);
@@ -406,14 +409,19 @@ DecodedJWT verifyJWT(PublicKey publicKey,
 
         // We verify issuer when retrieving the PublicKey, so it is not verified here.
         // The claim presence requirements are based on https://openid.net/specs/openid-connect-basic-1_0.html#IDToken
-        JWTVerifier verifier = JWT.require(alg)
+         Verification verifierBuilder = JWT.require(alg)
                 .acceptLeeway(acceptedTimeLeewaySeconds)
                 .withAnyOfAudience(allowedAudiences)
                 .withClaimPresence(RegisteredClaims.ISSUED_AT)
                 .withClaimPresence(RegisteredClaims.EXPIRES_AT)
                 .withClaimPresence(RegisteredClaims.NOT_BEFORE)
-                .withClaimPresence(RegisteredClaims.SUBJECT)
-                .build();
+                .withClaimPresence(RegisteredClaims.SUBJECT);
+
+        if (isRoleClaimNotSubject) {
+            verifierBuilder = verifierBuilder.withClaimPresence(roleClaim);
+        }
+
+        JWTVerifier verifier = verifierBuilder.build();
 
         try {
             return verifier.verify(jwt);
@@ -432,7 +440,7 @@ DecodedJWT verifyJWT(PublicKey publicKey,
         } catch (JWTDecodeException e) {
             incrementFailureMetric(AuthenticationExceptionCode.ERROR_DECODING_JWT);
             throw new AuthenticationException("Error while decoding JWT: " + e.getMessage());
-        } catch (JWTVerificationException e) {
+        } catch (JWTVerificationException | IllegalArgumentException e) {
             incrementFailureMetric(AuthenticationExceptionCode.ERROR_VERIFYING_JWT);
             throw new AuthenticationException("JWT verification failed: " + e.getMessage());
         }

pulsar-broker-auth-oidc/src/test/java/org/apache/pulsar/broker/authentication/oidc/AuthenticationProviderOpenIDIntegrationTest.java

@@ -40,6 +40,7 @@
 import java.security.interfaces.RSAPublicKey;
 import java.util.Base64;
 import java.util.Date;
+import java.util.HashMap;
 import java.util.Properties;
 import java.util.Set;
 import java.util.concurrent.ExecutionException;
@@ -437,6 +438,49 @@ public void testAuthenticationStateOpenIDForTokenExpiration() throws Exception {
         assertTrue(state.isExpired());
     }
 
+    @Test
+    void ensureRoleClaimForNonSubClaimReturnsRole() throws Exception {
+        AuthenticationProviderOpenID provider = new AuthenticationProviderOpenID();
+        Properties props = new Properties();
+        props.setProperty(AuthenticationProviderOpenID.ALLOWED_TOKEN_ISSUERS, issuer);
+        props.setProperty(AuthenticationProviderOpenID.ALLOWED_AUDIENCES, "allowed-audience");
+        props.setProperty(AuthenticationProviderOpenID.ROLE_CLAIM, "test");
+        props.setProperty(AuthenticationProviderOpenID.REQUIRE_HTTPS, "false");
+        ServiceConfiguration config = new ServiceConfiguration();
+        config.setProperties(props);
+        provider.initialize(config);
+
+        // Build a JWT with a custom claim
+        HashMap<String, Object> claims = new HashMap();
+        claims.put("test", "my-role");
+        String token = generateToken(validJwk, issuer, "not-my-role", "allowed-audience", 0L,
+                0L, 10000L, claims);
+        assertEquals(provider.authenticateAsync(new AuthenticationDataCommand(token)).get(), "my-role");
+    }
+
+    @Test
+    void ensureRoleClaimForNonSubClaimFailsWhenClaimIsMissing() throws Exception {
+        AuthenticationProviderOpenID provider = new AuthenticationProviderOpenID();
+        Properties props = new Properties();
+        props.setProperty(AuthenticationProviderOpenID.ALLOWED_TOKEN_ISSUERS, issuer);
+        props.setProperty(AuthenticationProviderOpenID.ALLOWED_AUDIENCES, "allowed-audience");
+        props.setProperty(AuthenticationProviderOpenID.ROLE_CLAIM, "test");
+        props.setProperty(AuthenticationProviderOpenID.REQUIRE_HTTPS, "false");
+        ServiceConfiguration config = new ServiceConfiguration();
+        config.setProperties(props);
+        provider.initialize(config);
+
+        // Build a JWT without the "test" claim, which should cause the authentication to fail
+        String token = generateToken(validJwk, issuer, "not-my-role", "allowed-audience", 0L,
+                0L, 10000L);
+        try {
+            provider.authenticateAsync(new AuthenticationDataCommand(token)).get();
+            fail("Expected exception");
+        } catch (ExecutionException e) {
+            assertTrue(e.getCause() instanceof AuthenticationException, "Found exception: " + e.getCause());
+        }
+    }
+
     // This test is somewhat counterintuitive. We allow the state object to change roles, but then we fail it
     // in the ServerCnx handling of the state object. As such, it is essential that the state object allow
     // the role to change.
@@ -462,6 +506,11 @@ public void testAuthenticationStateOpenIDAllowsRoleChange() throws Exception {
 
     private String generateToken(String kid, String issuer, String subject, String audience,
                                  Long iatOffset, Long nbfOffset, Long expOffset) {
+        return generateToken(kid, issuer, subject, audience, iatOffset, nbfOffset, expOffset, new HashMap<>());
+    }
+
+    private String generateToken(String kid, String issuer, String subject, String audience,
+                                 Long iatOffset, Long nbfOffset, Long expOffset, HashMap<String, Object> extraClaims) {
         long now = System.currentTimeMillis();
         DefaultJwtBuilder defaultJwtBuilder = new DefaultJwtBuilder();
         defaultJwtBuilder.setHeaderParam("kid", kid);
@@ -473,6 +522,7 @@ private String generateToken(String kid, String issuer, String subject, String a
         defaultJwtBuilder.setIssuedAt(iatOffset != null ? new Date(now + iatOffset) : null);
         defaultJwtBuilder.setNotBefore(nbfOffset != null ? new Date(now + nbfOffset) : null);
         defaultJwtBuilder.setExpiration(expOffset != null ? new Date(now + expOffset) : null);
+        defaultJwtBuilder.addClaims(extraClaims);
         defaultJwtBuilder.signWith(privateKey);
         return defaultJwtBuilder.compact();
     }