diff --git a/src/main/java/com/webank/ddcms/controller/FileController.java b/src/main/java/com/webank/ddcms/controller/FileController.java index 9d61f287..219fbe97 100644 --- a/src/main/java/com/webank/ddcms/controller/FileController.java +++ b/src/main/java/com/webank/ddcms/controller/FileController.java @@ -1,9 +1,12 @@ package com.webank.ddcms.controller; +import com.google.common.base.Strings; +import com.google.common.io.Files; import com.webank.ddcms.enums.CodeEnum; import com.webank.ddcms.service.FileService; import com.webank.ddcms.vo.common.CommonResponse; import lombok.extern.slf4j.Slf4j; +import lombok.var; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.core.io.Resource; import org.springframework.http.ResponseEntity; @@ -12,6 +15,9 @@ import org.springframework.web.multipart.MultipartFile; import java.io.IOException; +import java.util.HashMap; +import java.util.HashSet; +import java.util.Set; @RestController @RequestMapping("/api/file") @@ -20,12 +26,23 @@ public class FileController { @Autowired private FileService fileService; + private final Set ALLOWED_FILE_TYPES = new HashSet(){{ + add("png"); + add("jpg"); + add("jpeg"); + add("pdf"); + }}; + @PostMapping("/upload") public CommonResponse handleFileUpload(@RequestParam("file") MultipartFile file) throws Exception { String contentType = file.getContentType(); - log.info("containt type {}",contentType); - if (contentType == null || !contentType.startsWith("image/")){ + String fileName = file.getName(); + if (contentType == null || !contentType.startsWith("image/")) { + return CommonResponse.error(CodeEnum.PARAMETER_ERROR); + } + String ext = Files.getFileExtension(fileName); + if (Strings.isNullOrEmpty(ext) || !ALLOWED_FILE_TYPES.contains(ext)) { return CommonResponse.error(CodeEnum.PARAMETER_ERROR); } @@ -33,6 +50,7 @@ public CommonResponse handleFileUpload(@RequestParam("file") MultipartFile file) return CommonResponse.success(filename); } + @GetMapping("/download") public ResponseEntity handleFileDownload(@RequestParam("filename") String filename) throws IOException {