diff --git a/polaris-core/src/main/java/org/apache/polaris/core/auth/PolarisPrincipal.java b/polaris-core/src/main/java/org/apache/polaris/core/auth/PolarisPrincipal.java index e3be050147..07e36b8018 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/auth/PolarisPrincipal.java +++ b/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 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. + * + *

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 roles, Optional token) { + return of( + principalEntity.getName(), principalEntity.getInternalPropertiesAsMap(), roles, token); } /** @@ -51,9 +74,24 @@ static PolarisPrincipal of(PrincipalEntity principalEntity, Set roles) { * @param roles the set of roles associated with the principal */ static PolarisPrincipal of(String name, Map properties, Set 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 properties, Set roles, Optional token) { return ImmutablePolarisPrincipal.builder() .name(name) .properties(properties) + .token(token) .roles(roles) .build(); } @@ -74,4 +112,8 @@ static PolarisPrincipal of(String name, Map properties, Set getProperties(); + + /** Optionally returns the access token of the current user. */ + @Value.Redacted + Optional getToken(); } diff --git a/runtime/service/src/main/java/org/apache/polaris/service/auth/DefaultAuthenticator.java b/runtime/service/src/main/java/org/apache/polaris/service/auth/DefaultAuthenticator.java index 83a3a419d7..74867054ff 100644 --- a/runtime/service/src/main/java/org/apache/polaris/service/auth/DefaultAuthenticator.java +++ b/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 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; diff --git a/runtime/service/src/main/java/org/apache/polaris/service/auth/PolarisCredential.java b/runtime/service/src/main/java/org/apache/polaris/service/auth/PolarisCredential.java index f86565e669..05cc5c8547 100644 --- a/runtime/service/src/main/java/org/apache/polaris/service/auth/PolarisCredential.java +++ b/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 principalRoles) { + return of(principalId, principalName, principalRoles, null); + } + + static PolarisCredential of( + @Nullable Long principalId, + @Nullable String principalName, + Set 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 getPrincipalRoles(); + + /** The access token of the current user, or null if not applicable. */ + @Nullable + String getToken(); } diff --git a/runtime/service/src/main/java/org/apache/polaris/service/auth/external/OidcPolarisCredentialAugmentor.java b/runtime/service/src/main/java/org/apache/polaris/service/auth/external/OidcPolarisCredentialAugmentor.java index 30fe86cccd..0a8a8eeb33 100644 --- a/runtime/service/src/main/java/org/apache/polaris/service/auth/external/OidcPolarisCredentialAugmentor.java +++ b/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 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(); diff --git a/runtime/service/src/test/java/org/apache/polaris/service/auth/external/OidcPolarisCredentialAugmentorTest.java b/runtime/service/src/test/java/org/apache/polaris/service/auth/external/OidcPolarisCredentialAugmentorTest.java index 588ead41ca..8aebaa88b2 100644 --- a/runtime/service/src/test/java/org/apache/polaris/service/auth/external/OidcPolarisCredentialAugmentorTest.java +++ b/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"); }