Refactor GCS fixture multipart parser (#125828)

Author: mhl-b

modules/repository-gcs/src/test/java/org/elasticsearch/repositories/gcs/GoogleCloudStorageBlobContainerRetriesTests.java

@@ -10,6 +10,7 @@
 
 import fixture.gcs.FakeOAuth2HttpHandler;
 import fixture.gcs.GoogleCloudStorageHttpHandler;
+import fixture.gcs.MultipartUpload;
 
 import com.google.api.client.http.HttpExecuteInterceptor;
 import com.google.api.client.http.HttpRequestInitializer;
@@ -43,7 +44,6 @@
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.SuppressForbidden;
 import org.elasticsearch.core.TimeValue;
-import org.elasticsearch.core.Tuple;
 import org.elasticsearch.http.ResponseInjectingHttpHandler;
 import org.elasticsearch.repositories.blobstore.AbstractBlobContainerRetriesTestCase;
 import org.elasticsearch.repositories.blobstore.ESMockAPIBasedRepositoryIntegTestCase;
@@ -62,15 +62,13 @@
 import java.util.Locale;
 import java.util.Map;
 import java.util.Objects;
-import java.util.Optional;
 import java.util.Queue;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 
-import static fixture.gcs.GoogleCloudStorageHttpHandler.parseMultipartRequestBody;
 import static fixture.gcs.TestUtils.createServiceAccount;
 import static java.nio.charset.StandardCharsets.UTF_8;
 import static org.elasticsearch.repositories.blobstore.BlobStoreTestUtil.randomPurpose;
@@ -80,7 +78,6 @@
 import static org.elasticsearch.repositories.gcs.GoogleCloudStorageClientSettings.ENDPOINT_SETTING;
 import static org.elasticsearch.repositories.gcs.GoogleCloudStorageClientSettings.READ_TIMEOUT_SETTING;
 import static org.elasticsearch.repositories.gcs.GoogleCloudStorageClientSettings.TOKEN_URI_SETTING;
-import static org.elasticsearch.test.hamcrest.OptionalMatchers.isPresent;
 import static org.hamcrest.Matchers.anyOf;
 import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.equalTo;
@@ -268,17 +265,16 @@ public void testWriteBlobWithRetries() throws Exception {
         final CountDown countDown = new CountDown(maxRetries);
 
         final BlobContainer blobContainer = blobContainerBuilder().maxRetries(maxRetries).build();
-        final byte[] bytes = randomBlobContent();
+        final byte[] bytes = randomBlobContent(0);
         httpServer.createContext("/upload/storage/v1/b/bucket/o", safeHandler(exchange -> {
             assertThat(exchange.getRequestURI().getQuery(), containsString("uploadType=multipart"));
             if (countDown.countDown()) {
-                Optional<Tuple<String, BytesReference>> content = parseMultipartRequestBody(exchange.getRequestBody());
-                assertThat(content, isPresent());
-                assertThat(content.get().v1(), equalTo(blobContainer.path().buildAsString() + "write_blob_max_retries"));
-                if (Objects.deepEquals(bytes, BytesReference.toBytes(content.get().v2()))) {
+                MultipartUpload multipartUpload = MultipartUpload.parseBody(exchange, exchange.getRequestBody());
+                assertEquals(multipartUpload.name(), blobContainer.path().buildAsString() + "write_blob_max_retries");
+                if (multipartUpload.content().equals(new BytesArray(bytes))) {
                     byte[] response = Strings.format("""
                         {"bucket":"bucket","name":"%s"}
-                        """, content.get().v1()).getBytes(UTF_8);
+                        """, multipartUpload.name()).getBytes(UTF_8);
                     exchange.getResponseHeaders().add("Content-Type", "application/json");
                     exchange.sendResponseHeaders(RestStatus.OK.getStatus(), response.length);
                     exchange.getResponseBody().write(response);

test/fixtures/gcs-fixture/src/main/java/fixture/gcs/GoogleCloudStorageHttpHandler.java

@@ -18,7 +18,6 @@
 import org.elasticsearch.common.io.Streams;
 import org.elasticsearch.common.regex.Regex;
 import org.elasticsearch.core.SuppressForbidden;
-import org.elasticsearch.core.Tuple;
 import org.elasticsearch.rest.RestStatus;
 import org.elasticsearch.rest.RestUtils;
 import org.elasticsearch.test.fixture.HttpHeaderParser;
@@ -27,26 +26,17 @@
 import org.elasticsearch.xcontent.XContentFactory;
 import org.elasticsearch.xcontent.XContentType;
 
-import java.io.BufferedReader;
 import java.io.IOException;
-import java.io.InputStream;
-import java.io.InputStreamReader;
 import java.net.URLDecoder;
 import java.util.HashMap;
 import java.util.Locale;
 import java.util.Map;
 import java.util.Objects;
-import java.util.Optional;
 import java.util.concurrent.atomic.AtomicInteger;
-import java.util.regex.Matcher;
-import java.util.regex.Pattern;
 import java.util.stream.Collectors;
-import java.util.zip.GZIPInputStream;
 
 import static fixture.gcs.MockGcsBlobStore.failAndThrow;
 import static java.nio.charset.StandardCharsets.UTF_8;
-import static java.util.stream.Collectors.joining;
-import static org.elasticsearch.core.Strings.format;
 
 /**
  * Minimal HTTP handler that acts as a Google Cloud Storage compliant server
@@ -183,26 +173,18 @@ public void handle(final HttpExchange exchange) throws IOException {
                 exchange.getResponseBody().write(response);
 
             } else if (Regex.simpleMatch("POST /upload/storage/v1/b/" + bucket + "/*uploadType=multipart*", request)) {
-                // Multipart upload
-                Optional<Tuple<String, BytesReference>> content = parseMultipartRequestBody(requestBody.streamInput());
-                if (content.isPresent()) {
+                try {
+                    final var multipartUpload = MultipartUpload.parseBody(exchange, requestBody.streamInput());
                     final Long ifGenerationMatch = parseOptionalLongParameter(exchange, IF_GENERATION_MATCH);
                     final MockGcsBlobStore.BlobVersion newBlobVersion = mockGcsBlobStore.updateBlob(
-                        content.get().v1(),
+                        multipartUpload.name(),
                         ifGenerationMatch,
-                        content.get().v2()
+                        multipartUpload.content()
                     );
                     writeBlobVersionAsJson(exchange, newBlobVersion);
-                } else {
-                    throw new AssertionError(
-                        "Could not read multi-part request to ["
-                            + request
-                            + "] with headers ["
-                            + new HashMap<>(exchange.getRequestHeaders())
-                            + "]"
-                    );
+                } catch (IllegalArgumentException e) {
+                    throw new AssertionError(e);
                 }
-
             } else if (Regex.simpleMatch("POST /upload/storage/v1/b/" + bucket + "/*uploadType=resumable*", request)) {
                 // Resumable upload initialization https://cloud.google.com/storage/docs/json_api/v1/how-tos/resumable-upload
                 final Map<String, String> params = new HashMap<>();
@@ -328,81 +310,6 @@ private static String httpServerUrl(final HttpExchange exchange) {
         return "http://" + exchange.getRequestHeaders().get("HOST").get(0);
     }
 
-    private static final Pattern NAME_PATTERN = Pattern.compile("\"name\":\"([^\"]*)\"");
-
-    public static Optional<Tuple<String, BytesReference>> parseMultipartRequestBody(final InputStream requestBody) throws IOException {
-        Tuple<String, BytesReference> content = null;
-        final BytesReference fullRequestBody;
-        try (InputStream in = new GZIPInputStream(requestBody)) {
-            fullRequestBody = Streams.readFully(in);
-        }
-        String name = null;
-        boolean skippedEmptyLine = false;
-        int startPos = 0;
-        int endPos = 0;
-        while (startPos < fullRequestBody.length()) {
-            do {
-                endPos = fullRequestBody.indexOf((byte) '\r', endPos + 1);
-            } while (endPos >= 0 && fullRequestBody.get(endPos + 1) != '\n');
-            boolean markAndContinue = false;
-            final String bucketPrefix = "{\"bucket\":";
-            if (startPos > 0) {
-                startPos += 2;
-            }
-            if (name == null || skippedEmptyLine == false) {
-                if ((skippedEmptyLine == false && endPos == startPos)
-                    || (fullRequestBody.get(startPos) == '-' && fullRequestBody.get(startPos + 1) == '-')) {
-                    markAndContinue = true;
-                } else {
-                    final String start = fullRequestBody.slice(startPos, Math.min(endPos - startPos, bucketPrefix.length())).utf8ToString();
-                    if (start.toLowerCase(Locale.ROOT).startsWith("content")) {
-                        markAndContinue = true;
-                    } else if (start.startsWith(bucketPrefix)) {
-                        markAndContinue = true;
-                        final String line = fullRequestBody.slice(
-                            startPos + bucketPrefix.length(),
-                            endPos - startPos - bucketPrefix.length()
-                        ).utf8ToString();
-                        Matcher matcher = NAME_PATTERN.matcher(line);
-                        if (matcher.find()) {
-                            name = matcher.group(1);
-                        }
-                    }
-                }
-                skippedEmptyLine = markAndContinue && endPos == startPos;
-                startPos = endPos;
-            } else {
-                while (isEndOfPart(fullRequestBody, endPos) == false) {
-                    endPos = fullRequestBody.indexOf((byte) '\r', endPos + 1);
-                }
-                content = Tuple.tuple(name, fullRequestBody.slice(startPos, endPos - startPos));
-                break;
-            }
-        }
-        if (content == null) {
-            final InputStream stream = fullRequestBody.streamInput();
-            logger.warn(
-                () -> format(
-                    "Failed to find multi-part upload in [%s]",
-                    new BufferedReader(new InputStreamReader(stream)).lines().collect(joining("\n"))
-                )
-            );
-        }
-        return Optional.ofNullable(content);
-    }
-
-    private static final byte[] END_OF_PARTS_MARKER = "\r\n--__END_OF_PART__".getBytes(UTF_8);
-
-    private static boolean isEndOfPart(BytesReference fullRequestBody, int endPos) {
-        for (int i = 0; i < END_OF_PARTS_MARKER.length; i++) {
-            final byte b = END_OF_PARTS_MARKER[i];
-            if (fullRequestBody.get(endPos + i) != b) {
-                return false;
-            }
-        }
-        return true;
-    }
-
     private static String requireHeader(HttpExchange exchange, String headerName) {
         final String headerValue = exchange.getRequestHeaders().getFirst(headerName);
         if (headerValue != null) {