Prevent DoS attacks by rejecting unknown realms (#594)

Author: adutra

quarkus/service/src/main/resources/application.properties

@@ -72,8 +72,9 @@ quarkus.otel.sdk.disabled=false
 # quarkus.otel.traces.sampler=parentbased_always_on
 # quarkus.otel.traces.sampler.arg=1.0d
 
-polaris.realm-context.default-realm=default-realm
 polaris.realm-context.type=default
+polaris.realm-context.realms=realm1,realm2,realm3
+polaris.realm-context.header-name=Polaris-Realm
 
 polaris.features.defaults."ENFORCE_PRINCIPAL_CREDENTIAL_ROTATION_REQUIRED_CHECKING"=false
 polaris.features.defaults."SUPPORTED_CATALOG_STORAGE_TYPES"=["S3","GCS","AZURE","FILE"]
@@ -85,7 +86,7 @@ polaris.persistence.type=in-memory
 
 polaris.file-io.type=default
 
-polaris.log.request-id-header-name=request_id
+polaris.log.request-id-header-name=Polaris-Request-Id
 # polaris.log.mdc.aid=polaris
 # polaris.log.mdc.sid=polaris-service
 
@@ -138,7 +139,8 @@ polaris.authentication.token-broker.max-token-generation=PT1H
 %test.polaris.features.defaults."INITIALIZE_DEFAULT_CATALOG_FILEIO_FOR_TEST"=true
 %test.polaris.features.defaults."SKIP_CREDENTIAL_SUBSCOPING_INDIRECTION"=true
 %test.polaris.features.defaults."SUPPORTED_CATALOG_STORAGE_TYPES"=["FILE","S3","GCS","AZURE"]
-%test.polaris.realm-context.default-realm=POLARIS
+%test.polaris.realm-context.realms=POLARIS
+%test.polaris.realm-context.type=test
 %test.polaris.storage.aws.access-key=accessKey
 %test.polaris.storage.aws.secret-key=secretKey
 %test.polaris.storage.gcp.token=token

quarkus/service/src/test/java/org/apache/polaris/service/quarkus/TimedApplicationEventListenerTest.java

@@ -18,7 +18,7 @@
  */
 package org.apache.polaris.service.quarkus;
 
-import static org.apache.polaris.service.context.DefaultRealmContextResolver.REALM_PROPERTY_KEY;
+import static org.apache.polaris.service.context.TestRealmContextResolver.REALM_PROPERTY_KEY;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.InstanceOfAssertFactories.type;
 

quarkus/service/src/test/java/org/apache/polaris/service/quarkus/auth/TokenUtils.java

@@ -19,7 +19,7 @@
 package org.apache.polaris.service.quarkus.auth;
 
 import static org.apache.polaris.service.auth.BasePolarisAuthenticator.PRINCIPAL_ROLE_ALL;
-import static org.apache.polaris.service.context.DefaultRealmContextResolver.REALM_PROPERTY_KEY;
+import static org.apache.polaris.service.context.TestRealmContextResolver.REALM_PROPERTY_KEY;
 import static org.assertj.core.api.Assertions.assertThat;
 
 import jakarta.ws.rs.client.Client;

quarkus/service/src/test/java/org/apache/polaris/service/quarkus/catalog/TestUtil.java

@@ -18,7 +18,7 @@
  */
 package org.apache.polaris.service.quarkus.catalog;
 
-import static org.apache.polaris.service.context.DefaultRealmContextResolver.REALM_PROPERTY_KEY;
+import static org.apache.polaris.service.context.TestRealmContextResolver.REALM_PROPERTY_KEY;
 import static org.assertj.core.api.Assertions.assertThat;
 
 import com.google.common.collect.ImmutableMap;

quarkus/service/src/test/java/org/apache/polaris/service/quarkus/ratelimiter/TestUtil.java

@@ -18,7 +18,7 @@
  */
 package org.apache.polaris.service.quarkus.ratelimiter;
 
-import static org.apache.polaris.service.context.DefaultRealmContextResolver.REALM_PROPERTY_KEY;
+import static org.apache.polaris.service.context.TestRealmContextResolver.REALM_PROPERTY_KEY;
 import static org.assertj.core.api.Assertions.assertThat;
 
 import jakarta.ws.rs.core.Response;

quarkus/service/src/test/java/org/apache/polaris/service/quarkus/test/PolarisIntegrationTestFixture.java

@@ -18,7 +18,7 @@
  */
 package org.apache.polaris.service.quarkus.test;
 
-import static org.apache.polaris.service.context.DefaultRealmContextResolver.REALM_PROPERTY_KEY;
+import static org.apache.polaris.service.context.TestRealmContextResolver.REALM_PROPERTY_KEY;
 import static org.assertj.core.api.Assertions.assertThat;
 
 import com.fasterxml.jackson.core.JsonProcessingException;

service/common/src/main/java/org/apache/polaris/service/context/DefaultRealmContextResolver.java

@@ -18,28 +18,15 @@
  */
 package org.apache.polaris.service.context;
 
-import com.google.common.base.Splitter;
 import io.smallrye.common.annotation.Identifier;
 import jakarta.enterprise.context.ApplicationScoped;
 import jakarta.inject.Inject;
-import java.util.HashMap;
 import java.util.Map;
 import org.apache.polaris.core.context.RealmContext;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 
-/**
- * For local/dev testing, this resolver simply expects a custom bearer-token format that is a
- * semicolon-separated list of colon-separated key/value pairs that constitute the realm properties.
- *
- * <p>Example: principal:data-engineer;password:test;realm:acct123
- */
 @ApplicationScoped
 @Identifier("default")
 public class DefaultRealmContextResolver implements RealmContextResolver {
-  private static final Logger LOGGER = LoggerFactory.getLogger(DefaultRealmContextResolver.class);
-
-  public static final String REALM_PROPERTY_KEY = "realm";
 
   private final RealmContextConfiguration configuration;
 
@@ -51,47 +38,18 @@ public DefaultRealmContextResolver(RealmContextConfiguration configuration) {
   @Override
   public RealmContext resolveRealmContext(
       String requestURL, String method, String path, Map<String, String> headers) {
-    // Since this default resolver is strictly for use in test/dev environments, we'll consider
-    // it safe to log all contents. Any "real" resolver used in a prod environment should make
-    // sure to only log non-sensitive contents.
-    LOGGER.debug(
-        "Resolving RealmContext for method: {}, path: {}, headers: {}", method, path, headers);
-    Map<String, String> parsedProperties = parseBearerTokenAsKvPairs(headers);
-
-    if (!parsedProperties.containsKey(REALM_PROPERTY_KEY)
-        && headers.containsKey(REALM_PROPERTY_KEY)) {
-      parsedProperties.put(REALM_PROPERTY_KEY, headers.get(REALM_PROPERTY_KEY));
-    }
 
-    if (!parsedProperties.containsKey(REALM_PROPERTY_KEY)) {
-      LOGGER.warn(
-          "Failed to parse {} from headers; using {}",
-          REALM_PROPERTY_KEY,
-          configuration.defaultRealm());
-      parsedProperties.put(REALM_PROPERTY_KEY, configuration.defaultRealm());
-    }
-    String realmId = parsedProperties.get(REALM_PROPERTY_KEY);
-    return () -> realmId;
-  }
+    String realm;
 
-  /**
-   * Returns kv pairs parsed from the "Authorization: Bearer k1:v1;k2:v2;k3:v3" header if it exists;
-   * if missing, returns empty map.
-   */
-  static Map<String, String> parseBearerTokenAsKvPairs(Map<String, String> headers) {
-    Map<String, String> parsedProperties = new HashMap<>();
-    if (headers != null) {
-      String authHeader = headers.get("Authorization");
-      if (authHeader != null) {
-        String[] parts = authHeader.split(" ");
-        if (parts.length == 2 && "Bearer".equalsIgnoreCase(parts[0])) {
-          if (parts[1].matches("[\\w\\d=_+-]+:[\\w\\d=+_-]+(?:;[\\w\\d=+_-]+:[\\w\\d=+_-]+)*")) {
-            parsedProperties.putAll(
-                Splitter.on(';').trimResults().withKeyValueSeparator(':').split(parts[1]));
-          }
-        }
+    if (headers.containsKey(configuration.headerName())) {
+      realm = headers.get(configuration.headerName());
+      if (!configuration.realms().contains(realm)) {
+        throw new IllegalArgumentException("Unknown realm: " + realm);
       }
+    } else {
+      realm = configuration.defaultRealm();
     }
-    return parsedProperties;
+
+    return () -> realm;
   }
 }

service/common/src/main/java/org/apache/polaris/service/context/RealmContextConfiguration.java

@@ -18,8 +18,23 @@
  */
 package org.apache.polaris.service.context;
 
+import jakarta.validation.constraints.Size;
+import java.util.Set;
+
 public interface RealmContextConfiguration {
 
+  /**
+   * The set of realms that are supported by the realm context resolver. The first realm is
+   * considered the default realm.
+   */
+  @Size(min = 1)
+  Set<String> realms();
+
+  /** The header name that contains the realm identifier. */
+  String headerName();
+
   /** The default realm to use when no realm is specified. */
-  String defaultRealm();
+  default String defaultRealm() {
+    return realms().iterator().next();
+  }
 }

service/common/src/main/java/org/apache/polaris/service/context/TestRealmContextResolver.java

@@ -0,0 +1,97 @@
+/*
+ * 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.context;
+
+import com.google.common.base.Splitter;
+import io.smallrye.common.annotation.Identifier;
+import jakarta.enterprise.context.ApplicationScoped;
+import jakarta.inject.Inject;
+import java.util.HashMap;
+import java.util.Map;
+import org.apache.polaris.core.context.RealmContext;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * For local/dev testing, this resolver simply expects a custom bearer-token format that is a
+ * semicolon-separated list of colon-separated key/value pairs that constitute the realm properties.
+ *
+ * <p>Example: principal:data-engineer;password:test;realm:acct123
+ */
+@ApplicationScoped
+@Identifier("test")
+public class TestRealmContextResolver implements RealmContextResolver {
+  private static final Logger LOGGER = LoggerFactory.getLogger(DefaultRealmContextResolver.class);
+
+  public static final String REALM_PROPERTY_KEY = "realm";
+
+  private final RealmContextConfiguration configuration;
+
+  @Inject
+  public TestRealmContextResolver(RealmContextConfiguration configuration) {
+    this.configuration = configuration;
+  }
+
+  @Override
+  public RealmContext resolveRealmContext(
+      String requestURL, String method, String path, Map<String, String> headers) {
+    // Since this default resolver is strictly for use in test/dev environments, we'll consider
+    // it safe to log all contents. Any "real" resolver used in a prod environment should make
+    // sure to only log non-sensitive contents.
+    LOGGER.debug(
+        "Resolving RealmContext for method: {}, path: {}, headers: {}", method, path, headers);
+    Map<String, String> parsedProperties = parseBearerTokenAsKvPairs(headers);
+
+    if (!parsedProperties.containsKey(REALM_PROPERTY_KEY)
+        && headers.containsKey(REALM_PROPERTY_KEY)) {
+      parsedProperties.put(REALM_PROPERTY_KEY, headers.get(REALM_PROPERTY_KEY));
+    }
+
+    if (!parsedProperties.containsKey(REALM_PROPERTY_KEY)) {
+      LOGGER.warn(
+          "Failed to parse {} from headers; using {}",
+          REALM_PROPERTY_KEY,
+          configuration.defaultRealm());
+      parsedProperties.put(REALM_PROPERTY_KEY, configuration.defaultRealm());
+    }
+    String realmId = parsedProperties.get(REALM_PROPERTY_KEY);
+    return () -> realmId;
+  }
+
+  /**
+   * Returns kv pairs parsed from the "Authorization: Bearer k1:v1;k2:v2;k3:v3" header if it exists;
+   * if missing, returns empty map.
+   */
+  private static Map<String, String> parseBearerTokenAsKvPairs(Map<String, String> headers) {
+    Map<String, String> parsedProperties = new HashMap<>();
+    if (headers != null) {
+      String authHeader = headers.get("Authorization");
+      if (authHeader != null) {
+        String[] parts = authHeader.split(" ");
+        if (parts.length == 2 && "Bearer".equalsIgnoreCase(parts[0])) {
+          if (parts[1].matches("[\\w\\d=_+-]+:[\\w\\d=+_-]+(?:;[\\w\\d=+_-]+:[\\w\\d=+_-]+)*")) {
+            parsedProperties.putAll(
+                Splitter.on(';').trimResults().withKeyValueSeparator(':').split(parts[1]));
+          }
+        }
+      }
+    }
+    return parsedProperties;
+  }
+}