Add a config option to limit request body size (#276)

polaris-server.yml

@@ -163,3 +163,6 @@ logging:
       archivedLogFilenamePattern: ./logs/polaris-%d.log.gz
       # The maximum number of log files to archive.
       archivedFileCount: 14
+
+# Limits the size of request bodies sent to Polaris. -1 means no limit.
+maxRequestBodyBytes: -1

polaris-service/src/main/java/org/apache/polaris/service/PolarisApplication.java

@@ -20,9 +20,11 @@
 
 import static java.nio.charset.StandardCharsets.UTF_8;
 import static java.util.Objects.requireNonNull;
+import static org.apache.polaris.service.config.PolarisApplicationConfig.REQUEST_BODY_BYTES_NO_LIMIT;
 
 import com.fasterxml.jackson.annotation.JsonAutoDetect;
 import com.fasterxml.jackson.annotation.PropertyAccessor;
+import com.fasterxml.jackson.core.StreamReadConstraints;
 import com.fasterxml.jackson.databind.DeserializationFeature;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.fasterxml.jackson.databind.PropertyNamingStrategies;
@@ -104,6 +106,7 @@
 import org.apache.polaris.service.task.TableCleanupTaskHandler;
 import org.apache.polaris.service.task.TaskExecutorImpl;
 import org.apache.polaris.service.task.TaskFileIOSupplier;
+import org.apache.polaris.service.throttling.StreamReadConstraintsExceptionMapper;
 import org.apache.polaris.service.tracing.OpenTelemetryAware;
 import org.apache.polaris.service.tracing.TracingFilter;
 import org.eclipse.jetty.servlets.CrossOriginFilter;
@@ -273,6 +276,17 @@ public void run(PolarisApplicationConfig configuration, Environment environment)
     objectMapper.setVisibility(PropertyAccessor.FIELD, JsonAutoDetect.Visibility.ANY);
     objectMapper.configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false);
     objectMapper.setPropertyNamingStrategy(new PropertyNamingStrategies.KebabCaseStrategy());
+
+    long maxRequestBodyBytes = configuration.getMaxRequestBodyBytes();
+    if (maxRequestBodyBytes != REQUEST_BODY_BYTES_NO_LIMIT) {
+      objectMapper
+          .getFactory()
+          .setStreamReadConstraints(
+              StreamReadConstraints.builder().maxDocumentLength(maxRequestBodyBytes).build());
+      LOGGER.info("Limiting request body size to {} bytes", maxRequestBodyBytes);
+    }
+
+    environment.jersey().register(new StreamReadConstraintsExceptionMapper());
     RESTSerializers.registerAll(objectMapper);
     Serializers.registerSerializers(objectMapper);
     environment.jersey().register(new IcebergJsonProcessingExceptionMapper());

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

@@ -19,6 +19,7 @@
 package org.apache.polaris.service.config;
 
 import com.fasterxml.jackson.annotation.JsonProperty;
+import com.google.common.base.Preconditions;
 import io.dropwizard.core.Configuration;
 import java.util.HashMap;
 import java.util.List;
