hotfix: add file type checking logic for presigned-url generating method

YAPP-Github · Oct 27, 2024 · 0a41edd · 0a41edd
1 parent 88cdcaa
commit 0a41edd
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 20 deletions.
diff --git a/photo-service/src/main/java/kr/mafoo/photo/exception/ErrorCode.java b/photo-service/src/main/java/kr/mafoo/photo/exception/ErrorCode.java
@@ -19,6 +19,7 @@ public enum ErrorCode {
     PHOTO_DISPLAY_INDEX_NOT_VALID("PE0005", "옮기려는 대상 사진 인덱스가 유효하지 않습니다"),
 
     PRE_SIGNED_URL_EXCEED_MAXIMUM("OE0001", "한 번에 생성할 수 있는 Pre-signed url 최대치를 초과했습니다"),
+    PRE_SIGNED_URL_BANNED_FILE_TYPE("OE0002", "Pre-signed url 발급이 허용되지 않는 파일 형식입니다")
 
     ;
     private final String code;

diff --git a/photo-service/src/main/java/kr/mafoo/photo/exception/PreSignedUrlBannedFileType.java b/photo-service/src/main/java/kr/mafoo/photo/exception/PreSignedUrlBannedFileType.java
@@ -0,0 +1,7 @@
+package kr.mafoo.photo.exception;
+
+public class PreSignedUrlBannedFileType extends DomainException {
+    public PreSignedUrlBannedFileType() {
+        super(ErrorCode.PRE_SIGNED_URL_BANNED_FILE_TYPE);
+    }
+}
diff --git a/photo-service/src/main/java/kr/mafoo/photo/service/ObjectStorageService.java b/photo-service/src/main/java/kr/mafoo/photo/service/ObjectStorageService.java
@@ -6,21 +6,23 @@
 import com.amazonaws.services.s3.model.GeneratePresignedUrlRequest;
 import com.amazonaws.services.s3.model.ObjectMetadata;
 import com.amazonaws.services.s3.model.PutObjectRequest;
+import kr.mafoo.photo.exception.PreSignedUrlBannedFileType;
 import kr.mafoo.photo.exception.PreSignedUrlExceedMaximum;
 import lombok.extern.slf4j.Slf4j;
 import org.springframework.beans.factory.annotation.Value;
 import org.springframework.stereotype.Service;
 
 import lombok.RequiredArgsConstructor;
+import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
 import java.io.ByteArrayInputStream;
 import java.io.InputStream;
 
 import java.net.URL;
 import java.util.Date;
+import java.util.List;
 import java.util.UUID;
-import java.util.stream.Stream;
 
 @Slf4j
 @Service
@@ -59,25 +61,35 @@ public Mono<String> uploadFile(byte[] fileByte) {
     }
 
     public Mono<String[]> createPreSignedUrls(String[] fileNames, String memberId) {
-        return Mono.fromCallable(() -> {
-            if (fileNames.length > 30) {
-                throw new PreSignedUrlExceedMaximum();
-            }
-
-            return Stream.of(fileNames)
-                    .map(fileName -> generatePresignedUrl(fileName, memberId).toString())
-                    .toArray(String[]::new);
-        });
+        if (fileNames.length > 30) {
+            return Mono.error(new PreSignedUrlExceedMaximum());
+        }
+
+        return Flux.fromArray(fileNames)
+                .flatMap(fileName -> generatePresignedUrlForImage(fileName, memberId)
+                        .map(URL::toString))
+                .collectList()
+                .map(list -> list.toArray(new String[0]));
     }
 
-    private URL generatePresignedUrl(String fileName, String memberId) {
-        String fileType = extractFileType(fileName);
-        String filePath = String.format("%s/photo/%s.%s", memberId, UUID.randomUUID(), fileType);
-        Date expiration = new Date(System.currentTimeMillis() + presignedUrlExpiration);
-
-        return amazonS3Client.generatePresignedUrl(new GeneratePresignedUrlRequest(bucketName, filePath)
-                .withMethod(HttpMethod.PUT)
-                .withExpiration(expiration));
+    private Mono<URL> generatePresignedUrlForImage(String fileName, String memberId) {
+        List<String> allowedFileTypes = List.of("jpg", "jpeg", "png");
+
+        return extractFileType(fileName)
+                .flatMap(fileType -> {
+                    if (allowedFileTypes.contains(fileType)) {
+                        String filePath = String.format("%s/photo/%s.%s", memberId, UUID.randomUUID(), fileType);
+                        Date expiration = new Date(System.currentTimeMillis() + presignedUrlExpiration);
+
+                        return Mono.just(
+                                amazonS3Client.generatePresignedUrl(new GeneratePresignedUrlRequest(bucketName, filePath)
+                                        .withMethod(HttpMethod.PUT)
+                                        .withExpiration(expiration))
+                        );
+                    } else {
+                        return Mono.error(new PreSignedUrlBannedFileType());
+                    }
+                });
     }
 
     public Mono<String> setObjectPublicRead(String filePath) {
@@ -93,8 +105,10 @@ public Mono<String> setObjectPublicRead(String filePath) {
         });
     }
 
-    private String extractFileType(String fileName) {
-        return fileName.substring(fileName.lastIndexOf(".") + 1);
+    private Mono<String> extractFileType(String fileName) {
+        return Mono.just(
+                fileName.substring(fileName.lastIndexOf(".") + 1)
+        );
     }
 
     private String generateFileLink(String keyName) {