Make JWT brokers injectable

dropwizard/service/src/main/java/org/apache/polaris/service/dropwizard/config/PolarisApplicationConfig.java

@@ -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;
@@ -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;
@@ -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;
@@ -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;
 
@@ -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))
@@ -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() {

dropwizard/service/src/main/java/org/apache/polaris/service/dropwizard/config/TokenBrokerConfig.java

@@ -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;
+  }
+}

dropwizard/service/src/main/resources/META-INF/hk2-locator/default

@@ -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

service/common/src/main/java/org/apache/polaris/service/auth/JWTRSAKeyPairFactory.java

@@ -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());
   }
 }

service/common/src/main/java/org/apache/polaris/service/auth/JWTSymmetricKeyFactory.java

@@ -18,57 +18,52 @@
  */
 package org.apache.polaris.service.auth;
 
+import static com.google.common.base.Preconditions.checkState;
+
 import io.smallrye.common.annotation.Identifier;
 import jakarta.inject.Inject;
 import java.io.IOException;
 import java.nio.file.Files;
+import java.nio.file.Path;
 import java.nio.file.Paths;
 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;
+  private final Supplier<String> secretSupplier;
+
+  @Inject
+  public JWTSymmetricKeyFactory(
+      MetaStoreManagerFactory metaStoreManagerFactory, TokenBrokerFactoryConfig config) {
+    this.metaStoreManagerFactory = metaStoreManagerFactory;
+    this.config = config;
+
+    String secret = config.secret();
+    String file = config.file();
+    checkState(secret != null || file != null, "Either file or secret must be set");
+    this.secretSupplier = secret != null ? () -> secret : readSecretFromDisk(Paths.get(file));
+  }
 
   @Override
   public TokenBroker apply(RealmContext realmContext) {
-    if (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() {
+  private Supplier<String> readSecretFromDisk(Path path) {
     return () -> {
       try {
-        return Files.readString(Paths.get(file));
+        return Files.readString(path);
       } 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;
-  }
 }

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

@@ -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;
+      }
+    };
+  }
+}

service/common/src/main/java/org/apache/polaris/service/auth/TokenBrokerFactoryConfig.java

@@ -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();
+}