Add full conditional write support to S3 test fixture (#138542)

In #133030 we added limited support for conditional writes in
`S3HttpHandler`, allowing callers to prevent overwriting an existing
blob with an `If-None-Match: *` precondition header. This commit extends
the implementation to include support for the `If-Match: <etag>`
precondition header allowing callers to perform atomic compare-and-set
operations which overwrite existing objects.

Author: DaveCTurner

test/fixtures/s3-fixture/src/main/java/fixture/s3/S3HttpHandler.java

@@ -45,6 +45,7 @@
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
@@ -248,6 +249,9 @@ public void handle(final HttpExchange exchange) throws IOException {
                     if (isProtectOverwrite(exchange)) {
                         throw new AssertionError("If-None-Match: * header is not supported here");
                     }
+                    if (getRequiredExistingETag(exchange) != null) {
+                        throw new AssertionError("If-Match: * header is not supported here");
+                    }
 
                     var sourceBlob = blobs.get(copySource);
                     if (sourceBlob == null) {
@@ -406,8 +410,9 @@ public void handle(final HttpExchange exchange) throws IOException {
      * Update the blob contents if and only if the preconditions in the request are satisfied.
      *
      * @return {@link RestStatus#OK} if the blob contents were updated, or else a different status code to indicate the error: possibly
-     *         {@link RestStatus#CONFLICT} or {@link RestStatus#PRECONDITION_FAILED} if the object exists and the precondition requires it
-     *         not to.
+     *         {@link RestStatus#CONFLICT} in any case, but if not then either  {@link RestStatus#PRECONDITION_FAILED} if the object exists
+     *         but doesn't match the specified precondition, or {@link RestStatus#NOT_FOUND} if the object doesn't exist but is required to
+     *         do so by the precondition.
      *
      * @see <a href="https://docs.aws.amazon.com/AmazonS3/latest/userguide/conditional-writes.html#conditional-error-response">AWS docs</a>
      */
@@ -418,6 +423,25 @@ private RestStatus updateBlobContents(HttpExchange exchange, String path, BytesR
                 : ESTestCase.randomFrom(RestStatus.PRECONDITION_FAILED, RestStatus.CONFLICT);
         }
 
+        final var requireExistingETag = getRequiredExistingETag(exchange);
+        if (requireExistingETag != null) {
+            final var responseCode = new AtomicReference<>(RestStatus.OK);
+            blobs.compute(path, (ignoredPath, existingContents) -> {
+                if (existingContents != null && requireExistingETag.equals(getEtagFromContents(existingContents))) {
+                    return newContents;
+                }
+
+                responseCode.set(
+                    ESTestCase.randomFrom(
+                        existingContents == null ? RestStatus.NOT_FOUND : RestStatus.PRECONDITION_FAILED,
+                        RestStatus.CONFLICT
+                    )
+                );
+                return existingContents;
+            });
+            return responseCode.get();
+        }
+
         blobs.put(path, newContents);
         return RestStatus.OK;
     }
@@ -598,6 +622,9 @@ private static boolean isProtectOverwrite(final HttpExchange exchange) {
             return false;
         }
 
+        if (exchange.getRequestHeaders().get("If-Match") != null) {
+            throw new AssertionError("Handling both If-None-Match and If-Match headers is not supported");
+        }
         if (ifNoneMatch.size() != 1) {
             throw new AssertionError("multiple If-None-Match headers found: " + ifNoneMatch);
         }
@@ -609,6 +636,29 @@ private static boolean isProtectOverwrite(final HttpExchange exchange) {
         throw new AssertionError("invalid If-None-Match header: " + ifNoneMatch);
     }
 
+    @Nullable // if no If-Match header found
+    private static String getRequiredExistingETag(final HttpExchange exchange) {
+        final var ifMatch = exchange.getRequestHeaders().get("If-Match");
+
+        if (ifMatch == null) {
+            return null;
+        }
+
+        if (exchange.getRequestHeaders().get("If-None-Match") != null) {
+            throw new AssertionError("Handling both If-None-Match and If-Match headers is not supported");
+        }
+
+        final var iterator = ifMatch.iterator();
+        if (iterator.hasNext()) {
+            final var result = iterator.next();
+            if (iterator.hasNext() == false) {
+                return result;
+            }
+        }
+
+        throw new AssertionError("multiple If-Match headers found: " + ifMatch);
+    }
+
     MultipartUpload putUpload(String path) {
         final var upload = new MultipartUpload(UUIDs.randomBase64UUID(), path);
         synchronized (uploads) {

test/fixtures/s3-fixture/src/test/java/fixture/s3/S3HttpHandlerTests.java

@@ -412,13 +412,29 @@ public void testExtractPartEtags() {
     }
 
     public void testPreventObjectOverwrite() {
+        ensureExactlyOneSuccess(new S3HttpHandler("bucket", "path"), null);
+    }
+
+    public void testConditionalOverwrite() {
         final var handler = new S3HttpHandler("bucket", "path");
 
-        var tasks = List.of(
-            createPutObjectTask(handler),
-            createPutObjectTask(handler),
-            createMultipartUploadTask(handler),
-            createMultipartUploadTask(handler)
+        final var originalBody = new BytesArray(randomAlphaOfLength(50).getBytes(StandardCharsets.UTF_8));
+        final var originalETag = S3HttpHandler.getEtagFromContents(originalBody);
+        assertEquals(RestStatus.OK, handleRequest(handler, "PUT", "/bucket/path/blob", originalBody).status());
+        assertEquals(
+            new TestHttpResponse(RestStatus.OK, originalBody, addETag(originalETag, TestHttpExchange.EMPTY_HEADERS)),
+            handleRequest(handler, "GET", "/bucket/path/blob")
+        );
+
+        ensureExactlyOneSuccess(handler, originalETag);
+    }
+
+    private static void ensureExactlyOneSuccess(S3HttpHandler handler, String originalETag) {
+        final var tasks = List.of(
+            createPutObjectTask(handler, originalETag),
+            createPutObjectTask(handler, originalETag),
+            createMultipartUploadTask(handler, originalETag),
+            createMultipartUploadTask(handler, originalETag)
         );
 
         runInParallel(tasks.size(), i -> tasks.get(i).consumer.run());
@@ -445,13 +461,38 @@ public void testPreventObjectOverwrite() {
         );
     }
 
-    private static TestWriteTask createPutObjectTask(S3HttpHandler handler) {
+    public void testPutObjectIfMatchWithBlobNotFound() {
+        final var handler = new S3HttpHandler("bucket", "path");
+        while (true) {
+            final var task = createPutObjectTask(handler, randomIdentifier());
+            task.consumer.run();
+            if (task.status == RestStatus.NOT_FOUND) {
+                break;
+            }
+            assertEquals(RestStatus.CONFLICT, task.status); // chosen randomly so eventually we will escape the loop
+        }
+    }
+
+    public void testCompleteMultipartUploadIfMatchWithBlobNotFound() {
+        final var handler = new S3HttpHandler("bucket", "path");
+        while (true) {
+            final var task = createMultipartUploadTask(handler, randomIdentifier());
+            task.consumer.run();
+            if (task.status == RestStatus.NOT_FOUND) {
+                break;
+            }
+            assertEquals(RestStatus.CONFLICT, task.status); // chosen randomly so eventually we will escape the loop
+        }
+    }
+
+    private static TestWriteTask createPutObjectTask(S3HttpHandler handler, @Nullable String originalETag) {
         return new TestWriteTask(
-            (task) -> task.status = handleRequest(handler, "PUT", "/bucket/path/blob", task.body, ifNoneMatchHeader()).status()
+            (task) -> task.status = handleRequest(handler, "PUT", "/bucket/path/blob", task.body, conditionalWriteHeader(originalETag))
+                .status()
         );
     }
 
-    private static TestWriteTask createMultipartUploadTask(S3HttpHandler handler) {
+    private static TestWriteTask createMultipartUploadTask(S3HttpHandler handler, @Nullable String originalETag) {
         final var multipartUploadTask = new TestWriteTask(
             (task) -> task.status = handleRequest(
                 handler,
@@ -465,7 +506,7 @@ private static TestWriteTask createMultipartUploadTask(S3HttpHandler handler) {
                           <PartNumber>1</PartNumber>
                        </Part>
                     </CompleteMultipartUpload>""", task.etag)),
-                ifNoneMatchHeader()
+                conditionalWriteHeader(originalETag)
             ).status()
         );
 
@@ -594,9 +635,13 @@ private static Headers contentRangeHeader(long start, long end, long length) {
         return headers;
     }
 
-    private static Headers ifNoneMatchHeader() {
+    private static Headers conditionalWriteHeader(@Nullable String originalEtag) {
         var headers = new Headers();
-        headers.put("If-None-Match", List.of("*"));
+        if (originalEtag == null) {
+            headers.put("If-None-Match", List.of("*"));
+        } else {
+            headers.put("If-Match", List.of(originalEtag));
+        }
         return headers;
     }