Skip to content

Commit a75b9fe

Browse files
vishalyavishalya
committed
Add request filter for early user and query parsing
1 parent 4e7beb9 commit a75b9fe

23 files changed

+1261
-258
lines changed

‎gateway-ha/src/main/java/io/trino/gateway/baseapp/BaseApp.java‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,8 @@
3939
import io.trino.gateway.ha.router.RoutingRulesManager;
4040
import io.trino.gateway.ha.router.StochasticRoutingManager;
4141
import io.trino.gateway.ha.security.AuthorizedExceptionMapper;
42+
import io.trino.gateway.ha.security.QueryMetadataParser;
43+
import io.trino.gateway.ha.security.QueryUserInfoParser;
4244
import io.trino.gateway.proxyserver.ForProxy;
4345
import io.trino.gateway.proxyserver.ProxyRequestHandler;
4446
import io.trino.gateway.proxyserver.RouteToBackendResource;
@@ -187,6 +189,8 @@ private static void registerProxyResources(Binder binder)
187189
{
188190
jaxrsBinder(binder).bind(RouteToBackendResource.class);
189191
jaxrsBinder(binder).bind(RouterPreMatchContainerRequestFilter.class);
192+
jaxrsBinder(binder).bind(QueryUserInfoParser.class);
193+
jaxrsBinder(binder).bind(QueryMetadataParser.class);
190194
jaxrsBinder(binder).bind(ProxyRequestHandler.class);
191195
httpClientBinder(binder).bindHttpClient("proxy", ForProxy.class);
192196
httpClientBinder(binder).bindHttpClient("monitor", ForMonitor.class);

‎gateway-ha/src/main/java/io/trino/gateway/ha/handler/HttpUtils.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,8 @@ public class HttpUtils
2525
public static final String TRINO_UI_PATH = "/ui";
2626
public static final String OAUTH_PATH = "/oauth2";
2727
public static final String USER_HEADER = "X-Trino-User";
28+
public static final String TRINO_REQUEST_USER = "trinoRequestUser";
29+
public static final String TRINO_QUERY_PROPERTIES = "trinoQueryProperties";
2830

2931
private HttpUtils() {}
3032
}

‎gateway-ha/src/main/java/io/trino/gateway/ha/handler/ProxyUtils.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@
3030
import java.util.regex.Pattern;
3131

3232
import static com.google.common.base.Strings.isNullOrEmpty;
33+
import static io.trino.gateway.ha.handler.HttpUtils.TRINO_QUERY_PROPERTIES;
3334
import static io.trino.gateway.ha.handler.HttpUtils.TRINO_UI_PATH;
3435
import static io.trino.gateway.ha.handler.HttpUtils.V1_QUERY_PATH;
3536
import static java.nio.charset.StandardCharsets.UTF_8;
@@ -78,7 +79,7 @@ public static Optional<String> extractQueryIdIfPresent(
7879
throw new RuntimeException("Error reading request body", e);
7980
}
8081
if (!isNullOrEmpty(queryText) && queryText.toLowerCase(ENGLISH).contains("kill_query")) {
81-
TrinoQueryProperties trinoQueryProperties = new TrinoQueryProperties(request, requestAnalyserClientsUseV2Format, requestAnalyserMaxBodySize);
82+
TrinoQueryProperties trinoQueryProperties = (TrinoQueryProperties) request.getAttribute(TRINO_QUERY_PROPERTIES);
8283
return trinoQueryProperties.getQueryId();
8384
}
8485
return Optional.empty();

‎gateway-ha/src/main/java/io/trino/gateway/ha/handler/RoutingTargetHandler.java‎

Lines changed: 0 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -33,18 +33,10 @@
3333
import java.util.List;
3434
import java.util.Map;
3535
import java.util.Optional;
36-
import java.util.regex.Pattern;
3736
import java.util.stream.Stream;
3837

