Add cookie based routing for selected paths

docs/design.md

@@ -18,7 +18,8 @@
         <li><a href="gateway-api.md">Gateway API</a></li>
         <li><a href="resource-groups-api.md">Resource groups API</a></li>
         <li><a href="routing-rules.md">Routing rules</a></li>
-      </ul>
+        <li><a href="routing-logic.md">Routing logic</a></li>
+</ul>
     </td>
     <td>
       <ul>

docs/development.md

@@ -18,6 +18,7 @@
         <li><a href="gateway-api.md">Gateway API</a></li>
         <li><a href="resource-groups-api.md">Resource groups API</a></li>
         <li><a href="routing-rules.md">Routing rules</a></li>
+        <li><a href="routing-logic.md">Routing logic</a></li>
       </ul>
     </td>
     <td>

docs/gateway-api.md

@@ -18,6 +18,7 @@
         <li><a href="gateway-api.md">Gateway API</a></li>
         <li><a href="resource-groups-api.md">Resource groups API</a></li>
         <li><a href="routing-rules.md">Routing rules</a></li>
+        <li><a href="routing-logic.md">Routing logic</a></li>
       </ul>
     </td>
     <td>

docs/installation.md

@@ -18,6 +18,7 @@
         <li><a href="gateway-api.md">Gateway API</a></li>
         <li><a href="resource-groups-api.md">Resource groups API</a></li>
         <li><a href="routing-rules.md">Routing rules</a></li>
+        <li><a href="routing-logic.md">Routing logic</a></li>
       </ul>
     </td>
     <td>

docs/operation.md

@@ -18,6 +18,7 @@
         <li><a href="gateway-api.md">Gateway API</a></li>
         <li><a href="resource-groups-api.md">Resource groups API</a></li>
         <li><a href="routing-rules.md">Routing rules</a></li>
+        <li><a href="routing-logic.md">Routing logic</a></li>
       </ul>
     </td>
     <td>

docs/quickstart.md

@@ -18,6 +18,7 @@
         <li><a href="gateway-api.md">Gateway API</a></li>
         <li><a href="resource-groups-api.md">Resource groups API</a></li>
         <li><a href="routing-rules.md">Routing rules</a></li>
+        <li><a href="routing-logic.md">Routing logic</a></li>
       </ul>
     </td>
     <td>

docs/references.md

@@ -18,6 +18,7 @@
         <li><a href="gateway-api.md">Gateway API</a></li>
         <li><a href="resource-groups-api.md">Resource groups API</a></li>
         <li><a href="routing-rules.md">Routing rules</a></li>
+        <li><a href="routing-logic.md">Routing logic</a></li>
       </ul>
     </td>
     <td>

docs/release-notes.md

@@ -18,6 +18,7 @@
         <li><a href="gateway-api.md">Gateway API</a></li>
         <li><a href="resource-groups-api.md">Resource groups API</a></li>
         <li><a href="routing-rules.md">Routing rules</a></li>
+        <li><a href="routing-logic.md">Routing logic</a></li>
       </ul>
     </td>
     <td>

docs/resource-groups-api.md

@@ -18,6 +18,7 @@
         <li><a href="gateway-api.md">Gateway API</a></li>
         <li><a href="resource-groups-api.md">Resource groups API</a></li>
         <li><a href="routing-rules.md">Routing rules</a></li>
+        <li><a href="routing-logic.md">Routing logic</a></li>
       </ul>
     </td>
     <td>

docs/routing-logic.md

