Handle trailer headers and unsigned aws chunked payloads

Author: thinaih

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

@@ -13,7 +13,10 @@
  */
 package io.trino.aws.proxy.spi.rest;
 
+import com.google.common.collect.ImmutableList;
+
 import java.io.InputStream;
+import java.util.List;
 import java.util.Optional;
 
 @FunctionalInterface
@@ -27,7 +30,9 @@ enum ContentType
         STANDARD,
         W3C_CHUNKED,
         AWS_CHUNKED,
+        AWS_CHUNKED_UNSIGNED,
         AWS_CHUNKED_IN_W3C_CHUNKED,
+        AWS_CHUNKED_IN_W3C_CHUNKED_UNSIGNED,
     }
 
     default ContentType contentType()
@@ -48,5 +53,10 @@ default Optional<Integer> contentLength()
         return Optional.empty();
     }
 
+    default List<String> trailerHeaders()
+    {
+        return ImmutableList.of();
+    }
+
     Optional<InputStream> inputStream();
 }

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

@@ -14,11 +14,13 @@
 package io.trino.aws.proxy.server.rest;
 
 import com.google.common.base.Splitter;
+import com.google.common.collect.ImmutableList;
 import io.trino.aws.proxy.spi.signing.ChunkSigningSession;
 import jakarta.ws.rs.WebApplicationException;
 
 import java.io.IOException;
 import java.io.InputStream;
+import java.nio.charset.StandardCharsets;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -30,7 +32,7 @@ class AwsChunkedInputStream
         extends InputStream
 {
     private final InputStream delegate;
-    private final ChunkSigningSession chunkSigningSession;
+    private final Optional<ChunkSigningSession> chunkSigningSession;
 
     private enum State
     {
@@ -44,12 +46,14 @@ private enum State
     private int bytesRemainingInChunk;
     private int bytesAccountedFor;
     private final int decodedContentLength;
+    private final List<String> trailerHeaders;
 
-    AwsChunkedInputStream(InputStream delegate, ChunkSigningSession chunkSigningSession, int decodedContentLength)
+    AwsChunkedInputStream(InputStream delegate, Optional<ChunkSigningSession> chunkSigningSession, int decodedContentLength, List<String> trailerHeaders)
     {
         this.delegate = requireNonNull(delegate, "delegate is null");
         this.chunkSigningSession = requireNonNull(chunkSigningSession, "chunkSigningSession is null");
         this.decodedContentLength = decodedContentLength;
+        this.trailerHeaders = requireNonNull(ImmutableList.copyOf(trailerHeaders), "trailerHeaders is null");
     }
 
     @Override
@@ -65,7 +69,7 @@ public int read()
             throw new WebApplicationException("Unexpected end of stream", BAD_REQUEST);
         }
 
-        chunkSigningSession.write((byte) (i & 0xff));
+        chunkSigningSession.ifPresent(chunkSigningSession -> chunkSigningSession.write((byte) (i & 0xff)));
         updateBytesRemaining(1);
 
         return i;
@@ -86,7 +90,7 @@ public int read(byte[] b, int off, int len)
             throw new WebApplicationException("Unexpected end of stream", BAD_REQUEST);
         }
 
-        chunkSigningSession.write(b, off, count);
+        chunkSigningSession.ifPresent(chunkSigningSession -> chunkSigningSession.write(b, off, count));
         updateBytesRemaining(count);
 
         return count;