@@ -55,6 +56,9 @@ public class PolarisApplicationConfig extends Configuration {
   private String awsSecretKey;
   private FileIOFactory fileIOFactory;
 
+  public static final long REQUEST_BODY_BYTES_NO_LIMIT = -1;
+  private long maxRequestBodyBytes = REQUEST_BODY_BYTES_NO_LIMIT;
+
   @JsonProperty("metaStoreManager")
   public void setMetaStoreManagerFactory(MetaStoreManagerFactory metaStoreManagerFactory) {
     this.metaStoreManagerFactory = metaStoreManagerFactory;
@@ -146,6 +150,22 @@ public void setFeatureConfiguration(Map<String, Object> featureConfiguration) {
     this.configurationStore = new DefaultConfigurationStore(featureConfiguration);
   }
 
+  @JsonProperty("maxRequestBodyBytes")
+  public void setMaxRequestBodyBytes(long maxRequestBodyBytes) {
+    // The underlying library that we use to implement the limit treats all values <= 0 as the
+    // same, so we block all but -1 to prevent ambiguity.
+    Preconditions.checkArgument(
+        maxRequestBodyBytes == -1 || maxRequestBodyBytes > 0,
+        "maxRequestBodyBytes must be a positive integer or %s to specify no limit.",
+        REQUEST_BODY_BYTES_NO_LIMIT);
+
+    this.maxRequestBodyBytes = maxRequestBodyBytes;
+  }
+
+  public long getMaxRequestBodyBytes() {
+    return maxRequestBodyBytes;
+  }
+
   public PolarisConfigurationStore getConfigurationStore() {
     return configurationStore;
   }

polaris-service/src/main/java/org/apache/polaris/service/throttling/RequestThrottlingErrorResponse.java

@@ -0,0 +1,45 @@
+/*
+ * 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.throttling;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+/**
+ * Response object for errors caused by DoS-prevention throttling mechanisms, such as request size
+ * limits
+ */
+public class RequestThrottlingErrorResponse {
+  public enum RequestThrottlingErrorType {
+    REQUEST_TOO_LARGE,
+    ;
+  }
+
+  private final RequestThrottlingErrorType errorType;
+
+  @JsonCreator
+  public RequestThrottlingErrorResponse(
+      @JsonProperty("error_type") RequestThrottlingErrorType errorType) {
+    this.errorType = errorType;
+  }
+
+  public RequestThrottlingErrorType getErrorType() {
+    return errorType;
+  }
+}

polaris-service/src/main/java/org/apache/polaris/service/throttling/StreamReadConstraintsExceptionMapper.java

@@ -0,0 +1,42 @@
+/*
+ * 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.throttling;
+
+import com.fasterxml.jackson.core.exc.StreamConstraintsException;
+import jakarta.ws.rs.core.MediaType;
+import jakarta.ws.rs.core.Response;
+import jakarta.ws.rs.ext.ExceptionMapper;
+
+/**
+ * Handles exceptions during the request that are a result of stream constraints such as the request
+ * being too large
+ */
+public class StreamReadConstraintsExceptionMapper
+    implements ExceptionMapper<StreamConstraintsException> {
+
+  @Override
+  public Response toResponse(StreamConstraintsException exception) {
+    return Response.status(Response.Status.BAD_REQUEST)
+        .type(MediaType.APPLICATION_JSON_TYPE)
+        .entity(
+            new RequestThrottlingErrorResponse(
+                RequestThrottlingErrorResponse.RequestThrottlingErrorType.REQUEST_TOO_LARGE))
+        .build();
+  }
+}

polaris-service/src/test/java/org/apache/polaris/service/PolarisApplicationIntegrationTest.java

@@ -26,7 +26,9 @@
 import io.dropwizard.testing.ResourceHelpers;
 import io.dropwizard.testing.junit5.DropwizardAppExtension;
 import io.dropwizard.testing.junit5.DropwizardExtensionsSupport;
+import jakarta.ws.rs.ProcessingException;
 import jakarta.ws.rs.client.Entity;
+import jakarta.ws.rs.client.Invocation;
 import jakarta.ws.rs.core.Response;
 import java.io.IOException;
 import java.nio.file.Files;
@@ -74,6 +76,7 @@
 import org.apache.polaris.service.test.PolarisConnectionExtension;
 import org.apache.polaris.service.test.PolarisRealm;
 import org.apache.polaris.service.test.SnowmanCredentialsExtension;
+import org.apache.polaris.service.throttling.RequestThrottlingErrorResponse;
 import org.assertj.core.api.Assertions;
 import org.assertj.core.api.InstanceOfAssertFactories;
 import org.junit.jupiter.api.AfterAll;
@@ -638,4 +641,63 @@ public void testWarehouseNotSpecified() throws IOException {
           .hasMessage("Malformed request: Please specify a warehouse");
     }
   }
+
+  @Test
+  public void testRequestHeaderTooLarge() {
+    Invocation.Builder request =
+        EXT.client()
+            .target(
+                String.format(
+                    "http://localhost:%d/api/management/v1/principal-roles", EXT.getLocalPort()))
+            .request("application/json");
+
+    // The default limit is 8KiB and each of these headers is at least 8 bytes, so 1500 definitely
+    // exceeds the limit
+    for (int i = 0; i < 1500; i++) {
+      request = request.header("header" + i, "" + i);
+    }
+
+    try {
+      try (Response response =
+          request
+              .header("Authorization", "Bearer " + userToken)
+              .header(REALM_PROPERTY_KEY, realm)
+              .post(Entity.json(new PrincipalRole("r")))) {
+        assertThat(response)
+            .returns(
+                Response.Status.REQUEST_HEADER_FIELDS_TOO_LARGE.getStatusCode(),
+                Response::getStatus);
+      }
+    } catch (ProcessingException e) {
+      // In some runtime environments the request above will return a 431 but in others it'll result
+      // in a ProcessingException from the socket being closed. The test asserts that one of those
+      // things happens.
+    }
+  }
+
+  @Test
+  public void testRequestBodyTooLarge() {
+    // The size is set to be higher than the limit in polaris-server-integrationtest.yml
+    Entity<PrincipalRole> largeRequest = Entity.json(new PrincipalRole("r".repeat(1000001)));
+
+    try (Response response =
+        EXT.client()
+            .target(
+                String.format(
+                    "http://localhost:%d/api/management/v1/principal-roles", EXT.getLocalPort()))
+            .request("application/json")
+            .header("Authorization", "Bearer " + userToken)
+            .header(REALM_PROPERTY_KEY, realm)
+            .post(largeRequest)) {
+      assertThat(response)
+          .returns(Response.Status.BAD_REQUEST.getStatusCode(), Response::getStatus)
+          .matches(
+              r ->
+                  r.readEntity(RequestThrottlingErrorResponse.class)
+                      .getErrorType()
+                      .equals(
+                          RequestThrottlingErrorResponse.RequestThrottlingErrorType
+                              .REQUEST_TOO_LARGE));
+    }
+  }
 }

polaris-service/src/test/resources/polaris-server-integrationtest.yml

@@ -149,3 +149,6 @@ logging:
       archivedLogFilenamePattern: ./logs/iceberg-rest-%d.log.gz
       # The maximum number of log files to archive.
       archivedFileCount: 14
+
+# Limits the size of request bodies sent to Polaris. -1 means no limit.
+maxRequestBodyBytes: 1000000