@@ -0,0 +1,111 @@
+**Trino Gateway documentation**
+
+<table>
+  <tr>
+    <td>
+      <img src="./assets/logos/trino-gateway-v.png"/>
+    </td>
+    <td>
+      <ul>
+        <li><a href="quickstart.md">Quickstart</a></li>
+        <li><a href="installation.md">Installation</a></li>
+        <li><a href="security.md">Security</a></li>
+        <li><a href="operation.md">Operation</a></li>
+      </ul>
+    </td>
+    <td>
+      <ul>
+        <li><a href="gateway-api.md">Gateway API</a></li>
+        <li><a href="resource-groups-api.md">Resource groups API</a></li>
+        <li><a href="routing-rules.md">Routing rules</a></li>
+        <li><a href="routing-logic.md">Routing logic</a></li>
+      </ul>
+    </td>
+    <td>
+      <ul>
+        <li><a href="design.md">Design</a></li>
+        <li><a href="development.md">Development</a></li>
+        <li><a href="release-notes.md">Release notes</a></li>
+        <li><a href="references.md">References</a></li>
+      </ul>
+    </td>
+  </tr>
+</table>
+
+# Routing Logic
+
+## Overview
+
+Trino Gateway checks incoming requests to see if they're related to previous 
+ones it handled. If they are, then Trino Gateway sends them to the same backend 
+that dealt with the earlier requests.
+
+If it is a new request, the Trino Gateway refers to [Routing rules](routing-rules.md) 
+to decide which group of backends, called a 'Routing Group,' should handle it. 
+It then picks a backend from that Routing Group to handle the request using 
+either an adaptive or round-robin strategy.
+
+![Request Routing Flow](assets/gateway-routing-flow.svg)
+
+## Sticky routing
+
+A request related to an ongoing process, or to state maintained on a single 
+backend cluster, must be routed to that backend for proper handling. Two 
+mechanisms for identifying related requests are currently implemented. By default,
+only routing based on query identifier is enabled.
+
+### Routing based on query identifier (default)
+
+When a query is initiated through the Trino Gateway, the query id will be 
+extracted from the response and mapped to the backend that provided the 
+response. Any subsequent request containing that query id will be forwarded 
+to that backend. For example, to retrieve query results, the trino client 
+polls a URI of the form 
+`v1/statement/executing/queryid/nonce/counter`. The Trino Gateway will extract
+the queryid from this URI.
+
+### Routing based on cookies
+
+OAuth2 authentication requires that the same backend is used for each step of 
+the handshake. When `gatewayCookieConfiguration.enabled` is set to true, a cookie 
+will be added to requests made to paths beginning with `/oauth2` unless they already have 
+a cookie present, which is used to route further `/oauth2/*` requests to the correct backend. 
+Cookies are not added to requests to `v1/*` and other Trino endpoints.
+
+Trino Gateway signs its cookies to ensure that they are not tampered with. You 
+must set a `cookieSigningSecret` string in your configuration
+```yaml
+gatewayCookieConfiguration:
+    enabled: true
+    cookieSigningSecret: "ahighentropystring"
+```
+when making use of this feature. If you load balance request across multiple Trino Gateway
+instances, ensure each instance has the same `cookieSigningSecret`.
+
+The Trino Gateway handles standard Trino OAuth2 handshakes with no additional 
+configuration. If you are using a customized or commercial Trino distribution, then
+the paths used to define the OAuth handshake may be modified.
+
+`routingPaths`: If the request URI starts with a path in this list, then
+* If no cookie is present, add a routing cookie
+* If a cookie is present, route the request to the backend defined by that cookie
+
+`deletePaths`: If the request URI starts with a path in this list,
+return a response that instructs the client to delete the cookie.
+
+Additionally, the `lifetime` property sets the duration for which a cookie remains in 
+effect after creation. Ensure that it is greater than
+the time required to complete the handshake. Default `lifetime` is 10 minutes.
+
+These properties are defined under the `oauth2GatewayCookieConfiguration` node: 
+
+```yaml
+oauth2GatewayCookieConfiguration:
+  routingPaths:
+    - "/oauth2"
+    - "/custom/oauth2/callback"
+    - "/alternative/oauth2/initiate"
+  deletePaths:
+    - "/custom/logout"
+  lifetime: "5m"
+```

docs/routing-rules.md

@@ -18,6 +18,7 @@
         <li><a href="gateway-api.md">Gateway API</a></li>
         <li><a href="resource-groups-api.md">Resource groups API</a></li>
         <li><a href="routing-rules.md">Routing rules</a></li>
+        <li><a href="routing-logic.md">Routing logic</a></li>
       </ul>
     </td>
     <td>

docs/security.md

@@ -18,6 +18,7 @@
         <li><a href="gateway-api.md">Gateway API</a></li>
         <li><a href="resource-groups-api.md">Resource groups API</a></li>
         <li><a href="routing-rules.md">Routing rules</a></li>
+        <li><a href="routing-logic.md">Routing logic</a></li>
       </ul>
     </td>
     <td>

gateway-ha/pom.xml

@@ -126,6 +126,12 @@
             <artifactId>stats</artifactId>
         </dependency>
 
