Handle streaming exceptions as bad requests instead of internal errors

thinaih · mosiac1 · commit 2b5f0306650d · 2025-04-01T10:29:30.000+01:00
diff --git a/trino-aws-proxy/src/main/java/io/trino/aws/proxy/server/rest/AwsChunkedInputStream.java b/trino-aws-proxy/src/main/java/io/trino/aws/proxy/server/rest/AwsChunkedInputStream.java
@@ -15,14 +15,15 @@
 
 import com.google.common.base.Splitter;
 import io.trino.aws.proxy.spi.signing.ChunkSigningSession;
+import jakarta.ws.rs.WebApplicationException;
 
-import java.io.EOFException;
 import java.io.IOException;
 import java.io.InputStream;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 
+import static jakarta.ws.rs.core.Response.Status.BAD_REQUEST;
 import static java.util.Objects.requireNonNull;
 
 class AwsChunkedInputStream
@@ -61,7 +62,7 @@ public int read()
 
         int i = delegate.read();
         if (i < 0) {
-            throw new EOFException("Unexpected end of stream");
+            throw new WebApplicationException("Unexpected end of stream", BAD_REQUEST);
         }
 
         chunkSigningSession.write((byte) (i & 0xff));
@@ -82,7 +83,7 @@ public int read(byte[] b, int off, int len)
 
         int count = delegate.read(b, off, len);
         if (count < 0) {
-            throw new EOFException("Unexpected end of stream");
+            throw new WebApplicationException("Unexpected end of stream", BAD_REQUEST);
         }
 
         chunkSigningSession.write(b, off, count);
@@ -116,7 +117,7 @@ private void updateBytesRemaining(int count)
             nextChunk();
         }
         else if (bytesRemainingInChunk < 0) {
-            throw new IllegalStateException("bytesRemainingInChunk has gone negative: " + bytesRemainingInChunk);
+            throw new WebApplicationException("bytesRemainingInChunk has gone negative: " + bytesRemainingInChunk, BAD_REQUEST);
         }
     }
 
@@ -195,10 +196,11 @@ private void nextChunk()
         while (false);
 
         if (!success) {
-            throw new IOException("Invalid chunk header: " + header);
+            throw new WebApplicationException("Invalid chunk header: " + header, BAD_REQUEST);
         }
         if (bytesAccountedFor > decodedContentLength) {
-            throw new IllegalStateException("chunked data headers report a larger size than originally declared in the request: declared %s sent %s".formatted(decodedContentLength, bytesAccountedFor));
+            throw new WebApplicationException("chunked data headers report a larger size than originally declared in the request: declared %s sent %s".formatted(decodedContentLength, bytesAccountedFor),
+                    BAD_REQUEST);
         }
     }
 
@@ -207,7 +209,7 @@ private void readEmptyLine()
     {
         String crLf = readLine();
         if (!crLf.isEmpty()) {
-            throw new IOException("Expected CR/LF. Instead read: " + crLf);
+            throw new WebApplicationException("Expected CR/LF. Instead read: " + crLf, BAD_REQUEST);
         }
     }
 
