Add option for extra whitelisted paths

trinodb · Oct 25, 2023 · ac633f2 · ac633f2
1 parent a74b0bc
commit ac633f2
Show file tree

Hide file tree

Showing 15 changed files with 105 additions and 4 deletions.
diff --git a/docs/configuration.md b/docs/configuration.md
@@ -0,0 +1,45 @@
+**Trino Gateway documentation**
+
+<table>
+  <tr>
+    <td><a href="configuration.md">Configuration</a></td>
+    <td><a href="design.md">Design</a></td>
+    <td><b><a href="development.md">Development</a></b></td>
+    <td><a href="security.md">Security</a></td>
+    <td><a href="operation.md">Operation</a></td>
+    <td><a href="gateway-api.md">Gateway API</a></td>
+    <td><a href="resource-groups-api.md">Resource groups API</a></td>
+    <td><a href="routing-rules.md">Routing rules</a></td>
+    <td><a href="references.md">References</a></td>
+    <td><a href="release-notes.md">Release notes</a></td>
+  </tr>
+</table>
+
+# Configuration
+
+The Trino Gateway is configured by passing a yaml when running the start command.
+```shell
+java -XX:MinRAMPercentage=50 -XX:MaxRAMPercentage=80 --add-opens=java.base/java.lang=ALL-UNNAMED --add-opens=java.base/java.net=ALL-UNNAMED -jar gateway-ha.jar server gateway-config.yml
+```
+Each component of the Trino Gateway will have a corresponding node in the configuration yaml. 
+
+## Proxying additional paths
+
+By default, Trino Gateway only proxies requests to paths starting with 
+`/v1/statement`, `/v1/query`, `/ui`, `/v1/info`, `/v1/node`, 
+`/ui/api/stats` and `/oauth`. 
+
+If you want to proxy additional paths, 
+you can add them by adding the `extraWhitelistPaths` node to your gateway 
+configuration yaml:
+
+```yaml
+extraWhitelistPaths:
+  - "/ui/insights"
+  - "/api/v1/biac"
+  - "/api/v1/dataProduct"
+  - "/api/v1/dataproduct"
+  - "/ext/faster"
+```
+
+This example enables additional proxying of any requests to path starting with the specified paths. 
diff --git a/docs/design.md b/docs/design.md
@@ -2,6 +2,7 @@
 
 <table>
   <tr>
+    <td><a href="configuration.md">Configuration</a></td>
     <td><b><a href="design.md">Design</a></b></td>
     <td><a href="development.md">Development</a></td>
     <td><a href="security.md">Security</a></td>

diff --git a/docs/development.md b/docs/development.md
@@ -2,6 +2,7 @@
 
 <table>
   <tr>
+    <td><a href="configuration.md">Configuration</a></td>
     <td><a href="design.md">Design</a></td>
     <td><b><a href="development.md">Development</a></b></td>
     <td><a href="security.md">Security</a></td>

diff --git a/docs/gateway-api.md b/docs/gateway-api.md
@@ -2,6 +2,7 @@
 
 <table>
   <tr>
+    <td><a href="configuration.md">Configuration</a></td>
     <td><b><a href="design.md">Design</a></b></td>
     <td><a href="development.md">Development</a></td>
     <td><a href="security.md">Security</a></td>

diff --git a/docs/operation.md b/docs/operation.md
@@ -2,6 +2,7 @@
 
 <table>
   <tr>
+    <td><a href="configuration.md">Configuration</a></td>
     <td><b><a href="design.md">Design</a></b></td>
     <td><a href="development.md">Development</a></td>
     <td><a href="security.md">Security</a></td>

diff --git a/docs/references.md b/docs/references.md
@@ -2,6 +2,7 @@
 
 <table>
   <tr>
+    <td><a href="configuration.md">Configuration</a></td>
     <td><b><a href="design.md">Design</a></b></td>
     <td><a href="development.md">Development</a></td>
     <td><a href="security.md">Security</a></td>

diff --git a/docs/release-notes.md b/docs/release-notes.md
@@ -2,6 +2,7 @@
 
 <table>
   <tr>
+    <td><a href="configuration.md">Configuration</a></td>
     <td><a href="design.md">Design</a></td>
     <td><a href="development.md">Development</a></b></td>
     <td><a href="security.md">Security</a></td>

