Fix GCS Mock Range Downloads (#52804)

We were not correctly respecting the download range which lead to the GCS SDK client closing the connection at times. Also, fixes another instance of failing to drain the request fully before sending the response headers. Closes #51446
elastic · Feb 26, 2020 · 97f8f4a · 97f8f4a
1 parent 649cd63
commit 97f8f4a
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 3 deletions.
diff --git a/test/fixtures/gcs-fixture/src/main/java/fixture/gcs/FakeOAuth2HttpHandler.java b/test/fixtures/gcs-fixture/src/main/java/fixture/gcs/FakeOAuth2HttpHandler.java
@@ -35,11 +35,11 @@ public class FakeOAuth2HttpHandler implements HttpHandler {
     @Override
     public void handle(final HttpExchange exchange) throws IOException {
         try {
+            while (exchange.getRequestBody().read(BUFFER) >= 0) ;
             byte[] response = ("{\"access_token\":\"foo\",\"token_type\":\"Bearer\",\"expires_in\":3600}").getBytes(UTF_8);
             exchange.getResponseHeaders().add("Content-Type", "application/json");
             exchange.sendResponseHeaders(RestStatus.OK.getStatus(), response.length);
             exchange.getResponseBody().write(response);
-            while (exchange.getRequestBody().read(BUFFER) >= 0) ;
         } finally {
             int read = exchange.getRequestBody().read();
             assert read == -1 : "Request body should have been fully read here but saw [" + read + "]";

diff --git a/test/fixtures/gcs-fixture/src/main/java/fixture/gcs/GoogleCloudStorageHttpHandler.java b/test/fixtures/gcs-fixture/src/main/java/fixture/gcs/GoogleCloudStorageHttpHandler.java
@@ -142,9 +142,14 @@ public void handle(final HttpExchange exchange) throws IOException {
                     if (matcher.find() == false) {
                         throw new AssertionError("Range bytes header does not match expected format: " + range);
                     }
-
-                    BytesReference response = Integer.parseInt(matcher.group(1)) == 0 ? blob : BytesArray.EMPTY;
+                    final int offset = Integer.parseInt(matcher.group(1));
+                    final int end = Integer.parseInt(matcher.group(2));
+                    BytesReference response = blob;
                     exchange.getResponseHeaders().add("Content-Type", "application/octet-stream");
+                    final int bufferedLength = response.length();
+                    if (offset > 0 || bufferedLength > end) {
+                        response = response.slice(offset, Math.min(end + 1 - offset, bufferedLength - offset));
+                    }
                     exchange.sendResponseHeaders(RestStatus.OK.getStatus(), response.length());
                     response.writeTo(exchange.getResponseBody());
                 } else {