@@ -219,7 +221,7 @@ private String readLine()
             int i = delegate.read();
             if (i < 0) {
                 delegateIsDone = true;
-                throw new EOFException("Unexpected end of stream");
+                throw new WebApplicationException("Unexpected end of stream", BAD_REQUEST);
             }
             if (i == '\r') {
                 break;
@@ -229,7 +231,7 @@ private String readLine()
 
         int i = delegate.read();
         if (i != '\n') {
-            throw new IOException("Expected LF. Instead read: " + i);
+            throw new WebApplicationException("Expected LF. Instead read: " + i, BAD_REQUEST);
         }
 
         return line.toString();
diff --git a/trino-aws-proxy/src/main/java/io/trino/aws/proxy/server/rest/HashCheckInputStream.java b/trino-aws-proxy/src/main/java/io/trino/aws/proxy/server/rest/HashCheckInputStream.java
@@ -22,6 +22,7 @@
 import java.io.InputStream;
 import java.util.Optional;
 
+import static jakarta.ws.rs.core.Response.Status.BAD_REQUEST;
 import static jakarta.ws.rs.core.Response.Status.UNAUTHORIZED;
 import static java.util.Objects.requireNonNull;
 
@@ -106,8 +107,7 @@ private void updateBytesRead(int count)
         bytesRead += count;
         expectedLength.ifPresent(expected -> {
             if (bytesRead > expected) {
-                log.debug("More bytes read than expected. Expected: %s, Actual: %s", expected, bytesRead);
-                throw new WebApplicationException(UNAUTHORIZED);
+                throw new WebApplicationException("More bytes read than expected. Expected: %s, Actual: %s".formatted(expected, bytesRead), BAD_REQUEST);
             }
 
             if (bytesRead == expected) {
diff --git a/trino-aws-proxy/src/main/java/io/trino/aws/proxy/server/rest/StreamingResponseHandler.java b/trino-aws-proxy/src/main/java/io/trino/aws/proxy/server/rest/StreamingResponseHandler.java
@@ -101,8 +101,10 @@ private void resume(Object result)
     {
         switch (result) {
             case WebApplicationException exception -> resume(exception.getResponse());
-            case Throwable exception when Throwables.getRootCause(exception) instanceof WebApplicationException webApplicationException -> resume(webApplicationException.getResponse());
-            case Throwable exception -> resume(jakarta.ws.rs.core.Response.status(INTERNAL_SERVER_ERROR.getStatusCode(), Optional.ofNullable(exception.getMessage()).orElse("Unknown error")).build());
+            case Throwable exception when Throwables.getRootCause(exception) instanceof WebApplicationException webApplicationException ->
+                    resume(webApplicationException.getResponse());
+            case Throwable exception ->
+                    resume(jakarta.ws.rs.core.Response.status(INTERNAL_SERVER_ERROR.getStatusCode(), Optional.ofNullable(exception.getMessage()).orElse("Unknown error")).build());
             default -> {
                 if (hasBeenResumed.compareAndSet(false, true)) {
                     asyncResponse.resume(result);
diff --git a/trino-aws-proxy/src/test/java/io/trino/aws/proxy/server/TestGenericRestRequests.java b/trino-aws-proxy/src/test/java/io/trino/aws/proxy/server/TestGenericRestRequests.java
@@ -182,7 +182,7 @@ public void testAwsChunkedUploadInvalidContent()
 
         // Final chunk has an invalid size
         Function<String, String> changeSizeOfFinalChunk = chunked -> chunked.replaceFirst("\\r\\n0;chunk-signature=(\\w+)", "\r\n1;chunk-signature=$1");
-        assertThat(doAwsChunkedUpload(bucket, fileKey, LOREM_IPSUM, 2, validCredential, changeSizeOfFinalChunk).getStatusCode()).isEqualTo(500);
+        assertThat(doAwsChunkedUpload(bucket, fileKey, LOREM_IPSUM, 2, validCredential, changeSizeOfFinalChunk).getStatusCode()).isEqualTo(400);
         assertFileNotInS3(storageClient, bucket, fileKey);
 
         // First chunk has an invalid size
@@ -198,7 +198,7 @@ public void testAwsChunkedUploadInvalidContent()
             }
             return "%s%s".formatted(newSizeAsString, chunked.substring(firstChunkIdx));
         };
-        assertThat(doAwsChunkedUpload(bucket, fileKey, LOREM_IPSUM, 2, validCredential, changeSizeOfFirstChunk).getStatusCode()).isEqualTo(500);
+        assertThat(doAwsChunkedUpload(bucket, fileKey, LOREM_IPSUM, 2, validCredential, changeSizeOfFirstChunk).getStatusCode()).isEqualTo(400);
         assertFileNotInS3(storageClient, bucket, fileKey);
 
         // Change the signature of each of the chunks
@@ -269,24 +269,24 @@ public void testAwsChunkedCornerCases()
         storageClient.createBucket(r -> r.bucket(bucket).build());
 
         // Illegal signature and no final chunk
-        testAwsChunkedIllegalChunks(bucket, "no-final-chunk", buildFakeChunk(longDummyContent, longDummyContent.length()), longDummyContent.length(), 500);
+        testAwsChunkedIllegalChunks(bucket, "no-final-chunk", buildFakeChunk(longDummyContent, longDummyContent.length()), longDummyContent.length(), 400);
         // Illegal signature with a final chunk
         testAwsChunkedIllegalChunks(bucket, "with-final-chunk", "%s%s".formatted(buildFakeChunk(longDummyContent, longDummyContent.length()), buildFakeChunk("", 0)), longDummyContent.length(), 401);
         // Illegal signature and no final chunk - more chunked data than we report in the x-amz-decoded-content-length header
-        testAwsChunkedIllegalChunks(bucket, "no-final-chunk-more-data-than-headers-indicate", buildFakeChunk(longDummyContent, longDummyContent.length()), 4096, 500);
+        testAwsChunkedIllegalChunks(bucket, "no-final-chunk-more-data-than-headers-indicate", buildFakeChunk(longDummyContent, longDummyContent.length()), 4096, 400);
 
         // Illegal signature with a final chunk - more chunked data than we report in the x-amz-decoded-content-length header
-        testAwsChunkedIllegalChunks(bucket, "with-final-chunk-more-data-than-headers-indicate", "%s%s".formatted(buildFakeChunk(longDummyContent, longDummyContent.length()), buildFakeChunk("", 0)), 4096, 500);
+        testAwsChunkedIllegalChunks(bucket, "with-final-chunk-more-data-than-headers-indicate", "%s%s".formatted(buildFakeChunk(longDummyContent, longDummyContent.length()), buildFakeChunk("", 0)), 4096, 400);
 
         // Illegal signature and no final chunk - chunk misreports its size
-        testAwsChunkedIllegalChunks(bucket, "no-final-chunk-chunk-underreports-size", buildFakeChunk(longDummyContent, 4096), 4096, 500);
+        testAwsChunkedIllegalChunks(bucket, "no-final-chunk-chunk-underreports-size", buildFakeChunk(longDummyContent, 4096), 4096, 400);
         // Illegal signature with a final chunk - chunk misreports its size
-        testAwsChunkedIllegalChunks(bucket, "with-final-chunk-chunk-underreports-size", "%s%s".formatted(buildFakeChunk(longDummyContent, 4096), buildFakeChunk("", 0)), 4096, 500);
+        testAwsChunkedIllegalChunks(bucket, "with-final-chunk-chunk-underreports-size", "%s%s".formatted(buildFakeChunk(longDummyContent, 4096), buildFakeChunk("", 0)), 4096, 400);
 
         // Illegal signature and no final chunk - chunk misreports its size
-        testAwsChunkedIllegalChunks(bucket, "no-final-chunk-chunk-overreports-size", buildFakeChunk(longDummyContent, 9_000_000), 4096, 500);
+        testAwsChunkedIllegalChunks(bucket, "no-final-chunk-chunk-overreports-size", buildFakeChunk(longDummyContent, 9_000_000), 4096, 400);
         // Illegal signature with a final chunk - chunk misreports its size
-        testAwsChunkedIllegalChunks(bucket, "with-final-chunk-chunk-overreports-size", "%s%s".formatted(buildFakeChunk(longDummyContent, 9_000_000), buildFakeChunk("", 0)), 4096, 500);
+        testAwsChunkedIllegalChunks(bucket, "with-final-chunk-chunk-overreports-size", "%s%s".formatted(buildFakeChunk(longDummyContent, 9_000_000), buildFakeChunk("", 0)), 4096, 400);
         Thread.sleep(1000);
         assertThat(listFilesInS3Bucket(storageClient, bucket)).isEmpty();
     }
diff --git a/trino-aws-proxy/src/test/java/io/trino/aws/proxy/server/rest/TestAwsChunkedInputStream.java b/trino-aws-proxy/src/test/java/io/trino/aws/proxy/server/rest/TestAwsChunkedInputStream.java
@@ -282,7 +282,7 @@ private static void tryReadAwsChunkedData(String chunkedData, int decodedContent
     private static void testIllegalAwsChunkedData(String chunkedData, int decodedContentLength, TestingChunkSigningSession signingSession, ChunkReader readerMethod)
     {
         ByteArrayOutputStream testOutput = new ByteArrayOutputStream();
-        assertThatThrownBy(() -> readerMethod.read(chunkedData, decodedContentLength, signingSession, testOutput)).isInstanceOfAny(IllegalStateException.class, WebApplicationException.class, IOException.class);
+        assertThatThrownBy(() -> readerMethod.read(chunkedData, decodedContentLength, signingSession, testOutput)).isInstanceOfAny(WebApplicationException.class, IOException.class);
         assertThat(testOutput.toByteArray().length).isLessThan(decodedContentLength);
     }
 
@@ -377,7 +377,7 @@ public void testChunkedInputStreamNoClosingChunk()
         InputStream in = new AwsChunkedInputStream(inputStream, new DummyChunkSigningSession(), rawBytes.length);
         byte[] tmp = new byte[5];
         // altered from original test. Our AwsChunkedInputStream is improved and throws when the final chunk is missing or bad
-        assertThrows(IOException.class, () -> in.read(tmp));
+        assertThrows(WebApplicationException.class, () -> in.read(tmp));
     }
 
     // Truncated stream (missing closing CRLF)
@@ -393,7 +393,7 @@ public void testCorruptChunkedInputStreamTruncatedCRLF()
                     InputStream in = new AwsChunkedInputStream(inputStream, new DummyChunkSigningSession(), rawBytes.length);
                     byte[] tmp = new byte[5];
                     // altered from original test. Our AwsChunkedInputStream is improved and throws when the final chunk is missing or bad
-                    assertThrows(IOException.class, () -> in.read(tmp));
+                    assertThrows(WebApplicationException.class, () -> in.read(tmp));
                     try {
                         in.close();
                     }
@@ -413,7 +413,7 @@ public void testCorruptChunkedInputStreamMissingCRLF()
         InputStream in = new AwsChunkedInputStream(inputStream, new DummyChunkSigningSession(), rawBytes.length);
         byte[] buffer = new byte[300];
         ByteArrayOutputStream out = new ByteArrayOutputStream();
-        assertThrows(IOException.class, () -> {
+        assertThrows(WebApplicationException.class, () -> {
             int len;
             while ((len = in.read(buffer)) > 0) {
                 out.write(buffer, 0, len);
@@ -431,7 +431,7 @@ public void testCorruptChunkedInputStreamMissingLF()
         byte[] rawBytes = s.getBytes(UTF_8);
         ByteArrayInputStream inputStream = new ByteArrayInputStream(rawBytes);
         InputStream in = new AwsChunkedInputStream(inputStream, new DummyChunkSigningSession(), rawBytes.length);
-        assertThrows(IOException.class, in::read);
+        assertThrows(WebApplicationException.class, in::read);
         in.close();
     }
 
@@ -444,7 +444,7 @@ public void testCorruptChunkedInputStreamInvalidSize()
         byte[] rawBytes = s.getBytes(UTF_8);
         ByteArrayInputStream inputStream = new ByteArrayInputStream(rawBytes);
         InputStream in = new AwsChunkedInputStream(inputStream, new DummyChunkSigningSession(), rawBytes.length);
-        assertThrows(IOException.class, in::read);
+        assertThrows(WebApplicationException.class, in::read);
         in.close();
     }
 
@@ -457,7 +457,7 @@ public void testCorruptChunkedInputStreamNegativeSize()
         byte[] rawBytes = s.getBytes(UTF_8);
         ByteArrayInputStream inputStream = new ByteArrayInputStream(rawBytes);
         InputStream in = new AwsChunkedInputStream(inputStream, new DummyChunkSigningSession(), rawBytes.length);
-        assertThrows(IOException.class, in::read);
+        assertThrows(WebApplicationException.class, in::read);
         in.close();
     }
 
@@ -472,20 +472,18 @@ public void testCorruptChunkedInputStreamTruncatedChunk()
         InputStream in = new AwsChunkedInputStream(inputStream, new DummyChunkSigningSession(), rawBytes.length);
         byte[] buffer = new byte[300];
         assertEquals(2, in.read(buffer));
-        assertThrows(IOException.class, () -> in.read(buffer));
+        assertThrows(WebApplicationException.class, () -> in.read(buffer));
         in.close();
     }
 
     @Test
     public void testCorruptChunkedInputStreamClose()
-            throws IOException
     {
         String s = "whatever;chunk-signature=0\r\n01234\r\n5;chunk-signature=0\r\n56789\r\n0;chunk-signature=0\r\n\r\n";
         byte[] rawBytes = s.getBytes(UTF_8);
         ByteArrayInputStream inputStream = new ByteArrayInputStream(rawBytes);
-        try (InputStream in = new AwsChunkedInputStream(inputStream, new DummyChunkSigningSession(), rawBytes.length)) {
-            assertThrows(IOException.class, in::read);
-        }
+        InputStream in = new AwsChunkedInputStream(inputStream, new DummyChunkSigningSession(), rawBytes.length);
+        assertThrows(WebApplicationException.class, in::read);
     }
 
     @Test

Original file line number	Diff line number	Diff line change
`@@ -15,14 +15,15 @@`
`15`	`15`
`16`	`16`	`import com.google.common.base.Splitter;`
`17`	`17`	`import io.trino.aws.proxy.spi.signing.ChunkSigningSession;`
	`18`	`+import jakarta.ws.rs.WebApplicationException;`
`18`	`19`
`19`		`-import java.io.EOFException;`
`20`	`20`	`import java.io.IOException;`
`21`	`21`	`import java.io.InputStream;`
`22`	`22`	`import java.util.List;`
`23`	`23`	`import java.util.Map;`
`24`	`24`	`import java.util.Optional;`
`25`	`25`
	`26`	`+import static jakarta.ws.rs.core.Response.Status.BAD_REQUEST;`
`26`	`27`	`import static java.util.Objects.requireNonNull;`
`27`	`28`
`28`	`29`	`class AwsChunkedInputStream`
`@@ -61,7 +62,7 @@ public int read()`
`61`	`62`
`62`	`63`	`int i = delegate.read();`
`63`	`64`	`if (i < 0) {`
`64`		`- throw new EOFException("Unexpected end of stream");`
	`65`	`+ throw new WebApplicationException("Unexpected end of stream", BAD_REQUEST);`
`65`	`66`	`}`
`66`	`67`
`67`	`68`	`chunkSigningSession.write((byte) (i & 0xff));`
`@@ -82,7 +83,7 @@ public int read(byte[] b, int off, int len)`
`82`	`83`
`83`	`84`	`int count = delegate.read(b, off, len);`
`84`	`85`	`if (count < 0) {`
`85`		`- throw new EOFException("Unexpected end of stream");`
	`86`	`+ throw new WebApplicationException("Unexpected end of stream", BAD_REQUEST);`
`86`	`87`	`}`
`87`	`88`
`88`	`89`	`chunkSigningSession.write(b, off, count);`
`@@ -116,7 +117,7 @@ private void updateBytesRemaining(int count)`
`116`	`117`	`nextChunk();`
`117`	`118`	`}`
`118`	`119`	`else if (bytesRemainingInChunk < 0) {`
`119`		`- throw new IllegalStateException("bytesRemainingInChunk has gone negative: " + bytesRemainingInChunk);`
	`120`	`+ throw new WebApplicationException("bytesRemainingInChunk has gone negative: " + bytesRemainingInChunk, BAD_REQUEST);`
`120`	`121`	`}`
`121`	`122`	`}`
`122`	`123`
`@@ -195,10 +196,11 @@ private void nextChunk()`
`195`	`196`	`while (false);`
`196`	`197`
`197`	`198`	`if (!success) {`
`198`		`- throw new IOException("Invalid chunk header: " + header);`
	`199`	`+ throw new WebApplicationException("Invalid chunk header: " + header, BAD_REQUEST);`
`199`	`200`	`}`
`200`	`201`	`if (bytesAccountedFor > decodedContentLength) {`
`201`		`- throw new IllegalStateException("chunked data headers report a larger size than originally declared in the request: declared %s sent %s".formatted(decodedContentLength, bytesAccountedFor));`
	`202`	`+ throw new WebApplicationException("chunked data headers report a larger size than originally declared in the request: declared %s sent %s".formatted(decodedContentLength, bytesAccountedFor),`
	`203`	`+ BAD_REQUEST);`
`202`	`204`	`}`
`203`	`205`	`}`
`204`	`206`
`@@ -207,7 +209,7 @@ private void readEmptyLine()`
`207`	`209`	`{`
`208`	`210`	`String crLf = readLine();`
`209`	`211`	`if (!crLf.isEmpty()) {`
`210`		`- throw new IOException("Expected CR/LF. Instead read: " + crLf);`
	`212`	`+ throw new WebApplicationException("Expected CR/LF. Instead read: " + crLf, BAD_REQUEST);`
`211`	`213`	`}`
`212`	`214`	`}`
`213`	`215`
`@@ -219,7 +221,7 @@ private String readLine()`
`219`	`221`	`int i = delegate.read();`
`220`	`222`	`if (i < 0) {`
`221`	`223`	`delegateIsDone = true;`
`222`		`- throw new EOFException("Unexpected end of stream");`
	`224`	`+ throw new WebApplicationException("Unexpected end of stream", BAD_REQUEST);`
`223`	`225`	`}`
`224`	`226`	`if (i == '\r') {`
`225`	`227`	`break;`
`@@ -229,7 +231,7 @@ private String readLine()`
`229`	`231`
`230`	`232`	`int i = delegate.read();`
`231`	`233`	`if (i != '\n') {`
`232`		`- throw new IOException("Expected LF. Instead read: " + i);`
	`234`	`+ throw new WebApplicationException("Expected LF. Instead read: " + i, BAD_REQUEST);`
`233`	`235`	`}`
`234`	`236`
`235`	`237`	`return line.toString();`