3938
import static com.google.common.base.Strings.isNullOrEmpty;
40-
import static com.google.common.collect.ImmutableList.toImmutableList;
41-
import static io.trino.gateway.ha.handler.HttpUtils.OAUTH_PATH;
42-
import static io.trino.gateway.ha.handler.HttpUtils.TRINO_UI_PATH;
43-
import static io.trino.gateway.ha.handler.HttpUtils.UI_API_STATS_PATH;
4439
import static io.trino.gateway.ha.handler.HttpUtils.USER_HEADER;
45-
import static io.trino.gateway.ha.handler.HttpUtils.V1_INFO_PATH;
46-
import static io.trino.gateway.ha.handler.HttpUtils.V1_NODE_PATH;
47-
import static io.trino.gateway.ha.handler.HttpUtils.V1_QUERY_PATH;
4840
import static io.trino.gateway.ha.handler.ProxyUtils.buildUriWithNewCluster;
4941
import static io.trino.gateway.ha.handler.ProxyUtils.extractQueryIdIfPresent;
5042
import static java.util.Objects.requireNonNull;
@@ -56,7 +48,6 @@ public class RoutingTargetHandler
5648
private final RoutingGroupSelector routingGroupSelector;
5749
private final String defaultRoutingGroup;
5850
private final List<String> statementPaths;
59-
private final List<Pattern> extraWhitelistPaths;
6051
private final boolean requestAnalyserClientsUseV2Format;
6152
private final int requestAnalyserMaxBodySize;
6253
private final boolean cookiesEnabled;
@@ -71,7 +62,6 @@ public RoutingTargetHandler(
7162
this.routingGroupSelector = requireNonNull(routingGroupSelector);
7263
this.defaultRoutingGroup = haGatewayConfiguration.getRouting().getDefaultRoutingGroup();
7364
statementPaths = requireNonNull(haGatewayConfiguration.getStatementPaths());
74-
extraWhitelistPaths = requireNonNull(haGatewayConfiguration.getExtraWhitelistPaths()).stream().map(Pattern::compile).collect(toImmutableList());
7565
requestAnalyserClientsUseV2Format = haGatewayConfiguration.getRequestAnalyzerConfig().isClientsUseV2Format();
7666
requestAnalyserMaxBodySize = haGatewayConfiguration.getRequestAnalyzerConfig().getMaxBodySize();
7767
cookiesEnabled = GatewayCookieConfigurationPropertiesProvider.getInstance().isEnabled();
@@ -118,18 +108,6 @@ private RoutingTargetResponse getRoutingTargetResponse(HttpServletRequest reques
118108
modifiedRequest);
119109
}
120110

121-
public boolean isPathWhiteListed(String path)
122-
{
123-
return statementPaths.stream().anyMatch(path::startsWith)
124-
|| path.startsWith(V1_QUERY_PATH)
125-
|| path.startsWith(TRINO_UI_PATH)
126-
|| path.startsWith(V1_INFO_PATH)
127-
|| path.startsWith(V1_NODE_PATH)
128-
|| path.startsWith(UI_API_STATS_PATH)
129-
|| path.startsWith(OAUTH_PATH)
130-
|| extraWhitelistPaths.stream().anyMatch(pattern -> pattern.matcher(path).matches());
131-
}
132-
133111
/**
134112
* A wrapper for HttpServletRequest that allows modifying multiple headers.
135113
*/

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@
4747
import io.trino.gateway.ha.router.HaGatewayManager;
4848
import io.trino.gateway.ha.router.HaQueryHistoryManager;
4949
import io.trino.gateway.ha.router.HaResourceGroupsManager;
50+
import io.trino.gateway.ha.router.PathFilter;
5051
import io.trino.gateway.ha.router.QueryHistoryManager;
5152
import io.trino.gateway.ha.router.ResourceGroupsManager;
5253
import io.trino.gateway.ha.router.RoutingGroupSelector;
@@ -88,6 +89,7 @@ public class HaGatewayProviderModule
8889
private final ResourceGroupsManager resourceGroupsManager;
8990
private final GatewayBackendManager gatewayBackendManager;
9091
private final QueryHistoryManager queryHistoryManager;
92+
private final PathFilter pathFilter;
9193

9294
@Override
9395
protected void configure()
@@ -97,11 +99,13 @@ protected void configure()
9799
binder().bind(GatewayBackendManager.class).toInstance(gatewayBackendManager);
98100
binder().bind(QueryHistoryManager.class).toInstance(queryHistoryManager);
99101
binder().bind(BackendStateManager.class).in(Scopes.SINGLETON);
102+
binder().bind(PathFilter.class).toInstance(pathFilter);
100103
}
101104