@@ -155,9 +159,6 @@ private void nextChunk()
         boolean success = false;
         do {
             List<String> parts = Splitter.on(';').trimResults().limit(2).splitToList(header);
-            if (parts.size() != 2) {
-                break;
-            }
 
             int chunkSize;
             try {
@@ -170,23 +171,41 @@ private void nextChunk()
                 break;
             }
 
-            Optional<String> chunkSignature = Splitter.on(';').trimResults().withKeyValueSeparator('=').split(parts.get(1))
-                    .entrySet()
-                    .stream()
-                    .filter(entry -> entry.getKey().equalsIgnoreCase("chunk-signature"))
-                    .map(Map.Entry::getValue)
-                    .findFirst();
+            if (chunkSigningSession.isPresent()) {
+                if (parts.size() != 2) {
+                    break;
+                }
+
+                Optional<String> chunkSignature = Splitter.on(';').trimResults().withKeyValueSeparator('=').split(parts.get(1))
+                        .entrySet()
+                        .stream()
+                        .filter(entry -> entry.getKey().equalsIgnoreCase("chunk-signature"))
+                        .map(Map.Entry::getValue)
+                        .findFirst();
 
-            if (chunkSignature.isEmpty()) {
-                break;
+                if (chunkSignature.isEmpty()) {
+                    break;
+                }
+
+                chunkSigningSession.get().startChunk(chunkSignature.get());
+            }
+            else {
+                if (parts.size() != 1) {
+                    break;
+                }
             }
 
-            chunkSigningSession.startChunk(chunkSignature.get());
             bytesRemainingInChunk = chunkSize;
 
             if (chunkSize == 0) {
-                readEmptyLine();
-                chunkSigningSession.complete();
+                if (trailerHeaders.isEmpty()) {
+                    readEmptyLine();
+                    chunkSigningSession.ifPresent(ChunkSigningSession::complete);
+                }
+                else {
+                    readTrailingHeaders();
+                    readEmptyLine();
+                }
                 state = State.LAST_CHUNK;
             }
             bytesAccountedFor += chunkSize;
@@ -236,4 +255,44 @@ private String readLine()
 
         return line.toString();
     }
+
+    private TrailerHeaderChunk readTrailingHeadersChunk()
+            throws IOException
+    {
+        Optional<String> signature = Optional.empty();
+        StringBuilder trailerHeadersChunkBuilder = new StringBuilder();
+        for (int i = 0; i < this.trailerHeaders.size(); i++) {
+            String trailerHeaders = readLine();
+            List<String> trailerHeadersValues = Splitter.on(":").trimResults().limit(2).splitToList(trailerHeaders);
+            String trailerHeaderName = trailerHeadersValues.getFirst();
+            if ((trailerHeadersValues.size() != 2) || !this.trailerHeaders.contains(trailerHeaderName)) {
+                throw new WebApplicationException("Trailer header is invalid: " + trailerHeaders, BAD_REQUEST);
+            }
+            if (trailerHeaderName.equals("x-amz-trailer-signature")) {
+                signature = Optional.of(trailerHeadersValues.getLast());
+                break;
+            }
+            else {
+                trailerHeadersChunkBuilder.append(trailerHeaders);
+            }
+        }
+        return new TrailerHeaderChunk(trailerHeadersChunkBuilder.toString(), signature);
+    }
+
+    private void readTrailingHeaders()
+            throws IOException
+    {
+        TrailerHeaderChunk trailerHeaderChunk = readTrailingHeadersChunk();
+        chunkSigningSession.ifPresent(chunkSigningSession -> {
+            if (trailerHeaderChunk.signature.isEmpty()) {
+                throw new WebApplicationException("Expected x-amz-trailer-signature, none found", BAD_REQUEST);
+            }
+            chunkSigningSession.startChunk(trailerHeaderChunk.signature.get());
+            byte[] trailerHeaderContent = trailerHeaderChunk.trailerHeaders.getBytes(StandardCharsets.UTF_8);
+            chunkSigningSession.write(trailerHeaderContent, 0, trailerHeaderContent.length);
+            chunkSigningSession.complete();
+        });
+    }
+
+    private record TrailerHeaderChunk(String trailerHeaders, Optional<String> signature) {}
 }

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

@@ -15,6 +15,7 @@
 
 import com.google.common.base.Splitter;
 import com.google.common.base.Suppliers;
+import com.google.common.collect.ImmutableList;
 import io.airlift.log.Logger;
 import io.trino.aws.proxy.server.rest.RequestHeadersBuilder.InternalRequestHeaders;
 import io.trino.aws.proxy.server.signing.SigningQueryParameters;
