Add dynamic HTTP header size configuration

* Including defaults values
* For HTTP server and client

gateway-ha/gateway-ha-config.yml

@@ -6,6 +6,7 @@ requestRouter:
   port: 8080
   name: trinoRouter
   historySize: 1000
+  requestBufferSize: 8 * 1024
 
 dataStore:
   jdbcUrl: jdbc:postgresql://localhost:5432/trino_gateway_db

gateway-ha/src/main/java/io/trino/gateway/ha/config/RequestRouterConfiguration.java

@@ -19,4 +19,13 @@ public class RequestRouterConfiguration {
 
   // Use the certificate between gateway and trino?
   private boolean forwardKeystore;
+
+  // Set size for HttpConfiguration
+  private int outputBufferSize = 32 * 1024;
+  private int requestHeaderSize = 8 * 1024;
+  private int responseHeaderSize = 8 * 1024;
+
+  // Set size for HttpClient
+  private int requestBufferSize = 4 * 1024;
+  private int responseBufferSize = 16 * 1024;
 }

gateway-ha/src/main/java/io/trino/gateway/ha/module/HaGatewayProviderModule.java

@@ -195,6 +195,11 @@ public ProxyServer provideGateway() {
       routerProxyConfig.setKeystorePass(routerConfiguration.getKeystorePass());
       routerProxyConfig.setForwardKeystore(routerConfiguration.isForwardKeystore());
       routerProxyConfig.setPreserveHost("false");
+      routerProxyConfig.setOutputBufferSize(routerConfiguration.getOutputBufferSize());
+      routerProxyConfig.setRequestHeaderSize(routerConfiguration.getRequestHeaderSize());
+      routerProxyConfig.setResponseHeaderSize(routerConfiguration.getResponseHeaderSize());
+      routerProxyConfig.setRequestBufferSize(routerConfiguration.getRequestBufferSize());
+      routerProxyConfig.setResponseHeaderSize(routerConfiguration.getResponseBufferSize());
       ProxyHandler proxyHandler = getProxyHandler();
       gateway = new ProxyServer(routerProxyConfig, proxyHandler);
     }

pom.xml

@@ -69,6 +69,7 @@
         <dep.testng.version>6.10</dep.testng.version>
         <dep.trino.version>387</dep.trino.version>
         <dep.wiremock.version>3.0.1</dep.wiremock.version>
+        <dep.apache.commons.version>3.13.0</dep.apache.commons.version>
 
     </properties>
 
@@ -115,6 +116,12 @@
                 <version>${dep.testng.version}</version>
                 <scope>test</scope>
             </dependency>
+
+            <dependency>
+                <groupId>org.apache.commons</groupId>
+                <artifactId>commons-lang3</artifactId>
+                <version>${dep.apache.commons.version}</version>
+            </dependency>
         </dependencies>
     </dependencyManagement>
 

proxyserver/pom.xml

@@ -49,6 +49,11 @@
             <artifactId>mockwebserver</artifactId>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>org.apache.commons</groupId>
+            <artifactId>commons-lang3</artifactId>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
     <build>
         <plugins>

proxyserver/src/main/java/io/trino/gateway/proxyserver/ProxyServer.java