diff --git a/docs/resource-groups-api.md b/docs/resource-groups-api.md
@@ -2,6 +2,7 @@
 
 <table>
   <tr>
+    <td><a href="configuration.md">Configuration</a></td>
     <td><b><a href="design.md">Design</a></b></td>
     <td><a href="development.md">Development</a></td>
     <td><a href="security.md">Security</a></td>

diff --git a/docs/routing-rules.md b/docs/routing-rules.md
@@ -2,6 +2,7 @@
 
 <table>
   <tr>
+    <td><a href="configuration.md">Configuration</a></td>
     <td><b><a href="design.md">Design</a></b></td>
     <td><a href="development.md">Development</a></td>
     <td><a href="security.md">Security</a></td>

diff --git a/docs/security.md b/docs/security.md
@@ -2,6 +2,7 @@
 
 <table>
   <tr>
+    <td><a href="configuration.md">Configuration</a></td>
     <td><b><a href="design.md">Design</a></b></td>
     <td><a href="development.md">Development</a></td>
     <td><b><a href="security.md">Security</a></b></td>

diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/config/HaGatewayConfiguration.java b/gateway-ha/src/main/java/io/trino/gateway/ha/config/HaGatewayConfiguration.java
@@ -1,7 +1,10 @@
 package io.trino.gateway.ha.config;
 
 import io.trino.gateway.baseapp.AppConfiguration;
+
+import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import lombok.Data;
 import lombok.EqualsAndHashCode;
@@ -19,4 +22,5 @@ public class HaGatewayConfiguration extends AppConfiguration {
   private Map<String, UserConfiguration> presetUsers = new HashMap();
   private BackendStateConfiguration backendState;
   private ClusterStatsConfiguration clusterStatsConfiguration;
+  private List<String> extraWhitelistPaths = new ArrayList<>();
 }
diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/handler/QueryIdCachingProxyHandler.java b/gateway-ha/src/main/java/io/trino/gateway/ha/handler/QueryIdCachingProxyHandler.java
@@ -16,7 +16,6 @@
 import java.util.Base64;
 import java.util.HashMap;
 import java.util.List;