@@ -152,7 +153,7 @@ private static RequestContent buildRequestContent(InputStream requestEntityStrea
 
             // AWS does not mandate x-amz-decoded-content length is required for chunked transfer encoding
             // But we require it for simplicity (Content-Length is needed since we don't do chunking on outbound requests)
-            case AWS_CHUNKED, W3C_CHUNKED, AWS_CHUNKED_IN_W3C_CHUNKED -> () -> {
+            case AWS_CHUNKED, AWS_CHUNKED_UNSIGNED, W3C_CHUNKED, AWS_CHUNKED_IN_W3C_CHUNKED, AWS_CHUNKED_IN_W3C_CHUNKED_UNSIGNED -> () -> {
                 int contentLength = requestHeaders.decodedContentLength()
                         .orElseThrow(() -> new WebApplicationException(BAD_REQUEST));
                 return Optional.of(contentLength);
@@ -188,6 +189,12 @@ public Optional<InputStream> inputStream()
                         .map(bytes -> (InputStream) new ByteArrayInputStream(bytes))
                         .or(() -> Optional.of(requestEntityStream));
             }
+
+            @Override
+            public List<String> trailerHeaders()
+            {
+                return ImmutableList.copyOf(requestHeaders.requestHeaders().unmodifiedHeaders().get("x-amz-trailer"));
+            }
         };
     }
 }

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

@@ -50,7 +50,8 @@ private RequestHeadersBuilder() {}
             "connection",
             "amz-sdk-invocation-id",
             "amx-sdk-request",
