Added user token to the PolarisPrincipal (#3236)

* Added user token to the PolarisPrincipal

* added redacted

Author: cccs-cat001

polaris-core/src/main/java/org/apache/polaris/core/auth/PolarisPrincipal.java

@@ -20,9 +20,11 @@
 
 import java.security.Principal;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 import org.apache.polaris.core.entity.PrincipalEntity;
 import org.apache.polaris.immutables.PolarisImmutable;
+import org.immutables.value.Value;
 
 /** Represents a {@link Principal} in the Polaris system. */
 @PolarisImmutable
@@ -39,7 +41,28 @@ public interface PolarisPrincipal extends Principal {
    * @param roles the set of roles associated with the principal
    */
   static PolarisPrincipal of(PrincipalEntity principalEntity, Set<String> roles) {
-    return of(principalEntity.getName(), principalEntity.getInternalPropertiesAsMap(), roles);
+    return of(
+        principalEntity.getName(),
+        principalEntity.getInternalPropertiesAsMap(),
+        roles,
+        Optional.empty());
+  }
+
+  /**
+   * Creates a new instance of {@link PolarisPrincipal} from the given {@link PrincipalEntity} and
+   * roles.
+   *
+   * <p>The created principal will have the same ID and name as the {@link PrincipalEntity}, and its
+   * properties will be derived from the internal properties of the entity.
+   *
+   * @param principalEntity the principal entity representing the user or service
+   * @param roles the set of roles associated with the principal
+   * @param token the access token of the current user
+   */
+  static PolarisPrincipal of(
+      PrincipalEntity principalEntity, Set<String> roles, Optional<String> token) {
+    return of(
+        principalEntity.getName(), principalEntity.getInternalPropertiesAsMap(), roles, token);
   }
 
   /**
@@ -51,9 +74,24 @@ static PolarisPrincipal of(PrincipalEntity principalEntity, Set<String> roles) {
    * @param roles the set of roles associated with the principal
    */
   static PolarisPrincipal of(String name, Map<String, String> properties, Set<String> roles) {
+    return of(name, properties, roles, Optional.empty());
+  }
+
+  /**
+   * Creates a new instance of {@link PolarisPrincipal} with the specified ID, name, roles, and
+   * properties.
+   *
+   * @param name the name of the principal
+   * @param properties additional properties associated with the principal
+   * @param roles the set of roles associated with the principal
+   * @param token the access token of the current user
+   */
+  static PolarisPrincipal of(
+      String name, Map<String, String> properties, Set<String> roles, Optional<String> token) {
     return ImmutablePolarisPrincipal.builder()
         .name(name)
         .properties(properties)
+        .token(token)
         .roles(roles)
         .build();
   }
@@ -74,4 +112,8 @@ static PolarisPrincipal of(String name, Map<String, String> properties, Set<Stri
    * as permissions, preferences, or other metadata.
    */
   Map<String, String> getProperties();
+
+  /** Optionally returns the access token of the current user. */
+  @Value.Redacted
+  Optional<String> getToken();
 }

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

@@ -22,6 +22,7 @@
 import io.smallrye.common.annotation.Identifier;
 import jakarta.enterprise.context.RequestScoped;
 import jakarta.inject.Inject;
+import java.util.Optional;
 import java.util.Set;
 import java.util.function.Predicate;
 import java.util.stream.Collectors;
@@ -94,7 +95,9 @@ public PolarisPrincipal authenticate(PolarisCredential credentials) {
 
     PrincipalEntity principalEntity = resolvePrincipalEntity(credentials);
     Set<String> principalRoles = resolvePrincipalRoles(credentials, principalEntity);
-    PolarisPrincipal polarisPrincipal = PolarisPrincipal.of(principalEntity, principalRoles);
+    PolarisPrincipal polarisPrincipal =
+        PolarisPrincipal.of(
+            principalEntity, principalRoles, Optional.ofNullable(credentials.getToken()));
 
     LOGGER.debug("Resolved principal: {}", polarisPrincipal);
     return polarisPrincipal;

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

@@ -32,10 +32,19 @@ public interface PolarisCredential extends Credential {
 
   static PolarisCredential of(
       @Nullable Long principalId, @Nullable String principalName, Set<String> principalRoles) {
+    return of(principalId, principalName, principalRoles, null);
+  }
+
+  static PolarisCredential of(
+      @Nullable Long principalId,
+      @Nullable String principalName,
+      Set<String> principalRoles,
+      @Nullable String token) {
     return ImmutablePolarisCredential.builder()
         .principalId(principalId)
         .principalName(principalName)
         .principalRoles(principalRoles)
+        .token(token)
         .build();
   }
 
@@ -49,4 +58,8 @@ static PolarisCredential of(
 
   /** The principal roles, or empty if the principal has no roles. */
   Set<String> getPrincipalRoles();
+
+  /** The access token of the current user, or null if not applicable. */
+  @Nullable
+  String getToken();
 }

runtime/service/src/main/java/org/apache/polaris/service/auth/external/OidcPolarisCredentialAugmentor.java

@@ -90,7 +90,12 @@ protected SecurityIdentity setPolarisCredential(
         principalMapper.mapPrincipalId(identity).stream().boxed().findFirst().orElse(null);
     String principalName = principalMapper.mapPrincipalName(identity).orElse(null);
     Set<String> principalRoles = rolesMapper.mapPrincipalRoles(identity);
-    PolarisCredential credential = PolarisCredential.of(principalId, principalName, principalRoles);
+    String token = null;
+    if (identity.getPrincipal() instanceof JsonWebToken jwt) {
+      token = jwt.getRawToken();
+    }
+    PolarisCredential credential =
+        PolarisCredential.of(principalId, principalName, principalRoles, token);
     // Note: we don't change the identity roles here, this will be done later on
     // by the AuthenticatingAugmentor, which will also validate them.
     return QuarkusSecurityIdentity.builder(identity).addCredential(credential).build();

runtime/service/src/test/java/org/apache/polaris/service/auth/external/OidcPolarisCredentialAugmentorTest.java

@@ -100,7 +100,7 @@ public void testAugmentNonOidcPrincipal() {
   }
 
   @Test
-  public void testAugmentOidcPrincipal() {
+  public void testAugmentOidcPrincipalWithToken() {
     // Given
     JsonWebToken oidcPrincipal = mock(JsonWebToken.class);
     SecurityIdentity identity =
@@ -109,7 +109,35 @@ public void testAugmentOidcPrincipal() {
             .addRole("ROLE1")
             .addAttribute(TENANT_CONFIG_ATTRIBUTE, config)
             .build();
+    when(oidcPrincipal.getRawToken()).thenReturn("this_is_a_token");
+    when(principalMapper.mapPrincipalId(identity)).thenReturn(OptionalLong.of(123L));
+    when(principalMapper.mapPrincipalName(identity)).thenReturn(Optional.of("root"));
+    when(principalRolesMapper.mapPrincipalRoles(identity)).thenReturn(Set.of("MAPPED_ROLE1"));
+
+    // When
+    SecurityIdentity result =
+        augmentor.augment(identity, Uni.createFrom()::item).await().indefinitely();
+
+    // Then
+    assertThat(result).isNotNull();
+    assertThat(result.getPrincipal()).isSameAs(oidcPrincipal);
+    assertThat(result.getCredential(PolarisCredential.class))
+        .isEqualTo(PolarisCredential.of(123L, "root", Set.of("MAPPED_ROLE1"), "this_is_a_token"));
+    // the identity roles should not change, since this is done by the ActiveRolesAugmentor
+    assertThat(result.getRoles()).containsExactlyInAnyOrder("ROLE1");
+  }
 
+  @Test
+  public void testAugmentOidcPrincipalWithNoToken() {
+    // Given
+    JsonWebToken oidcPrincipal = mock(JsonWebToken.class);
+    SecurityIdentity identity =
+        QuarkusSecurityIdentity.builder()
+            .setPrincipal(oidcPrincipal)
+            .addRole("ROLE1")
+            .addAttribute(TENANT_CONFIG_ATTRIBUTE, config)
+            .build();
+    when(oidcPrincipal.getRawToken()).thenReturn(null);
     when(principalMapper.mapPrincipalId(identity)).thenReturn(OptionalLong.of(123L));
     when(principalMapper.mapPrincipalName(identity)).thenReturn(Optional.of("root"));
     when(principalRolesMapper.mapPrincipalRoles(identity)).thenReturn(Set.of("MAPPED_ROLE1"));
@@ -122,7 +150,7 @@ public void testAugmentOidcPrincipal() {
     assertThat(result).isNotNull();
     assertThat(result.getPrincipal()).isSameAs(oidcPrincipal);
     assertThat(result.getCredential(PolarisCredential.class))
-        .isEqualTo(PolarisCredential.of(123L, "root", Set.of("MAPPED_ROLE1")));
+        .isEqualTo(PolarisCredential.of(123L, "root", Set.of("MAPPED_ROLE1"), null));
     // the identity roles should not change, since this is done by the ActiveRolesAugmentor
     assertThat(result.getRoles()).containsExactlyInAnyOrder("ROLE1");
   }