102105
public HaGatewayProviderModule(HaGatewayConfiguration configuration)
103106
{
104107
this.configuration = requireNonNull(configuration, "configuration is null");
108+
pathFilter = new PathFilter(configuration.getStatementPaths(), configuration.getExtraWhitelistPaths());
105109
Map<String, UserConfiguration> presetUsers = configuration.getPresetUsers();
106110

107111
oauthManager = getOAuthManager(configuration);

‎gateway-ha/src/main/java/io/trino/gateway/ha/router/ExternalRoutingGroupSelector.java‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,8 @@
4646
import static io.airlift.http.client.JsonResponseHandler.createJsonResponseHandler;
4747
import static io.airlift.http.client.Request.Builder.preparePost;
4848
import static io.airlift.json.JsonCodec.jsonCodec;
49+
import static io.trino.gateway.ha.handler.HttpUtils.TRINO_QUERY_PROPERTIES;
50+
import static io.trino.gateway.ha.handler.HttpUtils.TRINO_REQUEST_USER;
4951
import static java.util.Collections.list;
5052
import static java.util.Objects.requireNonNull;
5153

@@ -58,7 +60,6 @@ public class ExternalRoutingGroupSelector
5860
private final boolean propagateErrors;
5961
private final HttpClient httpClient;
6062
private final RequestAnalyzerConfig requestAnalyzerConfig;
61-
private final TrinoRequestUser.TrinoRequestUserProvider trinoRequestUserProvider;
6263
private static final JsonCodec<RoutingGroupExternalBody> ROUTING_GROUP_EXTERNAL_BODY_JSON_CODEC = jsonCodec(RoutingGroupExternalBody.class);
6364
private static final JsonResponseHandler<ExternalRouterResponse> ROUTING_GROUP_EXTERNAL_RESPONSE_JSON_RESPONSE_HANDLER =
6465
createJsonResponseHandler(jsonCodec(ExternalRouterResponse.class));
@@ -74,7 +75,6 @@ public class ExternalRoutingGroupSelector
7475
propagateErrors = rulesExternalConfiguration.isPropagateErrors();
7576

7677
this.requestAnalyzerConfig = requestAnalyzerConfig;
77-
trinoRequestUserProvider = new TrinoRequestUser.TrinoRequestUserProvider(requestAnalyzerConfig);
7878
try {
7979
this.uri = new URI(requireNonNull(rulesExternalConfiguration.getUrlPath(),
8080
"Invalid URL provided, using routing group header as default."));
@@ -143,8 +143,8 @@ private RoutingGroupExternalBody createRequestBody(HttpServletRequest request)
143143
TrinoQueryProperties trinoQueryProperties = null;
144144
TrinoRequestUser trinoRequestUser = null;
145145
if (requestAnalyzerConfig.isAnalyzeRequest()) {
146-
trinoQueryProperties = new TrinoQueryProperties(request, requestAnalyzerConfig.isClientsUseV2Format(), requestAnalyzerConfig.getMaxBodySize());
147-
trinoRequestUser = trinoRequestUserProvider.getInstance(request);
146+
trinoQueryProperties = (TrinoQueryProperties) request.getAttribute(TRINO_QUERY_PROPERTIES);
147+
trinoRequestUser = (TrinoRequestUser) request.getAttribute(TRINO_REQUEST_USER);
148148
}
149149

150150
return new RoutingGroupExternalBody(

‎gateway-ha/src/main/java/io/trino/gateway/ha/router/FileBasedRoutingGroupSelector.java‎

Lines changed: 5 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,8 @@
3333
import java.util.Map;
3434

3535
import static com.google.common.base.Suppliers.memoizeWithExpiration;
36+
import static io.trino.gateway.ha.handler.HttpUtils.TRINO_QUERY_PROPERTIES;
37+
import static io.trino.gateway.ha.handler.HttpUtils.TRINO_REQUEST_USER;
3638
import static java.nio.charset.StandardCharsets.UTF_8;
3739
import static java.util.Collections.sort;
3840

@@ -46,16 +48,10 @@ public class FileBasedRoutingGroupSelector
4648

4749
private final Supplier<List<RoutingRule>> rules;
4850
private final boolean analyzeRequest;
49-
private final boolean clientsUseV2Format;
50-
private final int maxBodySize;
51-
private final TrinoRequestUser.TrinoRequestUserProvider trinoRequestUserProvider;
5251

5352
public FileBasedRoutingGroupSelector(String rulesPath, Duration rulesRefreshPeriod, RequestAnalyzerConfig requestAnalyzerConfig)
5453
{
5554
analyzeRequest = requestAnalyzerConfig.isAnalyzeRequest();
56-
clientsUseV2Format = requestAnalyzerConfig.isClientsUseV2Format();
57-
maxBodySize = requestAnalyzerConfig.getMaxBodySize();
58-
trinoRequestUserProvider = new TrinoRequestUser.TrinoRequestUserProvider(requestAnalyzerConfig);
5955

6056
rules = memoizeWithExpiration(() -> readRulesFromPath(Path.of(rulesPath)), rulesRefreshPeriod.toJavaTime());
6157
}
@@ -68,12 +64,9 @@ public RoutingSelectorResponse findRoutingDestination(HttpServletRequest request
6864

6965
Map<String, Object> data;
7066
if (analyzeRequest) {
71-
TrinoQueryProperties trinoQueryProperties = new TrinoQueryProperties(
72-
request,
73-
clientsUseV2Format,
74-
maxBodySize);
75-
TrinoRequestUser trinoRequestUser = trinoRequestUserProvider.getInstance(request);
76-
data = ImmutableMap.of("request", request, "trinoQueryProperties", trinoQueryProperties, "trinoRequestUser", trinoRequestUser);
67+
TrinoQueryProperties trinoQueryProperties = (TrinoQueryProperties) request.getAttribute(TRINO_QUERY_PROPERTIES);
68+
TrinoRequestUser trinoRequestUser = (TrinoRequestUser) request.getAttribute(TRINO_REQUEST_USER);
69+
data = ImmutableMap.of("request", request, TRINO_QUERY_PROPERTIES, trinoQueryProperties, TRINO_REQUEST_USER, trinoRequestUser);
7770
}
7871
else {
7972
data = ImmutableMap.of("request", request);

‎gateway-ha/src/main/java/io/trino/gateway/ha/router/PathFilter.java‎

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
/*
2+
* Licensed under the Apache License, Version 2.0 (the "License");
3+
* you may not use this file except in compliance with the License.
4+
* You may obtain a copy of the License at
5+
*
6+
* http://www.apache.org/licenses/LICENSE-2.0
7+
*
8+
* Unless required by applicable law or agreed to in writing, software
9+
* distributed under the License is distributed on an "AS IS" BASIS,
10+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11+
* See the License for the specific language governing permissions and
12+
* limitations under the License.
13+
*/
14+
package io.trino.gateway.ha.router;
15+
16+
import com.google.inject.Inject;
17+
import com.google.inject.Singleton;
18+
19+
import java.util.List;
20+
import java.util.Set;
21+
import java.util.regex.Pattern;
22+
import java.util.regex.PatternSyntaxException;
23+
24+
import static com.google.common.collect.ImmutableList.toImmutableList;
25+
import static io.trino.gateway.ha.handler.HttpUtils.OAUTH_PATH;
26+
import static io.trino.gateway.ha.handler.HttpUtils.TRINO_UI_PATH;
27+
import static io.trino.gateway.ha.handler.HttpUtils.UI_API_STATS_PATH;
28+
import static io.trino.gateway.ha.handler.HttpUtils.V1_INFO_PATH;
29+
import static io.trino.gateway.ha.handler.HttpUtils.V1_NODE_PATH;
30+
import static io.trino.gateway.ha.handler.HttpUtils.V1_QUERY_PATH;
31+
import static java.util.Objects.requireNonNull;
32+
33+
/**
34+
* A filter component that determines whether a given path should be whitelisted
35+
* for routing to Trino clusters.
36+
*/
37+
@Singleton
38+
public class PathFilter
39+
{
40+
private final Set<String> statementPaths;
41+
private final List<Pattern> extraWhitelistPatterns;
42+
43+
@Inject
44+
public PathFilter(
45+
List<String> statementPaths,
46+
List<String> extraWhitelistPaths)
47+
{
48+
this.statementPaths = Set.copyOf(requireNonNull(statementPaths, "Required configuration 'statementPaths' can't be null"));
49+
this.extraWhitelistPatterns = requireNonNull(extraWhitelistPaths, "extraWhitelistPaths cannot be null").stream()
50+
.map(pattern -> {
51+
try {
52+
return Pattern.compile(pattern);
53+
}
54+
catch (PatternSyntaxException e) {
55+
throw new IllegalArgumentException("Invalid regex pattern: " + pattern, e);
56+
}
57+
})
58+
.collect(toImmutableList());
59+
}
60+
61+
/**
62+
* Determines if the given path is whitelisted for routing to backend.
63+
*
64+
* @param path the request path to check
65+
* @return true if the path should be routed to backend, false otherwise
66+
*/
67+
public boolean isPathWhiteListed(String path)
68+
{
69+
return statementPaths.stream().anyMatch(path::startsWith)
70+
|| path.startsWith(V1_QUERY_PATH)
71+
|| path.startsWith(TRINO_UI_PATH)
72+
|| path.startsWith(V1_INFO_PATH)
73+
|| path.startsWith(V1_NODE_PATH)
74+
|| path.startsWith(UI_API_STATS_PATH)
75+
|| path.startsWith(OAUTH_PATH)
76+
|| extraWhitelistPatterns.stream().anyMatch(pattern -> pattern.matcher(path).matches());
77+
}
78+
}

0 commit comments

Comments
 (0)