From 21b752dd25df60b4fb3ef905a2d203ded1abd598 Mon Sep 17 00:00:00 2001
From: longjuan <769022681@qq.com>
Date: Sat, 28 Jan 2023 14:40:10 +0800
Subject: [PATCH] feat: check the file already exists before uploading (#11)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Fixes https://github.com/halo-dev/halo/issues/2945
```release-note
Add file check with the same name when uploading
```

我这里使用了ConcurrentHashMap避免**同时**上传两个同名文件导致文件覆盖问题，在本地存储策略中是由操作系统保证的
不知道这样是不是一个好方法

S3AsyncClient中没有`doesObjectExist`类似的方法，官方的文档也让用`headObject`捕获异常的办法来判断文件是否存在，详见https://github.com/aws/aws-sdk-java-v2/blob/master/docs/LaunchChangelog.md 中搜索`doesObjectExist`
---
 .../run/halo/s3os/S3OsAttachmentHandler.java  | 39 ++++++++++++++++---
 1 file changed, 34 insertions(+), 5 deletions(-)

diff --git a/src/main/java/run/halo/s3os/S3OsAttachmentHandler.java b/src/main/java/run/halo/s3os/S3OsAttachmentHandler.java
index ce8864d..f02874c 100644
--- a/src/main/java/run/halo/s3os/S3OsAttachmentHandler.java
+++ b/src/main/java/run/halo/s3os/S3OsAttachmentHandler.java
@@ -8,6 +8,7 @@ import org.springframework.core.io.buffer.DataBuffer;
 import org.springframework.http.MediaType;
 import org.springframework.http.MediaTypeFactory;
 import org.springframework.web.server.ServerErrorException;
+import org.springframework.web.server.ServerWebInputException;
 import org.springframework.web.util.UriUtils;
 import reactor.core.publisher.Mono;
 import run.halo.app.core.extension.attachment.Attachment;
@@ -21,6 +22,7 @@ import run.halo.app.infra.utils.JsonUtils;
 import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
 import software.amazon.awssdk.core.SdkResponse;
 import software.amazon.awssdk.core.async.AsyncRequestBody;
+import software.amazon.awssdk.http.SdkHttpResponse;
 import software.amazon.awssdk.regions.Region;
 import software.amazon.awssdk.services.s3.S3AsyncClient;
 import software.amazon.awssdk.services.s3.S3Configuration;
@@ -33,6 +35,7 @@ import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.UUID;
+import java.util.concurrent.ConcurrentHashMap;
 
 @Slf4j
 @Extension
@@ -40,6 +43,7 @@ public class S3OsAttachmentHandler implements AttachmentHandler {
 
     private static final String OBJECT_KEY = "s3os.plugin.halo.run/object-key";
     private static final int MULTIPART_MIN_PART_SIZE = 5 * 1024 * 1024;
+    private final Map<String, Object> uploadingFile = new ConcurrentHashMap<>();
 
     @Override
     public Mono<Attachment> upload(UploadContext uploadContext) {
@@ -126,23 +130,45 @@ public class S3OsAttachmentHandler implements AttachmentHandler {
     }
 
     Mono<ObjectDetail> upload(UploadContext uploadContext, S3OsProperties properties) {
-        var s3client = buildS3AsyncClient(properties);
-
         var originFilename = uploadContext.file().filename();
         var objectKey = properties.getObjectName(originFilename);
         var contentType = MediaTypeFactory.getMediaType(originFilename)
                 .orElse(MediaType.APPLICATION_OCTET_STREAM).toString();
+        var uploadingMapKey = properties.getBucket() + "/" + objectKey;
+        // deduplication of uploading files
+        if (uploadingFile.put(uploadingMapKey, uploadingMapKey) != null) {
+            return Mono.error(new ServerWebInputException("文件 " + originFilename + " 已存在，建议更名后重试。"));
+        }
+
+        var s3client = buildS3AsyncClient(properties);
 
         var uploadState = new UploadState(properties.getBucket(), objectKey);
 
         return Mono
+                // check whether file exists
+                .fromFuture(s3client.headObject(HeadObjectRequest.builder()
+                        .bucket(properties.getBucket())
+                        .key(objectKey)
+                        .build()))
+                .onErrorResume(NoSuchKeyException.class, e -> {
+                    var builder = HeadObjectResponse.builder();
+                    builder.sdkHttpResponse(SdkHttpResponse.builder().statusCode(404).build());
+                    return Mono.just(builder.build());
+                })
+                .flatMap(response -> {
+                    if (response != null && response.sdkHttpResponse() != null && response.sdkHttpResponse().isSuccessful()) {
+                        return Mono.error(new ServerWebInputException("文件 " + originFilename + " 已存在，建议更名后重试。"));
+                    }else {
+                        return Mono.just(uploadState);
+                    }
+                })
                 // init multipart upload
-                .fromFuture(s3client.createMultipartUpload(
+                .flatMap(state -> Mono.fromFuture(s3client.createMultipartUpload(
                         CreateMultipartUploadRequest.builder()
                                 .bucket(properties.getBucket())
                                 .contentType(contentType)
                                 .key(objectKey)
-                                .build()))
+                                .build())))
                 .flatMapMany((response) -> {
                     checkResult(response, "createMultipartUpload");
                     uploadState.setUploadId(response.uploadId());
@@ -194,7 +220,10 @@ public class S3OsAttachmentHandler implements AttachmentHandler {
                     return new ObjectDetail(properties.getBucket(), objectKey, response);
                 })
                 // close client
-                .doFinally((signalType) -> s3client.close());
+                .doFinally((signalType) -> {
+                    uploadingFile.remove(uploadingMapKey);
+                    s3client.close();
+                });
     }