-            "host");
+            "host",
+            "x-amz-trailer");
 
     record InternalRequestHeaders(
             RequestHeaders requestHeaders,
@@ -75,7 +76,9 @@ static InternalRequestHeaders parseHeaders(MultiMap allRequestHeaders)
         Builder builder = new Builder();
         allRequestHeaders.forEach((headerName, headerValues) -> {
             switch (headerName) {
-                case "authorization", "x-amz-security-token" -> {} // these get handled separately
+                case "authorization", "x-amz-security-token" -> {
+                    // these get handled separately
+                }
                 case "content-length" -> builder.contentLength(headerValues);
                 case "x-amz-decoded-content-length" -> builder.decodedContentLength(headerValues);
                 case "content-encoding" -> builder.contentEncoding(headerValues);
@@ -193,13 +196,22 @@ private void addPassthroughHeader(String headerName, List<String> headerValues)
             passthroughHeadersBuilder.addAll(headerName, headerValues);
         }
 
+        private String requiredContentSha256()
+        {
+            return contentSha256.orElseThrow(() -> new WebApplicationException(BAD_REQUEST));
+        }
+
         private void assertContentTypeValid(ContentType actualContentType)
         {
-            if (actualContentType == ContentType.AWS_CHUNKED || actualContentType == ContentType.AWS_CHUNKED_IN_W3C_CHUNKED || actualContentType == ContentType.W3C_CHUNKED) {
+            if (actualContentType == ContentType.AWS_CHUNKED
+                    || actualContentType == ContentType.AWS_CHUNKED_IN_W3C_CHUNKED
+                    || actualContentType == ContentType.W3C_CHUNKED
+                    || actualContentType == ContentType.AWS_CHUNKED_UNSIGNED
+                    || actualContentType == ContentType.AWS_CHUNKED_IN_W3C_CHUNKED_UNSIGNED) {
                 if (decodedContentLength.isEmpty()) {
                     throw new WebApplicationException(LENGTH_REQUIRED);
                 }
-                String sha256 = contentSha256.orElseThrow(() -> new WebApplicationException(BAD_REQUEST));
+                String sha256 = requiredContentSha256();
                 if (actualContentType != ContentType.W3C_CHUNKED && !sha256.startsWith("STREAMING-")) {
                     throw new WebApplicationException(BAD_REQUEST);
                 }
@@ -209,11 +221,21 @@ private void assertContentTypeValid(ContentType actualContentType)
         private InternalRequestHeaders build(MultiMap allHeaders)
         {
             Optional<ContentType> applicableContentType = switch (seenRequestPayloadContentTypes.size()) {
-                case 0, 1 -> seenRequestPayloadContentTypes.stream().findFirst();
+                case 0 -> Optional.empty();
+                case 1 -> {
+                    Optional<ContentType> contentType = seenRequestPayloadContentTypes.stream().findFirst();
+                    if (contentType.get().equals(ContentType.AWS_CHUNKED) && requiredContentSha256().startsWith("STREAMING-UNSIGNED-PAYLOAD")) {
+                        yield Optional.of(ContentType.AWS_CHUNKED_UNSIGNED);
+                    }
+                    yield contentType;
+                }
                 case 2 -> {
                     if (!seenRequestPayloadContentTypes.containsAll(ImmutableSet.of(ContentType.AWS_CHUNKED, ContentType.W3C_CHUNKED))) {
                         throw new WebApplicationException(BAD_REQUEST);
                     }
+                    if (requiredContentSha256().startsWith("STREAMING-UNSIGNED-PAYLOAD")) {
+                        yield Optional.of(ContentType.AWS_CHUNKED_IN_W3C_CHUNKED_UNSIGNED);
+                    }
                     yield Optional.of(ContentType.AWS_CHUNKED_IN_W3C_CHUNKED);
                 }
                 default -> throw new WebApplicationException(BAD_REQUEST);

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

@@ -212,8 +212,16 @@ private Optional<InputStream> contentInputStream(RequestContent requestContent,
     {
         return switch (requestContent.contentType()) {
             case AWS_CHUNKED, AWS_CHUNKED_IN_W3C_CHUNKED -> requestContent.inputStream()
-                    .map(inputStream -> new AwsChunkedInputStream(limitStreamController.wrap(inputStream), signingMetadata.requiredSigningContext().chunkSigningSession(), requestContent.contentLength().orElseThrow()));
-
+                    .map(inputStream -> new AwsChunkedInputStream(limitStreamController.wrap(inputStream),
+                            Optional.of(signingMetadata.requiredSigningContext().chunkSigningSession()),
+                            requestContent.contentLength().orElseThrow(),
+                            requestContent.trailerHeaders()));
+            case AWS_CHUNKED_UNSIGNED, AWS_CHUNKED_IN_W3C_CHUNKED_UNSIGNED -> requestContent.inputStream()
+                    .map(inputStream -> new AwsChunkedInputStream(
+                            limitStreamController.wrap(inputStream),
+                            Optional.empty(),
+                            requestContent.contentLength().orElseThrow(),
+                            requestContent.trailerHeaders()));
             case STANDARD, W3C_CHUNKED -> requestContent.inputStream().map(inputStream -> {
                 SigningContext signingContext = signingMetadata.requiredSigningContext();
                 return signingContext.contentHash()

trino-aws-proxy/src/test/java/io/trino/aws/proxy/server/TestGenericRestRequests.java

@@ -248,7 +248,7 @@ private StatusResponse doAwsChunkedUpload(
 
         URI requestUri = UriBuilder.fromUri(baseUri).path(bucket).path(key).build();
         RequestAuthorization requestAuthorization = signRequest(requestSigningCredential, requestUri, requestDate, "PUT", requestHeaderBuilder.build());
-        String chunkedContent = chunkedPayloadMutator.apply(TestingChunkSigningSession.build(chunkSigningCredential, requestAuthorization.signature(), requestDate).generateChunkedStream(contentToUpload, partitionCount));
+        String chunkedContent = chunkedPayloadMutator.apply(TestingChunkSigningSession.build(chunkSigningCredential, requestAuthorization.signature(), requestDate).generateChunkedStream(contentToUpload, partitionCount, Optional.empty()));
         Request.Builder requestBuilder = preparePut().setUri(requestUri);
 
         requestHeaderBuilder.add("Authorization", requestAuthorization.authorization());