Skip to content

Commit fa47229

Browse files
XN137XN137
authored
Embrace request-scoped TokenBroker (#3024)
* Embrace request-scoped TokenBroker `TokenBroker` and `CallContext` are both request-scoped, so instead of passing the former into the latter, we can do this via the `TokenBrokerFactory` and thus simplify the `TokenBroker` interface.
1 parent ac724d3 commit fa47229

File tree

12 files changed

+91
-153
lines changed

12 files changed

+91
-153
lines changed

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

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -49,10 +49,15 @@ public abstract class JWTBroker implements TokenBroker {
4949
private static final String CLAIM_KEY_SCOPE = "scope";
5050

5151
private final PolarisMetaStoreManager metaStoreManager;
52+
private final PolarisCallContext polarisCallContext;
5253
private final int maxTokenGenerationInSeconds;
5354

54-
JWTBroker(PolarisMetaStoreManager metaStoreManager, int maxTokenGenerationInSeconds) {
55+
JWTBroker(
56+
PolarisMetaStoreManager metaStoreManager,
57+
PolarisCallContext polarisCallContext,
58+
int maxTokenGenerationInSeconds) {
5559
this.metaStoreManager = metaStoreManager;
60+
this.polarisCallContext = polarisCallContext;
5661
this.maxTokenGenerationInSeconds = maxTokenGenerationInSeconds;
5762
}
5863

@@ -86,7 +91,6 @@ public TokenResponse generateFromToken(
8691
String subjectToken,
8792
String grantType,
8893
String scope,
89-
PolarisCallContext polarisCallContext,
9094
TokenType requestedTokenType) {
9195
if (requestedTokenType != null && !TokenType.ACCESS_TOKEN.equals(requestedTokenType)) {
9296
return TokenResponse.of(OAuthError.invalid_request);
@@ -125,7 +129,6 @@ public TokenResponse generateFromClientSecrets(
125129
String clientSecret,
126130
String grantType,
127131
String scope,
128-
PolarisCallContext polarisCallContext,
129132
TokenType requestedTokenType) {
130133
// Initial sanity checks
131134
TokenRequestValidator validator = new TokenRequestValidator();
@@ -135,8 +138,7 @@ public TokenResponse generateFromClientSecrets(
135138
return TokenResponse.of(initialValidationResponse.get());
136139
}
137140

138-
Optional<PrincipalEntity> principal =
139-
findPrincipalEntity(clientId, clientSecret, polarisCallContext);
141+
Optional<PrincipalEntity> principal = findPrincipalEntity(clientId, clientSecret);
140142
if (principal.isEmpty()) {
141143
return TokenResponse.of(OAuthError.unauthorized_client);
142144
}
@@ -176,8 +178,7 @@ private String scopes(String scope) {
176178
return scope == null || scope.isBlank() ? DefaultAuthenticator.PRINCIPAL_ROLE_ALL : scope;
177179
}
178180

179-
private Optional<PrincipalEntity> findPrincipalEntity(
180-
String clientId, String clientSecret, PolarisCallContext polarisCallContext) {
181+
private Optional<PrincipalEntity> findPrincipalEntity(String clientId, String clientSecret) {
181182
// Validate the principal is present and secrets match
182183
PrincipalSecretsResult principalSecrets =
183184
metaStoreManager.loadPrincipalSecrets(polarisCallContext, clientId);

‎runtime/service/src/main/java/org/apache/polaris/service/auth/internal/broker/RSAKeyPairJWTBroker.java‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import com.auth0.jwt.algorithms.Algorithm;
2222
import java.security.interfaces.RSAPrivateKey;
2323
import java.security.interfaces.RSAPublicKey;
24+
import org.apache.polaris.core.PolarisCallContext;
2425
import org.apache.polaris.core.persistence.PolarisMetaStoreManager;
2526

2627
/** Generates a JWT using a Public/Private RSA Key */
@@ -30,9 +31,10 @@ public class RSAKeyPairJWTBroker extends JWTBroker {
3031

3132
RSAKeyPairJWTBroker(
3233
PolarisMetaStoreManager metaStoreManager,
34+
PolarisCallContext polarisCallContext,
3335
int maxTokenGenerationInSeconds,
3436
KeyProvider keyProvider) {
35-
super(metaStoreManager, maxTokenGenerationInSeconds);
37+
super(metaStoreManager, polarisCallContext, maxTokenGenerationInSeconds);
3638
this.keyProvider = keyProvider;
3739
}
3840

‎runtime/service/src/main/java/org/apache/polaris/service/auth/internal/broker/RSAKeyPairJWTBrokerFactory.java‎

Lines changed: 15 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -25,8 +25,8 @@
2525
import java.time.Duration;
2626
import java.util.concurrent.ConcurrentHashMap;
2727
import java.util.concurrent.ConcurrentMap;
28+
import org.apache.polaris.core.PolarisCallContext;
2829
import org.apache.polaris.core.context.RealmContext;
29-
import org.apache.polaris.core.persistence.MetaStoreManagerFactory;
3030
import org.apache.polaris.core.persistence.PolarisMetaStoreManager;
3131
import org.apache.polaris.service.auth.AuthenticationConfiguration;
3232
import org.apache.polaris.service.auth.AuthenticationRealmConfiguration;
@@ -36,38 +36,32 @@
3636
@Identifier("rsa-key-pair")
3737
public class RSAKeyPairJWTBrokerFactory implements TokenBrokerFactory {
3838

39-
private final MetaStoreManagerFactory metaStoreManagerFactory;
4039
private final AuthenticationConfiguration authenticationConfiguration;
4140

42-
private final ConcurrentMap<String, RSAKeyPairJWTBroker> tokenBrokers = new ConcurrentHashMap<>();
41+
private final ConcurrentMap<String, KeyProvider> keyProviders = new ConcurrentHashMap<>();
4342

4443
@Inject
45-
public RSAKeyPairJWTBrokerFactory(
46-
MetaStoreManagerFactory metaStoreManagerFactory,
47-
AuthenticationConfiguration authenticationConfiguration) {
48-
this.metaStoreManagerFactory = metaStoreManagerFactory;
44+
public RSAKeyPairJWTBrokerFactory(AuthenticationConfiguration authenticationConfiguration) {
4945
this.authenticationConfiguration = authenticationConfiguration;
5046
}
5147

5248
@Override
53-
public TokenBroker apply(RealmContext realmContext) {
54-
return tokenBrokers.computeIfAbsent(
55-
realmContext.getRealmIdentifier(), k -> createTokenBroker(realmContext));
56-
}
57-
58-
private RSAKeyPairJWTBroker createTokenBroker(RealmContext realmContext) {
49+
public TokenBroker create(
50+
PolarisMetaStoreManager metaStoreManager, PolarisCallContext polarisCallContext) {
51+
RealmContext realmContext = polarisCallContext.getRealmContext();
5952
AuthenticationRealmConfiguration config = authenticationConfiguration.forRealm(realmContext);
6053
Duration maxTokenGeneration = config.tokenBroker().maxTokenGeneration();
6154
KeyProvider keyProvider =
62-
config
63-
.tokenBroker()
64-
.rsaKeyPair()
65-
.map(this::fileSystemKeyPair)
66-
.orElseGet(this::generateEphemeralKeyPair);
67-
PolarisMetaStoreManager metaStoreManager =
68-
metaStoreManagerFactory.getOrCreateMetaStoreManager(realmContext);
55+
keyProviders.computeIfAbsent(
56+
realmContext.getRealmIdentifier(),
57+
k ->
58+
config
59+
.tokenBroker()
60+
.rsaKeyPair()
61+
.map(this::fileSystemKeyPair)
62+
.orElseGet(this::generateEphemeralKeyPair));
6963
return new RSAKeyPairJWTBroker(
70-
metaStoreManager, (int) maxTokenGeneration.toSeconds(), keyProvider);
64+
metaStoreManager, polarisCallContext, (int) maxTokenGeneration.toSeconds(), keyProvider);
7165
}
7266

7367
private KeyProvider fileSystemKeyPair(RSAKeyPairConfiguration config) {

‎runtime/service/src/main/java/org/apache/polaris/service/auth/internal/broker/SymmetricKeyJWTBroker.java‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020

2121
import com.auth0.jwt.algorithms.Algorithm;
2222
import java.util.function.Supplier;
23+
import org.apache.polaris.core.PolarisCallContext;
2324
import org.apache.polaris.core.persistence.PolarisMetaStoreManager;
2425

2526
/** Generates a JWT using a Symmetric Key. */
@@ -28,9 +29,10 @@ public class SymmetricKeyJWTBroker extends JWTBroker {
2829

2930
public SymmetricKeyJWTBroker(
3031
PolarisMetaStoreManager metaStoreManager,
32+
PolarisCallContext polarisCallContext,
3133
int maxTokenGenerationInSeconds,
3234
Supplier<String> secretSupplier) {
33-
super(metaStoreManager, maxTokenGenerationInSeconds);
35+
super(metaStoreManager, polarisCallContext, maxTokenGenerationInSeconds);
3436
this.secretSupplier = secretSupplier;
3537
}
3638

‎runtime/service/src/main/java/org/apache/polaris/service/auth/internal/broker/SymmetricKeyJWTBrokerFactory.java‎

Lines changed: 31 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -27,11 +27,13 @@
2727
import java.nio.file.Files;
2828
import java.nio.file.Path;
2929
import java.time.Duration;
30+
import java.util.Objects;
3031
import java.util.concurrent.ConcurrentHashMap;
3132
import java.util.concurrent.ConcurrentMap;
3233
import java.util.function.Supplier;
34+
import org.apache.polaris.core.PolarisCallContext;
3335
import org.apache.polaris.core.context.RealmContext;
34-
import org.apache.polaris.core.persistence.MetaStoreManagerFactory;
36+
import org.apache.polaris.core.persistence.PolarisMetaStoreManager;
3537
import org.apache.polaris.service.auth.AuthenticationConfiguration;
3638
import org.apache.polaris.service.auth.AuthenticationRealmConfiguration;
3739
import org.apache.polaris.service.auth.AuthenticationRealmConfiguration.TokenBrokerConfiguration.SymmetricKeyConfiguration;
@@ -40,51 +42,46 @@
4042
@Identifier("symmetric-key")
4143
public class SymmetricKeyJWTBrokerFactory implements TokenBrokerFactory {
4244

43-
private final MetaStoreManagerFactory metaStoreManagerFactory;
4445
private final AuthenticationConfiguration authenticationConfiguration;
4546

46-
private final ConcurrentMap<String, SymmetricKeyJWTBroker> tokenBrokers =
47-
new ConcurrentHashMap<>();
47+
private final ConcurrentMap<String, Supplier<String>> secretSuppliers = new ConcurrentHashMap<>();
4848

4949
@Inject
50-
public SymmetricKeyJWTBrokerFactory(
51-
MetaStoreManagerFactory metaStoreManagerFactory,
52-
AuthenticationConfiguration authenticationConfiguration) {
53-
this.metaStoreManagerFactory = metaStoreManagerFactory;
50+
public SymmetricKeyJWTBrokerFactory(AuthenticationConfiguration authenticationConfiguration) {
5451
this.authenticationConfiguration = authenticationConfiguration;
5552
}
5653

5754
@Override
58-
public TokenBroker apply(RealmContext realmContext) {
59-
return tokenBrokers.computeIfAbsent(
60-
realmContext.getRealmIdentifier(), k -> createTokenBroker(realmContext));
61-
}
62-
63-
private SymmetricKeyJWTBroker createTokenBroker(RealmContext realmContext) {
55+
public TokenBroker create(
56+
PolarisMetaStoreManager metaStoreManager, PolarisCallContext polarisCallContext) {
57+
RealmContext realmContext = polarisCallContext.getRealmContext();
6458
AuthenticationRealmConfiguration config = authenticationConfiguration.forRealm(realmContext);
6559
Duration maxTokenGeneration = config.tokenBroker().maxTokenGeneration();
66-
SymmetricKeyConfiguration symmetricKeyConfiguration =
67-
config
68-
.tokenBroker()
69-
.symmetricKey()
70-
.orElseThrow(() -> new IllegalStateException("Symmetric key configuration is missing"));
71-
String secret = symmetricKeyConfiguration.secret().orElse(null);
72-
Path file = symmetricKeyConfiguration.file().orElse(null);
73-
checkState(secret != null || file != null, "Either file or secret must be set");
74-
Supplier<String> secretSupplier = secret != null ? () -> secret : readSecretFromDisk(file);
60+
Supplier<String> secretSupplier =
61+
secretSuppliers.computeIfAbsent(
62+
realmContext.getRealmIdentifier(),
63+
k -> {
64+
SymmetricKeyConfiguration symmetricKeyConfiguration =
65+
config
66+
.tokenBroker()
67+
.symmetricKey()
68+
.orElseThrow(
69+
() ->
70+
new IllegalStateException("Symmetric key configuration is missing"));
71+
String secret = symmetricKeyConfiguration.secret().orElse(null);
72+
Path file = symmetricKeyConfiguration.file().orElse(null);
73+
checkState(secret != null || file != null, "Either file or secret must be set");
74+
return () -> Objects.requireNonNullElseGet(secret, () -> readSecretFromDisk(file));
75+
});
7576
return new SymmetricKeyJWTBroker(
76-
metaStoreManagerFactory.getOrCreateMetaStoreManager(realmContext),
77-
(int) maxTokenGeneration.toSeconds(),
78-
secretSupplier);
77+
metaStoreManager, polarisCallContext, (int) maxTokenGeneration.toSeconds(), secretSupplier);
7978
}
8079

81-
private static Supplier<String> readSecretFromDisk(Path file) {
82-
return () -> {
83-
try {
84-
return Files.readString(file);
85-
} catch (IOException e) {
86-
throw new RuntimeException("Failed to read secret from file: " + file, e);
87-
}
88-
};
80+
private static String readSecretFromDisk(Path file) {
81+
try {
82+
return Files.readString(file);
83+
} catch (IOException e) {
84+
throw new RuntimeException("Failed to read secret from file: " + file, e);
85+
}
8986
}
9087
}

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

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,6 @@
1818
*/
1919
package org.apache.polaris.service.auth.internal.broker;
2020

21-
import org.apache.polaris.core.PolarisCallContext;
2221
import org.apache.polaris.service.auth.PolarisCredential;
2322
import org.apache.polaris.service.types.TokenType;
2423

@@ -39,7 +38,6 @@ TokenResponse generateFromClientSecrets(
3938
final String clientSecret,
4039
final String grantType,
4140
final String scope,
42-
PolarisCallContext polarisCallContext,
4341
TokenType requestedTokenType);
4442

4543
/**
@@ -52,7 +50,6 @@ TokenResponse generateFromToken(
5250
String subjectToken,
5351
final String grantType,
5452
final String scope,
55-
PolarisCallContext polarisCallContext,
5653
TokenType requestedTokenType);
5754

5855
/** Decodes and verifies the token, then returns the associated {@link PolarisCredential}. */

‎runtime/service/src/main/java/org/apache/polaris/service/auth/internal/broker/TokenBrokerFactory.java‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,11 +18,14 @@
1818
*/
1919
package org.apache.polaris.service.auth.internal.broker;
2020

21-
import java.util.function.Function;
22-
import org.apache.polaris.core.context.RealmContext;
21+
import org.apache.polaris.core.PolarisCallContext;
22+
import org.apache.polaris.core.persistence.PolarisMetaStoreManager;
2323

2424
/**
2525
* Factory that creates a {@link TokenBroker} for generating and parsing. The {@link TokenBroker} is
2626
* created based on the realm context.
2727
*/
28-
public interface TokenBrokerFactory extends Function<RealmContext, TokenBroker> {}
28+
public interface TokenBrokerFactory {
29+
TokenBroker create(
30+
PolarisMetaStoreManager metaStoreManager, PolarisCallContext polarisCallContext);
31+
}

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

Lines changed: 3 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,6 @@
2727
import jakarta.ws.rs.core.SecurityContext;
2828
import java.util.Base64;
2929
import org.apache.iceberg.rest.responses.OAuthTokenResponse;
30-
import org.apache.polaris.core.context.CallContext;
3130
import org.apache.polaris.core.context.RealmContext;
3231
import org.apache.polaris.service.auth.internal.broker.TokenBroker;
3332
import org.apache.polaris.service.auth.internal.broker.TokenResponse;
@@ -49,12 +48,10 @@ public class DefaultOAuth2ApiService implements IcebergRestOAuth2ApiService {
4948
private static final String BEARER = "bearer";
5049

5150
private final TokenBroker tokenBroker;
52-
private final CallContext callContext;
5351

5452
@Inject
55-
public DefaultOAuth2ApiService(TokenBroker tokenBroker, CallContext callContext) {
53+
public DefaultOAuth2ApiService(TokenBroker tokenBroker) {
5654
this.tokenBroker = tokenBroker;
57-
this.callContext = callContext;
5855
}
5956

6057
@Override
@@ -104,21 +101,11 @@ public Response getToken(
104101
if (clientSecret != null) {
105102
tokenResponse =
106103
tokenBroker.generateFromClientSecrets(
107-
clientId,
108-
clientSecret,
109-
grantType,
110-
scope,
111-
callContext.getPolarisCallContext(),
112-
requestedTokenType);
104+
clientId, clientSecret, grantType, scope, requestedTokenType);
113105
} else if (subjectToken != null) {
114106
tokenResponse =
115107
tokenBroker.generateFromToken(
116-
subjectTokenType,
117-
subjectToken,
118-
grantType,
119-
scope,
120-
callContext.getPolarisCallContext(),
121-
requestedTokenType);
108+
subjectTokenType, subjectToken, grantType, scope, requestedTokenType);
122109
} else {
123110
return OAuthUtils.getResponseFromError(OAuthError.invalid_request);
124111
}

‎runtime/service/src/main/java/org/apache/polaris/service/config/ServiceProducers.java‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -398,13 +398,14 @@ public IcebergRestOAuth2ApiService icebergRestOAuth2ApiService(
398398
@RequestScoped
399399
public TokenBroker tokenBroker(
400400
AuthenticationRealmConfiguration config,
401-
RealmContext realmContext,
402-
@Any Instance<TokenBrokerFactory> tokenBrokerFactories) {
401+
@Any Instance<TokenBrokerFactory> tokenBrokerFactories,
402+
PolarisMetaStoreManager polarisMetaStoreManager,
403+
CallContext callContext) {
403404
String type =
404405
config.type() == AuthenticationType.EXTERNAL ? "none" : config.tokenBroker().type();
405406
TokenBrokerFactory tokenBrokerFactory =
406407
tokenBrokerFactories.select(Identifier.Literal.of(type)).get();
407-
return tokenBrokerFactory.apply(realmContext);
408+
return tokenBrokerFactory.create(polarisMetaStoreManager, callContext.getPolarisCallContext());
408409
}
409410

410411
// other beans

‎runtime/service/src/test/java/org/apache/polaris/service/auth/internal/broker/JWTSymmetricKeyGeneratorTest.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -52,14 +52,14 @@ public void testJWTSymmetricKeyGenerator() {
5252
new PrincipalEntity.Builder().setId(principalId).setName("principal").build();
5353
Mockito.when(metastoreManager.findPrincipalById(polarisCallContext, principalId))
5454
.thenReturn(Optional.of(principal));
55-
TokenBroker generator = new SymmetricKeyJWTBroker(metastoreManager, 666, () -> "polaris");
55+
TokenBroker generator =
56+
new SymmetricKeyJWTBroker(metastoreManager, polarisCallContext, 666, () -> "polaris");
5657
TokenResponse token =
5758
generator.generateFromClientSecrets(
5859
clientId,
5960
mainSecret,
6061
TokenRequestValidator.CLIENT_CREDENTIALS,
6162
"PRINCIPAL_ROLE:TEST",
62-
polarisCallContext,
6363
TokenType.ACCESS_TOKEN);
6464
assertThat(token).isNotNull();
6565

0 commit comments

Comments
 (0)