Add S3 request rewriting capabilities

trino-aws-proxy-spi/src/main/java/io/trino/aws/proxy/spi/plugin/TrinoAwsProxyServerBinding.java

@@ -24,7 +24,9 @@
 import io.trino.aws.proxy.spi.plugin.config.AssumedRoleProviderConfig;
 import io.trino.aws.proxy.spi.plugin.config.CredentialsProviderConfig;
 import io.trino.aws.proxy.spi.plugin.config.PluginIdentifierConfig;
+import io.trino.aws.proxy.spi.plugin.config.S3RequestRewriterConfig;
 import io.trino.aws.proxy.spi.plugin.config.S3SecurityFacadeProviderConfig;
+import io.trino.aws.proxy.spi.rest.S3RequestRewriter;
 import io.trino.aws.proxy.spi.security.S3SecurityFacadeProvider;
 
 import static com.google.inject.multibindings.OptionalBinder.newOptionalBinder;
@@ -50,6 +52,11 @@ static Module s3SecurityFacadeProviderModule(String identifier, Class<? extends
  return optionalPluginModule(S3SecurityFacadeProviderConfig.class, identifier, S3SecurityFacadeProvider.class, implementationClass, module);
  }
 
+ static Module s3RequestRewriterModule(String identifier, Class<? extends S3RequestRewriter> implementationClass, Module module)
+ {
+ return optionalPluginModule(S3RequestRewriterConfig.class, identifier, S3RequestRewriter.class, implementationClass, module);
+ }
+
  static <T extends Identity> void bindIdentityType(Binder binder, Class<T> type)
  {
  newOptionalBinder(binder, new TypeLiteral<Class<? extends Identity>>() {}).setBinding().toProvider(() -> {

trino-aws-proxy-spi/src/main/java/io/trino/aws/proxy/spi/plugin/config/S3RequestRewriterConfig.java

@@ -0,0 +1,36 @@
+/*
+ * 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.aws.proxy.spi.plugin.config;
+
+import io.airlift.configuration.Config;
+
+import java.util.Optional;
+
+public class S3RequestRewriterConfig
+ implements PluginIdentifierConfig
+{
+ private Optional<String> identifier = Optional.empty();
+
+ @Override
+ public Optional<String> getPluginIdentifier()
+ {
+ return identifier;
+ }
+
+ @Config("s3-request-rewriter.type")
+ public void setPluginIdentifier(String identifier)
+ {
+ this.identifier = Optional.ofNullable(identifier);
+ }
+}

trino-aws-proxy-spi/src/main/java/io/trino/aws/proxy/spi/rest/S3RequestRewriter.java

@@ -0,0 +1,35 @@
+/*
+ * 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.aws.proxy.spi.rest;
+
+import io.trino.aws.proxy.spi.credentials.Credentials;
+
+import java.util.Optional;
+
+import static java.util.Objects.requireNonNull;
+
+public interface S3RequestRewriter
+{
+ S3RequestRewriter NOOP = (_, _) -> Optional.empty();
+
+ record S3RewriteResult(String finalRequestBucket, String finalRequestKey)
+ {
+ public S3RewriteResult {
+ requireNonNull(finalRequestBucket, "finalRequestBucket is null");
+ requireNonNull(finalRequestKey, "finalRequestKey is null");
+ }
+ }
+
+ Optional<S3RewriteResult> rewrite(Credentials credentials, ParsedS3Request request);
+}

trino-aws-proxy/src/main/java/io/trino/aws/proxy/server/TrinoAwsProxyServerModule.java

@@ -53,7 +53,9 @@
 import io.trino.aws.proxy.spi.plugin.TrinoAwsProxyServerPlugin;
 import io.trino.aws.proxy.spi.plugin.config.AssumedRoleProviderConfig;
 import io.trino.aws.proxy.spi.plugin.config.CredentialsProviderConfig;
+import io.trino.aws.proxy.spi.plugin.config.S3RequestRewriterConfig;
 import io.trino.aws.proxy.spi.plugin.config.S3SecurityFacadeProviderConfig;
+import io.trino.aws.proxy.spi.rest.S3RequestRewriter;
 import io.trino.aws.proxy.spi.security.S3SecurityFacadeProvider;
 import io.trino.aws.proxy.spi.signing.SigningServiceType;
 import org.glassfish.jersey.server.model.Resource;
@@ -121,6 +123,13 @@ protected void setup(Binder binder)
  });
  newSetBinder(binder, com.fasterxml.jackson.databind.Module.class).addBinding().toProvider(JsonIdentityProvider.class).in(Scopes.SINGLETON);
 
+ // RequestRewriter binder
+ configBinder(binder).bindConfig(S3RequestRewriterConfig.class);
+ newOptionalBinder(binder, S3RequestRewriter.class).setDefault().toProvider(() -> {
+ log.info("Using default %s NOOP implementation", S3RequestRewriter.class.getSimpleName());
+ return S3RequestRewriter.NOOP;
+ });
+
  // provided implementations
  install(new FileBasedCredentialsModule());
  install(new OpaS3SecurityModule());

trino-aws-proxy/src/main/java/io/trino/aws/proxy/server/rest/TrinoS3ProxyClient.java

@@ -25,6 +25,8 @@
 import io.trino.aws.proxy.spi.credentials.Credentials;
 import io.trino.aws.proxy.spi.rest.ParsedS3Request;
 import io.trino.aws.proxy.spi.rest.RequestContent;
+import io.trino.aws.proxy.spi.rest.S3RequestRewriter;
+import io.trino.aws.proxy.spi.rest.S3RequestRewriter.S3RewriteResult;
 import io.trino.aws.proxy.spi.security.SecurityResponse;
 import io.trino.aws.proxy.spi.security.SecurityResponse.Failure;
 import io.trino.aws.proxy.spi.signing.SigningContext;
@@ -38,6 +40,7 @@
 import jakarta.ws.rs.container.AsyncResponse;
 import jakarta.ws.rs.core.Response;
 import jakarta.ws.rs.core.UriBuilder;
+import software.amazon.awssdk.utils.http.SdkHttpUtils;
 
 import java.io.InputStream;
 import java.lang.annotation.Retention;
@@ -68,6 +71,7 @@ public class TrinoS3ProxyClient
  private final S3SecurityController s3SecurityController;
  private final S3PresignController s3PresignController;
  private final LimitStreamController limitStreamController;
+ private final S3RequestRewriter s3RequestRewriter;
  private final ExecutorService executorService = Executors.newVirtualThreadPerTaskExecutor();
  private final boolean generatePresignedUrlsOnHead;
 
@@ -84,14 +88,16 @@ public TrinoS3ProxyClient(
  S3SecurityController s3SecurityController,
  TrinoAwsProxyConfig trinoAwsProxyConfig,
  S3PresignController s3PresignController,
- LimitStreamController limitStreamController)
+ LimitStreamController limitStreamController,
+ S3RequestRewriter s3RequestRewriter)
  {
  this.httpClient = requireNonNull(httpClient, "httpClient is null");
  this.signingController = requireNonNull(signingController, "signingController is null");
  this.remoteS3Facade = requireNonNull(remoteS3Facade, "objectStore is null");
  this.s3SecurityController = requireNonNull(s3SecurityController, "securityController is null");
  this.s3PresignController = requireNonNull(s3PresignController, "presignController is null");
  this.limitStreamController = requireNonNull(limitStreamController, "quotaStreamController is null");
+ this.s3RequestRewriter = requireNonNull(s3RequestRewriter, "s3RequestRewriter is null");
 
  generatePresignedUrlsOnHead = trinoAwsProxyConfig.isGeneratePresignedUrlsOnHead();
  }
@@ -106,7 +112,13 @@ public void shutDown()
 
  public void proxyRequest(SigningMetadata signingMetadata, ParsedS3Request request, AsyncResponse asyncResponse, RequestLoggingSession requestLoggingSession)
  {
- URI remoteUri = remoteS3Facade.buildEndpoint(uriBuilder(request.queryParameters()), request.rawPath(), request.bucketName(), request.requestAuthorization().region());
+ Optional<S3RewriteResult> rewriteResult = s3RequestRewriter.rewrite(signingMetadata.credentials(), request);
+ String targetBucket = rewriteResult.map(S3RewriteResult::finalRequestBucket).orElse(request.bucketName());
+ String targetKey = rewriteResult
+ .map(S3RewriteResult::finalRequestKey)
+ .map(SdkHttpUtils::urlEncodeIgnoreSlashes)
+ .orElse(request.rawPath());
+ URI remoteUri = remoteS3Facade.buildEndpoint(uriBuilder(request.queryParameters()), targetKey, targetBucket, request.requestAuthorization().region());
 
  SecurityResponse securityResponse = s3SecurityController.apply(request, signingMetadata.credentials().identity());
  if (securityResponse instanceof Failure(var error)) {