+        <dependency>
+            <groupId>io.airlift</groupId>
+            <artifactId>units</artifactId>
+            <version>1.10</version>
+        </dependency>
+
         <dependency>
             <groupId>io.dropwizard</groupId>
             <artifactId>dropwizard-assets</artifactId>

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

@@ -0,0 +1,45 @@
+/*
+ * Licensed 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 io.trino.gateway.ha.config;
+
+import javax.crypto.SecretKey;
+import javax.crypto.spec.SecretKeySpec;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+
+public class GatewayCookieConfiguration
+{
+    private SecretKey cookieSigningKey;
+    private boolean enabled;
+
+    public boolean isEnabled()
+    {
+        return enabled;
+    }
+
+    public void setEnabled(boolean enabled)
+    {
+        this.enabled = enabled;
+    }
+
+    public SecretKey getCookieSigningKey()
+    {
+        return cookieSigningKey;
+    }
+
+    public void setCookieSigningSecret(String cookieSigningSecret)
+    {
+        cookieSigningKey = new SecretKeySpec(cookieSigningSecret.getBytes(UTF_8), "HmacSHA256");
+    }
+}

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

@@ -0,0 +1,57 @@
+/*
+ * Licensed 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 io.trino.gateway.ha.config;
+
+import javax.crypto.SecretKey;
+
+public class GatewayCookieConfigurationPropertiesProvider
+{
+    private static final GatewayCookieConfigurationPropertiesProvider instance = new GatewayCookieConfigurationPropertiesProvider();
+    private GatewayCookieConfiguration gatewayCookieConfiguration;
+
+    private GatewayCookieConfigurationPropertiesProvider()
+    {}
+
+    public void initialize(GatewayCookieConfiguration gatewayCookieConfiguration)
+    {
+        if (gatewayCookieConfiguration.isEnabled() && gatewayCookieConfiguration.getCookieSigningKey() == null) {
+            throw new IllegalArgumentException("gatewayCookieConfiguration.cookieSigningSecret must be provided when cookies are enabled");
+        }
+        this.gatewayCookieConfiguration = gatewayCookieConfiguration;
+    }
+
+    public static GatewayCookieConfigurationPropertiesProvider getInstance()
+    {
+        return instance;
+    }
+
+    public boolean isEnabled()
+    {
+        ensureInitialized();
+        return gatewayCookieConfiguration.isEnabled();
+    }
+
+    public SecretKey getCookieSigningKey()
+    {
+        ensureInitialized();
+        return gatewayCookieConfiguration.getCookieSigningKey();
+    }
+
+    private void ensureInitialized()
+    {
+        if (gatewayCookieConfiguration == null) {
+            throw new IllegalStateException("getInstance.initialize(GatewayCookieConfiguration) must be called before use");
+        }
+    }
+}

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

@@ -35,6 +35,8 @@ public class HaGatewayConfiguration
     private BackendStateConfiguration backendState;
     private ClusterStatsConfiguration clusterStatsConfiguration;
     private List<String> extraWhitelistPaths = new ArrayList<>();
+    private OAuth2GatewayCookieConfiguration oauth2GatewayCookieConfiguration = new OAuth2GatewayCookieConfiguration();
+    private GatewayCookieConfiguration gatewayCookieConfiguration = new GatewayCookieConfiguration();
 
     // List of Modules with FQCN (Fully Qualified Class Name)
     private List<String> modules;
@@ -164,6 +166,26 @@ public void setExtraWhitelistPaths(List<String> extraWhitelistPaths)
         this.extraWhitelistPaths = extraWhitelistPaths;
     }
 
+    public OAuth2GatewayCookieConfiguration getOauth2GatewayCookieConfiguration()
+    {
+        return oauth2GatewayCookieConfiguration;
+    }
+
+    public void setOauth2GatewayCookieConfiguration(OAuth2GatewayCookieConfiguration oauth2GatewayCookieConfiguration)
+    {
+        this.oauth2GatewayCookieConfiguration = oauth2GatewayCookieConfiguration;
+    }
+
+    public GatewayCookieConfiguration getGatewayCookieConfiguration()
+    {
+        return gatewayCookieConfiguration;
+    }
+
+    public void setGatewayCookieConfiguration(GatewayCookieConfiguration gatewayCookieConfiguration)
+    {
+        this.gatewayCookieConfiguration = gatewayCookieConfiguration;
+    }
+
     public List<String> getModules()
     {
         return this.modules;