Include principal name in Polaris tokens (#2389)

* Include principal name in Polaris tokens

Summary of changes:

- Instead of including the principal id twice in the token, the principal name is now used as the subject claim. While the default authenticator doesn't need the principal name and works with just the principal id, not having the "real" principal name available could be a problem for other authenticator implementations.

- `DecodedToken` has been refactored and renamed to `InternalPolarisCredential`. It is also now a package-private component.

- `TokenBroker.verify()` now returns PolarisCredential.

* rename to InternalPolarisToken

Author: adutra

runtime/service/src/main/java/org/apache/polaris/service/auth/DecodedToken.java

@@ -0,0 +1,72 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.polaris.service.auth;
+
+import com.google.common.base.Splitter;
+import jakarta.annotation.Nonnull;
+import java.util.Set;
+import java.util.stream.Collectors;
+import org.apache.polaris.immutables.PolarisImmutable;
+import org.immutables.value.Value;
+
+/**
+ * A specialized {@link PolarisCredential} used for internal authentication, when Polaris is the
+ * identity provider.
+ *
+ * <p>Such credentials are created by the Polaris service itself, from a JWT token previously issued
+ * by Polaris itself.
+ *
+ * @see JWTBroker
+ */
+@PolarisImmutable
+abstract class InternalPolarisToken implements PolarisCredential {
+
+  private static final Splitter SCOPE_SPLITTER = Splitter.on(' ').omitEmptyStrings().trimResults();
+
+  static InternalPolarisToken of(
+      String principalName, Long principalId, String clientId, String scope) {
+    return ImmutableInternalPolarisToken.builder()
+        .principalName(principalName)
+        .principalId(principalId)
+        .clientId(clientId)
+        .scope(scope)
+        .build();
+  }
+
+  @Nonnull // switch from nullable to non-nullable
+  @Override
+  @SuppressWarnings("NullableProblems")
+  public abstract String getPrincipalName();
+
+  @Nonnull // switch from nullable to non-nullable
+  @Override
+  @SuppressWarnings("NullableProblems")
+  public abstract Long getPrincipalId();
+
+  @Value.Lazy
+  @Override
+  public Set<String> getPrincipalRoles() {
+    // Polaris stores roles in the scope claim
+    return SCOPE_SPLITTER.splitToStream(getScope()).collect(Collectors.toSet());
+  }
+
+  abstract String getClientId();
+
+  abstract String getScope();
+}

runtime/service/src/main/java/org/apache/polaris/service/auth/InternalPolarisToken.java

@@ -20,12 +20,10 @@
 
 import com.auth0.jwt.JWT;
 import com.auth0.jwt.algorithms.Algorithm;
-import com.auth0.jwt.exceptions.JWTVerificationException;
 import com.auth0.jwt.interfaces.DecodedJWT;
 import com.auth0.jwt.interfaces.JWTVerifier;
 import java.time.Instant;
 import java.time.temporal.ChronoUnit;
-import java.util.Objects;
 import java.util.Optional;
 import java.util.UUID;
 import org.apache.iceberg.exceptions.NotAuthorizedException;
@@ -60,34 +58,22 @@ public abstract class JWTBroker implements TokenBroker {
   public abstract Algorithm getAlgorithm();
 
   @Override
-  public DecodedToken verify(String token) {
+  public PolarisCredential verify(String token) {
+    return verifyInternal(token);
+  }
+
+  private InternalPolarisToken verifyInternal(String token) {
     JWTVerifier verifier = JWT.require(getAlgorithm()).withClaim(CLAIM_KEY_ACTIVE, true).build();
 
     try {
       DecodedJWT decodedJWT = verifier.verify(token);
-      return new DecodedToken() {
-        @Override
-        public Long getPrincipalId() {
-          return decodedJWT.getClaim("principalId").asLong();
-        }
-
-        @Override
-        public String getClientId() {
-          return decodedJWT.getClaim("client_id").asString();
-        }
-
-        @Override
-        public String getSub() {
-          return decodedJWT.getSubject();
-        }
-
-        @Override
-        public String getScope() {
-          return decodedJWT.getClaim("scope").asString();
-        }
-      };
-
-    } catch (JWTVerificationException e) {
+      return InternalPolarisToken.of(
+          decodedJWT.getSubject(),
+          decodedJWT.getClaim(CLAIM_KEY_PRINCIPAL_ID).asLong(),
+          decodedJWT.getClaim(CLAIM_KEY_CLIENT_ID).asString(),
+          decodedJWT.getClaim(CLAIM_KEY_SCOPE).asString());
+
+    } catch (Exception e) {
       throw (NotAuthorizedException)
           new NotAuthorizedException("Failed to verify the token").initCause(e);
     }
@@ -110,26 +96,26 @@ public TokenResponse generateFromToken(
     if (subjectToken == null || subjectToken.isBlank()) {
       return new TokenResponse(OAuthTokenErrorResponse.Error.invalid_request);
     }
-    DecodedToken decodedToken;
+    InternalPolarisToken decodedToken;
     try {
-      decodedToken = verify(subjectToken);
+      decodedToken = verifyInternal(subjectToken);
     } catch (NotAuthorizedException e) {
       LOGGER.error("Failed to verify the token", e.getCause());
       return new TokenResponse(Error.invalid_client);
     }
     EntityResult principalLookup =
         metaStoreManager.loadEntity(
-            polarisCallContext,
-            0L,
-            Objects.requireNonNull(decodedToken.getPrincipalId()),
-            PolarisEntityType.PRINCIPAL);
+            polarisCallContext, 0L, decodedToken.getPrincipalId(), PolarisEntityType.PRINCIPAL);
     if (!principalLookup.isSuccess()
         || principalLookup.getEntity().getType() != PolarisEntityType.PRINCIPAL) {
       return new TokenResponse(OAuthTokenErrorResponse.Error.unauthorized_client);
     }
     String tokenString =
         generateTokenString(
-            decodedToken.getClientId(), decodedToken.getScope(), decodedToken.getPrincipalId());
+            decodedToken.getPrincipalName(),
+            decodedToken.getPrincipalId(),
+            decodedToken.getClientId(),
+            decodedToken.getScope());
     return new TokenResponse(
         tokenString, TokenType.ACCESS_TOKEN.getValue(), maxTokenGenerationInSeconds);
   }
@@ -156,16 +142,18 @@ public TokenResponse generateFromClientSecrets(
     if (principal.isEmpty()) {
       return new TokenResponse(OAuthTokenErrorResponse.Error.unauthorized_client);
     }
-    String tokenString = generateTokenString(clientId, scope, principal.get().getId());
+    String tokenString =
+        generateTokenString(principal.get().getName(), principal.get().getId(), clientId, scope);
     return new TokenResponse(
         tokenString, TokenType.ACCESS_TOKEN.getValue(), maxTokenGenerationInSeconds);
   }
 
-  private String generateTokenString(String clientId, String scope, Long principalId) {
+  private String generateTokenString(
+      String principalName, long principalId, String clientId, String scope) {
     Instant now = Instant.now();
     return JWT.create()
         .withIssuer(ISSUER_KEY)
-        .withSubject(String.valueOf(principalId))
+        .withSubject(principalName)
         .withIssuedAt(now)
         .withExpiresAt(now.plus(maxTokenGenerationInSeconds, ChronoUnit.SECONDS))
         .withJWTId(UUID.randomUUID().toString())

runtime/service/src/main/java/org/apache/polaris/service/auth/JWTBroker.java

@@ -64,7 +64,7 @@ public TokenResponse generateFromToken(
         }
 
         @Override
-        public DecodedToken verify(String token) {
+        public PolarisCredential verify(String token) {
           return null;
         }
       };

runtime/service/src/main/java/org/apache/polaris/service/auth/NoneTokenBrokerFactory.java

@@ -61,7 +61,8 @@ TokenResponse generateFromToken(
       PolarisCallContext polarisCallContext,
       TokenType requestedTokenType);
 
-  DecodedToken verify(String token);
+  /** Decodes and verifies the token, then returns the associated {@link PolarisCredential}. */
+  PolarisCredential verify(String token);
 
   static @Nonnull Optional<PrincipalEntity> findPrincipalEntity(
       PolarisMetaStoreManager metaStoreManager,

runtime/service/src/main/java/org/apache/polaris/service/auth/TokenBroker.java

@@ -38,7 +38,7 @@
 import java.util.Set;
 import org.apache.polaris.service.auth.AuthenticationRealmConfiguration;
 import org.apache.polaris.service.auth.AuthenticationType;
-import org.apache.polaris.service.auth.DecodedToken;
+import org.apache.polaris.service.auth.PolarisCredential;
 import org.apache.polaris.service.auth.TokenBroker;
 
 /**
@@ -90,7 +90,7 @@ public Uni<SecurityIdentity> authenticate(
 
     String credential = authHeader.substring(spaceIdx + 1);
 
-    DecodedToken token;
+    PolarisCredential token;
     try {
       token = tokenBroker.verify(credential);
     } catch (Exception e) {

runtime/service/src/main/java/org/apache/polaris/service/auth/internal/InternalAuthenticationMechanism.java

@@ -52,7 +52,7 @@ public void setUp() {
 
   @Test
   public void testFetchPrincipalThrowsServiceExceptionOnMetastoreException() {
-    DecodedToken token = Mockito.mock(DecodedToken.class);
+    PolarisCredential token = Mockito.mock(PolarisCredential.class);
     long principalId = 100L;
     when(token.getPrincipalId()).thenReturn(principalId);
     when(metaStoreManager.loadEntity(
@@ -69,10 +69,9 @@ public void testFetchPrincipalThrowsServiceExceptionOnMetastoreException() {
 
   @Test
   public void testFetchPrincipalThrowsNotAuthorizedWhenNotFound() {
-    DecodedToken token = Mockito.mock(DecodedToken.class);
+    PolarisCredential token = Mockito.mock(PolarisCredential.class);
     long principalId = 100L;
     when(token.getPrincipalId()).thenReturn(principalId);
-    when(token.getClientId()).thenReturn("abc");
     when(metaStoreManager.loadEntity(
             authenticator.callContext.getPolarisCallContext(),
             0L,

runtime/service/src/test/java/org/apache/polaris/service/auth/DefaultAuthenticatorTest.java

@@ -34,7 +34,7 @@
 import org.apache.iceberg.exceptions.NotAuthorizedException;
 import org.apache.polaris.service.auth.AuthenticationRealmConfiguration;
 import org.apache.polaris.service.auth.AuthenticationType;
-import org.apache.polaris.service.auth.DecodedToken;
+import org.apache.polaris.service.auth.PolarisCredential;
 import org.apache.polaris.service.auth.TokenBroker;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -156,7 +156,7 @@ public void testAuthenticateWithValidToken() {
     when(routingContext.request()).thenReturn(mock(io.vertx.core.http.HttpServerRequest.class));
     when(routingContext.request().getHeader("Authorization")).thenReturn("Bearer validToken");
 
-    DecodedToken decodedToken = mock(DecodedToken.class);
+    PolarisCredential decodedToken = mock(PolarisCredential.class);
     when(tokenBroker.verify("validToken")).thenReturn(decodedToken);
 
     SecurityIdentity securityIdentity = mock(SecurityIdentity.class);