-import java.util.Optional;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 import javax.servlet.http.HttpServletRequest;
@@ -56,18 +55,20 @@ public class QueryIdCachingProxyHandler extends ProxyHandler {
 
   private final Meter requestMeter;
   private final int serverApplicationPort;
+  private final List<String> extraWhitelistPaths;
 
   public QueryIdCachingProxyHandler(
       QueryHistoryManager queryHistoryManager,
       RoutingManager routingManager,
       RoutingGroupSelector routingGroupSelector,
       int serverApplicationPort,
-      Meter requestMeter) {
+      Meter requestMeter, List<String> extraWhitelistPaths) {
     this.requestMeter = requestMeter;
     this.routingManager = routingManager;
     this.routingGroupSelector = routingGroupSelector;
     this.queryHistoryManager = queryHistoryManager;
     this.serverApplicationPort = serverApplicationPort;
+    this.extraWhitelistPaths = extraWhitelistPaths;
   }
 
   protected static String extractQueryIdIfPresent(String path, String queryParams) {
@@ -218,7 +219,8 @@ private boolean isPathWhiteListed(String path) {
         || path.startsWith(V1_INFO_PATH)
         || path.startsWith(V1_NODE_PATH)
         || path.startsWith(UI_API_STATS_PATH)
-        || path.startsWith(OAUTH_PATH);
+        || path.startsWith(OAUTH_PATH)
+        || extraWhitelistPaths.stream().anyMatch(s -> path.startsWith(s));
   }
 
   public boolean isAuthEnabled() {

diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/module/HaGatewayProviderModule.java b/gateway-ha/src/main/java/io/trino/gateway/ha/module/HaGatewayProviderModule.java
@@ -59,6 +59,7 @@ public class HaGatewayProviderModule extends AppModule<HaGatewayConfiguration, E
   private final AuthorizationManager authorizationManager;
   private final BackendStateManager backendStateConnectionManager;
   private final AuthFilter authenticationFilter;
+  private final List<String> extraWhitelistPaths;
 
   public HaGatewayProviderModule(HaGatewayConfiguration configuration, Environment environment) {
     super(configuration, environment);
@@ -79,6 +80,7 @@ public HaGatewayProviderModule(HaGatewayConfiguration configuration, Environment
         presetUsers);
     authenticationFilter = getAuthFilter(configuration);
     backendStateConnectionManager = new BackendStateManager(configuration.getBackendState());
+    extraWhitelistPaths = configuration.getExtraWhitelistPaths();
   }
 
   private LbOAuthManager getOAuthManager(HaGatewayConfiguration configuration) {
@@ -156,7 +158,8 @@ protected ProxyHandler getProxyHandler() {
         getRoutingManager(),
         routingGroupSelector,
         getApplicationPort(),
-        requestMeter);
+        requestMeter,
+        extraWhitelistPaths);
   }
 
   protected AuthFilter getAuthFilter(HaGatewayConfiguration configuration) {

diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/TestGatewayHaMultipleBackend.java b/gateway-ha/src/test/java/io/trino/gateway/ha/TestGatewayHaMultipleBackend.java
@@ -15,20 +15,27 @@
 public class TestGatewayHaMultipleBackend {
   public static final String EXPECTED_RESPONSE1 = "{\"id\":\"testId1\"}";
   public static final String EXPECTED_RESPONSE2 = "{\"id\":\"testId2\"}";
+  public static final String CUSTOM_RESPONSE = "123";
+  public static final String CUSTOM_PATH = "/v1/custom/extra";
 
   final int routerPort = 20000 + (int) (Math.random() * 1000);
   final int backend1Port = 21000 + (int) (Math.random() * 1000);
   final int backend2Port = 21000 + (int) (Math.random() * 1000);
+  final int customBackendPort = 21000 + (int) (Math.random() * 1000);
 
   private final WireMockServer adhocBackend =
       new WireMockServer(WireMockConfiguration.options().port(backend1Port));
   private final WireMockServer scheduledBackend =
       new WireMockServer(WireMockConfiguration.options().port(backend2Port));
 
+  private final WireMockServer customBackend =
+          new WireMockServer(WireMockConfiguration.options().port(customBackendPort));
+
   @BeforeClass(alwaysRun = true)
   public void setup() throws Exception {
     HaGatewayTestUtils.prepareMockBackend(adhocBackend, "/v1/statement", EXPECTED_RESPONSE1);
     HaGatewayTestUtils.prepareMockBackend(scheduledBackend, "/v1/statement", EXPECTED_RESPONSE2);
+    HaGatewayTestUtils.prepareMockBackend(customBackend, CUSTOM_PATH, CUSTOM_RESPONSE);
 
     // seed database
     HaGatewayTestUtils.TestConfig testConfig =
@@ -43,6 +50,34 @@ public void setup() throws Exception {
     HaGatewayTestUtils.setUpBackend(
         "trino2", "http://localhost:" + backend2Port, "externalUrl", true, "scheduled",
         routerPort);
+    HaGatewayTestUtils.setUpBackend(
+            "custom", "http://localhost:" + customBackendPort, "externalUrl", true, "custom",
+            routerPort);
+
+  }
+
+  @Test
+  public void testCustomPath() throws Exception {
+    OkHttpClient httpClient = new OkHttpClient();
+    RequestBody requestBody =
+            RequestBody.create(MediaType.parse("application/json; charset=utf-8"), "abc");
+    Request request1 =
+            new Request.Builder()
+                    .url("http://localhost:" + routerPort + CUSTOM_PATH)
+                    .post(requestBody)
+                    .addHeader("X-Trino-Routing-Group", "custom")
+                    .build();
+    Response response1 = httpClient.newCall(request1).execute();
+    Assert.assertEquals(response1.body().string(), CUSTOM_RESPONSE);
+
+    Request request2 =
+            new Request.Builder()
+                    .url("http://localhost:" + routerPort + "/invalid")
+                    .post(requestBody)
+                    .addHeader("X-Trino-Routing-Group", "custom")
+                    .build();
+    Response response2 = httpClient.newCall(request2).execute();
+    Assert.assertEquals(response2.code(), 404);
   }
 
   @Test

diff --git a/gateway-ha/src/test/resources/test-config-template.yml b/gateway-ha/src/test/resources/test-config-template.yml
@@ -22,3 +22,6 @@ modules:
 
 managedApps:
   - io.trino.gateway.ha.GatewayManagedApp
+
+extraWhitelistPaths:
+  - "/v1/custom"