.

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,19 +26,31 @@ public class FileController {
 
   @Autowired private FileService fileService;
 
+  private final Set<String> ALLOWED_FILE_TYPES = new HashSet<String>(){{
+    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);
     }
 
     String filename = fileService.uploadFile(file);
     return CommonResponse.success(filename);
   }
 
+
   @GetMapping("/download")
   public ResponseEntity<Resource> handleFileDownload(@RequestParam("filename") String filename)
       throws IOException {