Skip to content

Commit

Permalink
Make JWT brokers injectable
Browse files Browse the repository at this point in the history
* Detach `TokenBrokerFactory` sub-classes from Jackson-based config parsing logic.

* Add a new unified `TokenBrokerConfig` under DW to hold JWT config data.

* Resolve JWT factories through HK2 named service lookup, using the type
  name from `TokenBrokerConfig`

* Add dedicated class `NoneTokenBrokerFactory` for the default do-nothing
  token broker implementation.

* Yaml config (`polaris-server.yml`) is not affected
  • Loading branch information
dimas-b committed Dec 17, 2024
1 parent cba706f commit 29bbfb3
Show file tree
Hide file tree
Showing 7 changed files with 210 additions and 75 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,6 @@
import java.util.Date;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.function.Supplier;
import org.apache.commons.lang3.StringUtils;
Expand All @@ -41,10 +40,8 @@
import org.apache.polaris.core.persistence.MetaStoreManagerFactory;
import org.apache.polaris.core.storage.PolarisStorageIntegrationProvider;
import org.apache.polaris.service.auth.Authenticator;
import org.apache.polaris.service.auth.DecodedToken;
import org.apache.polaris.service.auth.TokenBroker;
import org.apache.polaris.service.auth.TokenBrokerFactory;
import org.apache.polaris.service.auth.TokenResponse;
import org.apache.polaris.service.auth.TokenBrokerFactoryConfig;
import org.apache.polaris.service.catalog.api.IcebergRestOAuth2ApiService;
import org.apache.polaris.service.catalog.io.FileIOFactory;
import org.apache.polaris.service.config.DefaultConfigurationStore;
Expand All @@ -54,7 +51,6 @@
import org.apache.polaris.service.ratelimiter.RateLimiter;
import org.apache.polaris.service.ratelimiter.TokenBucketFactory;
import org.apache.polaris.service.storage.PolarisStorageIntegrationProviderImpl;
import org.apache.polaris.service.types.TokenType;
import org.glassfish.hk2.api.Factory;
import org.glassfish.hk2.api.ServiceLocator;
import org.glassfish.hk2.api.TypeLiteral;
Expand Down Expand Up @@ -92,7 +88,7 @@ public class PolarisApplicationConfig extends Configuration {
private FileIOFactory fileIOFactory;
private RateLimiter rateLimiter;
private TokenBucketFactory tokenBucketFactory;
private TokenBrokerFactory tokenBrokerFactory;
private TokenBrokerConfig tokenBroker = new TokenBrokerConfig();

private AccessToken gcpAccessToken;

Expand Down Expand Up @@ -131,7 +127,13 @@ protected void configure() {
bindFactory(SupplierFactory.create(serviceLocator, config::getPolarisAuthenticator))
.to(Authenticator.class)
.ranked(OVERRIDE_BINDING_RANK);
bindFactory(SupplierFactory.create(serviceLocator, config::getTokenBrokerFactory))
bindFactory(SupplierFactory.create(serviceLocator, () -> tokenBroker))
.to(TokenBrokerFactoryConfig.class);
bindFactory(
SupplierFactory.create(
serviceLocator,
() ->
serviceLocator.getService(TokenBrokerFactory.class, tokenBroker.getType())))
.to(TokenBrokerFactory.class)
.ranked(OVERRIDE_BINDING_RANK);
bindFactory(SupplierFactory.create(serviceLocator, config::getOauth2Service))
Expand Down Expand Up @@ -228,45 +230,8 @@ private Authenticator<String, AuthenticatedPolarisPrincipal> getPolarisAuthentic
}

@JsonProperty("tokenBroker")
@JsonTypeInfo(use = JsonTypeInfo.Id.NAME, include = JsonTypeInfo.As.PROPERTY, property = "type")
public void setTokenBrokerFactory(TokenBrokerFactory tokenBrokerFactory) {
this.tokenBrokerFactory = tokenBrokerFactory;
}

private TokenBrokerFactory getTokenBrokerFactory() {
// return a no-op implementation if none is specified
return Objects.requireNonNullElseGet(
tokenBrokerFactory,
() ->
(rc) ->
new TokenBroker() {
@Override
public boolean supportsGrantType(String grantType) {
return false;
}

@Override
public boolean supportsRequestedTokenType(TokenType tokenType) {
return false;
}

@Override
public TokenResponse generateFromClientSecrets(
String clientId, String clientSecret, String grantType, String scope) {
return null;
}

@Override
public TokenResponse generateFromToken(
TokenType tokenType, String subjectToken, String grantType, String scope) {
return null;
}

@Override
public DecodedToken verify(String token) {
return null;
}
});
public void setTokenBroker(TokenBrokerConfig tokenBroker) {
this.tokenBroker = tokenBroker;
}

private RealmContextResolver getRealmContextResolver() {
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,75 @@
/*
* 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.dropwizard.config;

import com.fasterxml.jackson.annotation.JsonProperty;
import jakarta.annotation.Nullable;
import org.apache.polaris.service.auth.TokenBrokerFactoryConfig;

/**
* This object receives the subsection of the server config YAML (`polaris-server.yml`) that is
* related to JWT brokers.
*/
public class TokenBrokerConfig implements TokenBrokerFactoryConfig {
private String type = "none";
private int maxTokenGenerationInSeconds = 3600;
private String file;
private String secret;

@JsonProperty("type")
public void setType(String type) {
this.type = type;
}

@JsonProperty("file")
public void setFile(String file) {
this.file = file;
}

@JsonProperty("secret")
public void setSecret(String secret) {
this.secret = secret;
}

@JsonProperty("maxTokenGenerationInSeconds")
public void setMaxTokenGenerationInSeconds(int maxTokenGenerationInSeconds) {
this.maxTokenGenerationInSeconds = maxTokenGenerationInSeconds;
}

public String getType() {
return type;
}

@Override
public int maxTokenGenerationInSeconds() {
return maxTokenGenerationInSeconds;
}

@Nullable
@Override
public String file() {
return file;
}

@Nullable
@Override
public String secret() {
return secret;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,11 @@ contract={org.apache.polaris.service.auth.TokenBrokerFactory}
name=rsa-key-pair
qualifier={io.smallrye.common.annotation.Identifier}

[org.apache.polaris.service.auth.NoneTokenBrokerFactory]S
contract={org.apache.polaris.service.auth.TokenBrokerFactory}
name=none
qualifier={io.smallrye.common.annotation.Identifier}

[org.apache.polaris.service.context.DefaultRealmContextResolver]S
contract={org.apache.polaris.service.context.RealmContextResolver}
name=default
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -25,18 +25,21 @@

@Identifier("rsa-key-pair")
public class JWTRSAKeyPairFactory implements TokenBrokerFactory {
private int maxTokenGenerationInSeconds = 3600;

@Inject private MetaStoreManagerFactory metaStoreManagerFactory;
private final TokenBrokerFactoryConfig config;
private final MetaStoreManagerFactory metaStoreManagerFactory;

public void setMaxTokenGenerationInSeconds(int maxTokenGenerationInSeconds) {
this.maxTokenGenerationInSeconds = maxTokenGenerationInSeconds;
@Inject
public JWTRSAKeyPairFactory(
TokenBrokerFactoryConfig config, MetaStoreManagerFactory metaStoreManagerFactory) {
this.config = config;
this.metaStoreManagerFactory = metaStoreManagerFactory;
}

@Override
public TokenBroker apply(RealmContext realmContext) {
return new JWTRSAKeyPair(
metaStoreManagerFactory.getOrCreateMetaStoreManager(realmContext),
maxTokenGenerationInSeconds);
config.maxTokenGenerationInSeconds());
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -23,52 +23,44 @@
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.Objects;
import java.util.function.Supplier;
import org.apache.polaris.core.context.RealmContext;
import org.apache.polaris.core.persistence.MetaStoreManagerFactory;

@Identifier("symmetric-key")
public class JWTSymmetricKeyFactory implements TokenBrokerFactory {
@Inject private MetaStoreManagerFactory metaStoreManagerFactory;
private int maxTokenGenerationInSeconds = 3600;
private String file;
private String secret;

private final MetaStoreManagerFactory metaStoreManagerFactory;
private final TokenBrokerFactoryConfig config;

@Inject
public JWTSymmetricKeyFactory(
MetaStoreManagerFactory metaStoreManagerFactory, TokenBrokerFactoryConfig config) {
this.metaStoreManagerFactory = metaStoreManagerFactory;
this.config = config;
}

@Override
public TokenBroker apply(RealmContext realmContext) {
if (file == null && secret == null) {
String secret = config.secret();
if (config.file() == null && secret == null) {
throw new IllegalStateException("Either file or secret must be set");
}
Supplier<String> secretSupplier = secret != null ? () -> secret : readSecretFromDisk();
return new JWTSymmetricKeyBroker(
metaStoreManagerFactory.getOrCreateMetaStoreManager(realmContext),
maxTokenGenerationInSeconds,
config.maxTokenGenerationInSeconds(),
secretSupplier);
}

private Supplier<String> readSecretFromDisk() {
return () -> {
try {
return Files.readString(Paths.get(file));
return Files.readString(Paths.get(Objects.requireNonNull(config.file())));
} catch (IOException e) {
throw new RuntimeException("Failed to read secret from file: " + file, e);
throw new RuntimeException("Failed to read secret from file: " + config.file(), e);
}
};
}

public void setMaxTokenGenerationInSeconds(int maxTokenGenerationInSeconds) {
this.maxTokenGenerationInSeconds = maxTokenGenerationInSeconds;
}

public void setFile(String file) {
this.file = file;
}

public void setSecret(String secret) {
this.secret = secret;
}

public void setMetaStoreManagerFactory(MetaStoreManagerFactory metaStoreManagerFactory) {
this.metaStoreManagerFactory = metaStoreManagerFactory;
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,59 @@
/*
* 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 io.smallrye.common.annotation.Identifier;
import org.apache.polaris.core.context.RealmContext;
import org.apache.polaris.service.types.TokenType;

/** Default {@link TokenBrokerFactory} that produces token brokers that do not do anything. */
@Identifier("none")
public class NoneTokenBrokerFactory implements TokenBrokerFactory {
@Override
public TokenBroker apply(RealmContext realmContext) {
return new TokenBroker() {
@Override
public boolean supportsGrantType(String grantType) {
return false;
}

@Override
public boolean supportsRequestedTokenType(TokenType tokenType) {
return false;
}

@Override
public TokenResponse generateFromClientSecrets(
String clientId, String clientSecret, String grantType, String scope) {
return null;
}

@Override
public TokenResponse generateFromToken(
TokenType tokenType, String subjectToken, String grantType, String scope) {
return null;
}

@Override
public DecodedToken verify(String token) {
return null;
}
};
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,36 @@
/*
* 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 jakarta.annotation.Nullable;

/**
* This is a union of configuration settings of all token brokers.
*
* @see TokenBrokerFactory
*/
public interface TokenBrokerFactoryConfig {
@Nullable
String file();

@Nullable
String secret();

int maxTokenGenerationInSeconds();
}

0 comments on commit 29bbfb3

Please sign in to comment.