From 97f8f4ae71f71b1024f5b8a377fc525cac87f1da Mon Sep 17 00:00:00 2001
From: Armin Braun <me@obrown.io>
Date: Wed, 26 Feb 2020 17:04:33 +0100
Subject: [PATCH] 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
---
 .../src/main/java/fixture/gcs/FakeOAuth2HttpHandler.java | 2 +-
 .../java/fixture/gcs/GoogleCloudStorageHttpHandler.java  | 9 +++++++--
 2 files changed, 8 insertions(+), 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
index 7dcaaf16f4a37..c49ae19dba0b5 100644
--- 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
index 944f025d6e3f5..e96a3db83133d 100644
--- 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 {