@@ -63,7 +63,9 @@ private void setupContext(ProxyServerConfiguration config) {
       HttpConfiguration httpsConfig = new HttpConfiguration();
       httpsConfig.setSecureScheme(HttpScheme.HTTPS.asString());
       httpsConfig.setSecurePort(config.getLocalPort());
-      httpsConfig.setOutputBufferSize(32768);
+      httpsConfig.setOutputBufferSize(config.getOutputBufferSize());
+      httpsConfig.setRequestHeaderSize(config.getRequestHeaderSize());
+      httpsConfig.setResponseHeaderSize(config.getResponseHeaderSize());
 
       SecureRequestCustomizer src = new SecureRequestCustomizer();
       src.setStsMaxAge(TimeUnit.SECONDS.toSeconds(2000));
@@ -111,7 +113,6 @@ public void addFilter(Class<? extends Filter> filterClass, String pathSpec) {
   }
 
   public void start() {
-
     try {
       this.server.start();
     } catch (Exception e) {

proxyserver/src/main/java/io/trino/gateway/proxyserver/ProxyServerConfiguration.java

@@ -14,6 +14,11 @@ public class ProxyServerConfiguration {
   private String keystorePath;
   private String keystorePass;
   private boolean forwardKeystore;
+  private int outputBufferSize = 2 * 1024 * 1024;
+  private int requestHeaderSize = 2 * 1024 * 1024;
+  private int responseHeaderSize = 8 * 1024;
+  private int requestBufferSize = 4 * 1024;
+  private int responseBufferSize = 16 * 1024;
 
   protected String getPrefix() {
     return prefix;
@@ -46,4 +51,24 @@ protected boolean isForwardKeystore() {
   protected int getLocalPort() {
     return localPort;
   }
+
+  protected int getOutputBufferSize() {
+    return outputBufferSize;
+  }
+
+  protected int getRequestHeaderSize() {
+    return requestHeaderSize;
+  }
+
+  protected int getResponseHeaderSize() {
+    return responseHeaderSize;
+  }
+
+  protected int getRequestBufferSize() {
+    return requestBufferSize;
+  }
+
+  protected int getResponseBufferSize() {
+    return responseBufferSize;
+  }
 }

proxyserver/src/main/java/io/trino/gateway/proxyserver/ProxyServletImpl.java

@@ -48,6 +48,8 @@ protected HttpClient newHttpClient() {
     HttpClient httpClient = new HttpClient(new HttpClientTransportDynamic(clientConnector));
     httpClient.setMaxConnectionsPerDestination(10000);
     httpClient.setConnectTimeout(TimeUnit.SECONDS.toMillis(60));
+    httpClient.setRequestBufferSize(serverConfig.getRequestBufferSize());
+    httpClient.setResponseBufferSize(serverConfig.getResponseBufferSize());
     return httpClient;
   }
 

proxyserver/src/test/java/io/trino/gateway/proxyserver/TestProxyServer.java

@@ -6,13 +6,10 @@
 import com.squareup.okhttp.mockwebserver.MockWebServer;
 import com.squareup.okhttp.mockwebserver.RecordedRequest;
 
-import io.trino.gateway.proxyserver.ProxyHandler;
-import io.trino.gateway.proxyserver.ProxyServer;
-import io.trino.gateway.proxyserver.ProxyServerConfiguration;
-
 import java.io.IOException;
 import java.util.Random;
 
+import org.apache.commons.lang3.RandomStringUtils;
 import org.apache.http.HttpResponse;
 import org.apache.http.client.methods.HttpGet;
 import org.apache.http.client.methods.HttpUriRequest;
@@ -23,19 +20,14 @@
 
 public class TestProxyServer {
 
+  private static int serverPort;
+  private static MockWebServer backend;
+  private static ProxyServer proxyServer;
+
   @Test
   public void testProxyServer() throws IOException {
     String mockResponseText = "Test1234";
-    int backendPort = 30000 + new Random().nextInt(1000);
-
-    MockWebServer backend = new MockWebServer();
-    backend.enqueue(new MockResponse().setBody(mockResponseText));
-    backend.play(backendPort);
-
-    int serverPort = backendPort + 1;
-    ProxyServerConfiguration config = buildConfig(backend.getUrl("/").toString(), serverPort);
-    ProxyServer proxyServer = new ProxyServer(config, new ProxyHandler());
-
+    setProxyServer(mockResponseText);
     try {
       proxyServer.start();
       CloseableHttpClient httpclient = HttpClientBuilder.create().build();
@@ -51,16 +43,7 @@ public void testProxyServer() throws IOException {
   @Test
   public void testCustomHeader() throws Exception {
     String mockResponseText = "CUSTOM HEADER TEST";
-    int backendPort = 30000 + new Random().nextInt(1000);
-
-    MockWebServer backend = new MockWebServer();
-    backend.enqueue(new MockResponse().setBody(mockResponseText));
-    backend.play(backendPort);
-
-    int serverPort = backendPort + 1;
-    ProxyServerConfiguration config = buildConfig(backend.getUrl("/").toString(), serverPort);
-    ProxyServer proxyServer = new ProxyServer(config, new ProxyHandler());
-
+    setProxyServer(mockResponseText);
     try {
       proxyServer.start();
       CloseableHttpClient httpclient = HttpClientBuilder.create().build();
@@ -79,13 +62,54 @@ public void testCustomHeader() throws Exception {
     }
   }
 
+  @Test
+  public void testLongHeader() throws Exception {
+    String mockResponseText = "CUSTOM LONG HEADER TEST";
+    setProxyServer(mockResponseText);
+    String mockLongHeaderKey = "HEADER_LONG";
+    // Mockserver has max 8k for HTTP Header values so test with header value larger than default 4k
+    int headerLength = 5 * 1024;
+    String mockLongHeaderValue = RandomStringUtils.random(headerLength, true, true);
+    try {
+      proxyServer.start();
+      CloseableHttpClient httpclient = HttpClientBuilder.create().build();
+      HttpUriRequest httpUriRequest = new HttpGet("http://localhost:" + serverPort);
+      httpUriRequest.setHeader(mockLongHeaderKey, mockLongHeaderValue);
+
+      HttpResponse response = httpclient.execute(httpUriRequest);
+      assertEquals(mockResponseText, EntityUtils.toString(response.getEntity()));
+      RecordedRequest recordedRequest = backend.takeRequest();
+      assertEquals(recordedRequest.getHeader(mockLongHeaderKey), mockLongHeaderValue);
+    } finally {
+      proxyServer.close();
+      backend.shutdown();
+    }
+  }
+
   private ProxyServerConfiguration buildConfig(String backendUrl, int localPort) {
     ProxyServerConfiguration config = new ProxyServerConfiguration();
     config.setName("MockBackend");
     config.setPrefix("/");
     config.setPreserveHost("true");
     config.setProxyTo(backendUrl);
     config.setLocalPort(localPort);
+    config.setOutputBufferSize(32 * 1024);
+    config.setResponseHeaderSize(8 * 1024);
+    config.setRequestHeaderSize(8 * 1024);
+    config.setRequestBufferSize(16 * 1024); // default 4 * 1024
+    config.setResponseBufferSize(16 * 1024);
     return config;
   }
+
+  private void setProxyServer(String mockResponseText) throws IOException {
+    int backendPort = 30000 + new Random().nextInt(1000);
+
+    backend = new MockWebServer();
+    backend.enqueue(new MockResponse().setBody(mockResponseText));
+    backend.play(backendPort);
+
+    serverPort = backendPort + 1;
+    ProxyServerConfiguration config = buildConfig(backend.getUrl("/").toString(), serverPort);
+    proxyServer = new ProxyServer(config, new ProxyHandler());
+  }
 }