Support object tagging when creating multipart uploads (#88)

Author: Robothy

local-s3-core/src/main/java/com/robothy/s3/core/model/answers/GetObjectAns.java

@@ -2,6 +2,7 @@
 
 import java.io.InputStream;
 import java.util.Map;
+import java.util.Optional;
 import lombok.Builder;
 import lombok.Getter;
 
@@ -29,4 +30,5 @@ public class GetObjectAns {
 
   private Map<String, String> userMetadata;
 
+  private int taggingCount;
 }

local-s3-core/src/main/java/com/robothy/s3/core/model/internal/UploadMetadata.java

@@ -3,6 +3,7 @@
 import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
 import com.robothy.s3.core.converters.deserializer.UploadPartMetadataMapConverter;
 import java.util.NavigableMap;
+import java.util.Optional;
 import java.util.concurrent.ConcurrentSkipListMap;
 import lombok.AllArgsConstructor;
 import lombok.Builder;
@@ -19,8 +20,14 @@ public class UploadMetadata {
 
   private String contentType;
 
+  private String[][] tagging;
+
   @JsonDeserialize(converter = UploadPartMetadataMapConverter.class)
   @Builder.Default
   private NavigableMap<Integer, UploadPartMetadata> parts = new ConcurrentSkipListMap<>();
 
+
+  public Optional<String[][]> getTagging() {
+    return Optional.ofNullable(tagging);
+  }
 }

local-s3-core/src/main/java/com/robothy/s3/core/model/request/CreateMultipartUploadOptions.java

@@ -1,5 +1,6 @@
 package com.robothy.s3.core.model.request;
 
+import java.util.Optional;
 import lombok.Builder;
 import lombok.Getter;
 
@@ -9,4 +10,9 @@ public class CreateMultipartUploadOptions {
 
   private String contentType;
 
+  private String[][] tagging;
+
+  public Optional<String[][]> getTagging() {
+    return Optional.ofNullable(tagging);
+  }
 }

local-s3-core/src/main/java/com/robothy/s3/core/service/CompleteMultipartUploadService.java

@@ -71,6 +71,7 @@ default CompleteMultipartUploadAns completeMultipartUpload(String bucket, String
           .size(size)
           .content(in)
           .contentType(uploadMetadata.getContentType())
+          .tagging(uploadMetadata.getTagging().orElse(null))
           .build();
 
       putObjectAns = putObject(bucket, key, putObjectOptions);

local-s3-core/src/main/java/com/robothy/s3/core/service/CreateMultipartUploadService.java

@@ -32,6 +32,7 @@ default String createMultipartUpload(String bucket, String key, CreateMultipartU
     uploads.get(key).put(uploadId, UploadMetadata.builder()
         .contentType(options.getContentType())
         .createDate(System.currentTimeMillis())
+        .tagging(options.getTagging().orElse(null))
         .build());
     return uploadId;
   }

local-s3-core/src/main/java/com/robothy/s3/core/service/GetObjectService.java

@@ -48,6 +48,7 @@ static GetObjectAns getObjectFromUnVersionedBucket(BucketMetadata bucketMetadata
         .content(metadataOnly ? null : storage.getInputStream(latestObject.getFileId()))
         .etag(latestObject.getEtag())
         .userMetadata(latestObject.getUserMetadata())
+        .taggingCount(latestObject.getTagging().map(tagging -> tagging.length).orElse(0))
         .build();
   }
 
@@ -108,6 +109,7 @@ static GetObjectAns getObject(BucketMetadata bucketMetadata, Storage storage,
           .size(versionedObjectMetadata.getSize())
           .content(metadataOnly ? null : storage.getInputStream(versionedObjectMetadata.getFileId()))
           .etag(versionedObjectMetadata.getEtag())
+          .taggingCount(versionedObjectMetadata.getTagging().map(tagging -> tagging.length).orElse(0))
           .userMetadata(versionedObjectMetadata.getUserMetadata())
           .build();
     }

local-s3-interationtest/src/test/java/com/robothy/s3/test/MultipartUploadIntegrationTest.java

@@ -29,6 +29,17 @@
 import java.net.URL;
 import java.util.List;
 import org.junit.jupiter.api.Test;
+import software.amazon.awssdk.core.ResponseInputStream;
+import software.amazon.awssdk.core.sync.RequestBody;
+import software.amazon.awssdk.services.s3.S3Client;
+import software.amazon.awssdk.services.s3.model.CompleteMultipartUploadResponse;
+import software.amazon.awssdk.services.s3.model.CompletedPart;
+import software.amazon.awssdk.services.s3.model.CreateBucketResponse;
+import software.amazon.awssdk.services.s3.model.CreateMultipartUploadResponse;
+import software.amazon.awssdk.services.s3.model.GetObjectResponse;
+import software.amazon.awssdk.services.s3.model.Tag;
+import software.amazon.awssdk.services.s3.model.Tagging;
+import software.amazon.awssdk.services.s3.model.UploadPartResponse;
 
 public class MultipartUploadIntegrationTest {
 
@@ -227,4 +238,34 @@ void testListParts(AmazonS3 s3) {
         s3.listParts(new ListPartsRequest(bucketName, "a.txt", initiateMultipartUploadResult.getUploadId())));
   }
 
+  @Test
+  @LocalS3
+  void testCreateMultipartUploadsWithTagging(S3Client s3) throws Exception {
+    s3.createBucket(builder -> builder.bucket("my-bucket"));
+    Tag tag1 = Tag.builder().key("k1").value("v1").build();
+    Tag tag2 = Tag.builder().key("k2").value("v2").build();
+    CreateMultipartUploadResponse multipartUpload = s3.createMultipartUpload(builder -> builder.bucket("my-bucket").key("a.txt")
+        .tagging(Tagging.builder().tagSet(tag1, tag2).build()));
+    UploadPartResponse part1 = s3.uploadPart(b -> b.bucket("my-bucket")
+            .uploadId(multipartUpload.uploadId()).key("a.txt").partNumber(1), RequestBody.fromString("Hello"));
+    UploadPartResponse part2 = s3.uploadPart(b -> b.bucket("my-bucket").uploadId(multipartUpload.uploadId()).key("a.txt").partNumber(2),
+            RequestBody.fromString("World"));
+
+    CompletedPart completedPart1 = CompletedPart.builder().partNumber(1).eTag(part1.eTag()).build();
+    CompletedPart completedPart2 = CompletedPart.builder().partNumber(2).eTag(part2.eTag()).build();
+    s3.completeMultipartUpload(b -> b.bucket("my-bucket").key("a.txt").multipartUpload(upload -> upload.parts(
+        completedPart1, completedPart2)).uploadId(multipartUpload.uploadId()));
+
+    ResponseInputStream<GetObjectResponse> completedObject =
+        s3.getObject(b -> b.bucket("my-bucket").key("a.txt"));
+    assertEquals("HelloWorld", new String(completedObject.readAllBytes()));
+    Integer tagCount = completedObject.response().tagCount();
+    assertEquals(2, tagCount);
+
+    List<Tag> tags = s3.getObjectTagging(b -> b.bucket("my-bucket").key("a.txt")).tagSet();
+    assertEquals(2, tags.size());
+    assertTrue(tags.contains(tag1));
+    assertTrue(tags.contains(tag2));
+  }
+
 }

local-s3-rest/src/main/java/com/robothy/s3/rest/constants/AmzHeaderNames.java

@@ -19,6 +19,8 @@ public class AmzHeaderNames {
 
   public static final String X_AMZ_TAGGING = "x-amz-tagging";
 
+  public static final String X_AMZ_TAGGING_COUNT = "x-amz-tagging-count";
+
   /**
    * Specifies the source object for the copy operation.
    */

local-s3-rest/src/main/java/com/robothy/s3/rest/handler/CreateMultipartUploadController.java

@@ -10,6 +10,7 @@
 import com.robothy.s3.rest.assertions.RequestAssertions;
 import com.robothy.s3.rest.model.response.InitiateMultipartUploadResult;
 import com.robothy.s3.rest.service.ServiceFactory;
+import com.robothy.s3.rest.utils.RequestUtils;
 import com.robothy.s3.rest.utils.ResponseUtils;
 import io.netty.handler.codec.http.HttpResponseStatus;
 
@@ -33,6 +34,7 @@ public void handle(HttpRequest request, HttpResponse response) throws Exception
     String key = RequestAssertions.assertObjectKeyProvided(request);
     String contentType = request.parameter("content-type").orElse("octet/stream");
     String uploadId = uploadService.createMultipartUpload(bucket, key, CreateMultipartUploadOptions.builder()
+        .tagging(RequestUtils.extractTagging(request).orElse(null))
         .contentType(contentType).build());
     InitiateMultipartUploadResult result = InitiateMultipartUploadResult.builder()
         .bucket(bucket)

local-s3-rest/src/main/java/com/robothy/s3/rest/handler/GetObjectController.java

@@ -47,6 +47,7 @@ public void handle(HttpRequest request, HttpResponse response) throws Exception
       ResponseUtils.addETag(response, getObjectAns.getEtag());
       response.status(HttpResponseStatus.OK)
           .write(content)
+          .putHeader(AmzHeaderNames.X_AMZ_TAGGING_COUNT, getObjectAns.getTaggingCount())
           .putHeader(HttpHeaderNames.CONTENT_TYPE.toString(), getObjectAns.getContentType())
           .putHeader(HttpHeaderNames.CONTENT_LENGTH.toString(), getObjectAns.getSize());
       getObjectAns.getUserMetadata().forEach((k, v) -> response.putHeader(AmzHeaderNames.X_AMZ_META_PREFIX + k, v));

local-s3-rest/src/main/java/com/robothy/s3/rest/handler/PutObjectController.java

@@ -3,7 +3,6 @@
 import com.robothy.netty.http.HttpRequest;
 import com.robothy.netty.http.HttpRequestHandler;
 import com.robothy.netty.http.HttpResponse;
-import com.robothy.s3.core.exception.LocalS3InvalidArgumentException;
 import com.robothy.s3.core.model.answers.PutObjectAns;
 import com.robothy.s3.core.model.request.PutObjectOptions;
 import com.robothy.s3.core.service.ObjectService;
@@ -18,8 +17,6 @@
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Objects;
-import java.util.Optional;
-import org.apache.commons.lang3.StringUtils;
 
 /**
  * Handle <a href="https://docs.aws.amazon.com/AmazonS3/latest/API/API_PutObject.html">PutObject<a/>.
@@ -43,7 +40,7 @@ public void handle(HttpRequest request, HttpResponse response) throws Exception
         .contentType(request.header(HttpHeaderNames.CONTENT_TYPE).orElse(null))
         .size(decodedBody.getDecodedContentLength())
         .content(decodedBody.getDecodedBody())
-        .tagging(extractTagging(request))
+        .tagging(RequestUtils.extractTagging(request).orElse(null))
         .userMetadata(extractUserMetadata(request))
         .build();
 
@@ -61,29 +58,6 @@ public void handle(HttpRequest request, HttpResponse response) throws Exception
     ResponseUtils.addAmzRequestId(response);
   }
 
-  // parse tagging from x-amz-tagging header in the put object request.
-  String[][] extractTagging(HttpRequest request) {
-    Optional<String> taggingOpt = request.header(AmzHeaderNames.X_AMZ_TAGGING);
-    String tagging;
-    if (taggingOpt.isEmpty() || StringUtils.isBlank(tagging = taggingOpt.get())) {
-      return null;
-    }
-
-    String[] tags = tagging.split("&");
-    String[][] tagSet = new String[tags.length][2];
-    for (int i = 0; i < tags.length; i++) {
-      String[] kv = tags[i].split("=");
-      if (kv.length != 2) {
-        throw new LocalS3InvalidArgumentException(AmzHeaderNames.X_AMZ_TAGGING, "Invalid tagging format.");
-      }
-
-      tagSet[i][0] = kv[0];
-      tagSet[i][1] = kv[1];
-    }
-
-    return tagSet;
-  }
-
   Map<String, String> extractUserMetadata(HttpRequest request) {
     Map<String, String> userMetadata = new HashMap<>();
     request.getHeaders()

local-s3-rest/src/main/java/com/robothy/s3/rest/utils/RequestUtils.java

@@ -1,13 +1,15 @@
 package com.robothy.s3.rest.utils;
 
 import com.robothy.netty.http.HttpRequest;
+import com.robothy.s3.core.exception.LocalS3InvalidArgumentException;
 import com.robothy.s3.rest.constants.AmzHeaderNames;
 import com.robothy.s3.rest.constants.AmzHeaderValues;
 import com.robothy.s3.rest.model.request.DecodedAmzRequestBody;
 import io.netty.buffer.ByteBufInputStream;
 import io.netty.handler.codec.http.HttpHeaderNames;
 import java.io.InputStream;
 import java.util.Optional;
+import org.apache.commons.lang3.StringUtils;
 
 /**
  * HTTP Request related utils.
@@ -56,4 +58,32 @@ public static Optional<String> getETag(HttpRequest request) {
     return request.header(HttpHeaderNames.ETAG.toString());
   }
 
+  /**
+   * Extract tagging from the HTTP header.
+   *
+   * @param request HTTP request.
+   * @return tagging.
+   */
+  public static Optional<String[][]> extractTagging(HttpRequest request) {
+    Optional<String> taggingOpt = request.header(AmzHeaderNames.X_AMZ_TAGGING);
+    String tagging;
+    if (taggingOpt.isEmpty() || StringUtils.isBlank(tagging = taggingOpt.get())) {
+      return Optional.empty();
+    }
+
+    String[] tags = tagging.split("&");
+    String[][] tagSet = new String[tags.length][2];
+    for (int i = 0; i < tags.length; i++) {
+      String[] kv = tags[i].split("=");
+      if (kv.length != 2) {
+        throw new LocalS3InvalidArgumentException(AmzHeaderNames.X_AMZ_TAGGING, "Invalid tagging format.");
+      }
+
+      tagSet[i][0] = kv[0];
+      tagSet[i][1] = kv[1];
+    }
+
+    return Optional.of(tagSet);
+  }
+
 }

local-s3-rest/src/test/java/com/robothy/s3/rest/handler/PutObjectControllerTest.java renamed to local-s3-rest/src/test/java/com/robothy/s3/rest/utils/RequestUtilsTest.java

@@ -1,27 +1,23 @@
-package com.robothy.s3.rest.handler;
+package com.robothy.s3.rest.utils;
 
-import static org.junit.jupiter.api.Assertions.assertEquals;
-import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.*;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 import com.robothy.netty.http.HttpRequest;
 import com.robothy.s3.core.exception.LocalS3InvalidArgumentException;
 import com.robothy.s3.rest.constants.AmzHeaderNames;
-import com.robothy.s3.rest.service.ServiceFactory;
 import java.util.Optional;
 import org.junit.jupiter.api.Test;
 
-class PutObjectControllerTest {
+class RequestUtilsTest {
 
   @Test
   void testExtractTagging() {
 
     HttpRequest request = mock(HttpRequest.class);
-    ServiceFactory serviceFactory = mock(ServiceFactory.class);
-    PutObjectController controller = new PutObjectController(serviceFactory);
-
     when(request.header(AmzHeaderNames.X_AMZ_TAGGING)).thenReturn(Optional.of("key1=value1&key2=value2"));
-    String[][] tagArray = controller.extractTagging(request);
+    String[][] tagArray = RequestUtils.extractTagging(request).orElse(null);
+    assertNotNull(tagArray);
     assertEquals(2, tagArray.length);
     assertEquals("key1", tagArray[0][0]);
     assertEquals("value1", tagArray[0][1]);
@@ -30,12 +26,13 @@ void testExtractTagging() {
 
 
     when(request.header(AmzHeaderNames.X_AMZ_TAGGING)).thenReturn(Optional.of("key1=value1"));
-    tagArray = controller.extractTagging(request);
+    tagArray = RequestUtils.extractTagging(request).orElse(null);
+    assertNotNull(tagArray);
     assertEquals(1, tagArray.length);
     assertEquals("key1", tagArray[0][0]);
     assertEquals("value1", tagArray[0][1]);
 
     when(request.header(AmzHeaderNames.X_AMZ_TAGGING)).thenReturn(Optional.of("invalid"));
-    assertThrows(LocalS3InvalidArgumentException.class, () -> controller.extractTagging(request));
+    assertThrows(LocalS3InvalidArgumentException.class, () -> RequestUtils.extractTagging(request));
   }
 }