From b498e2fe4b53334ba6f70103df2940d165782e65 Mon Sep 17 00:00:00 2001
From: yizhaorong <243653385@qq.com>
Date: Thu, 28 May 2020 23:28:02 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E5=A4=8D=E4=B8=8A=E4=BC=A0=E9=97=AE?=
 =?UTF-8?q?=E9=A2=98?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../org/yzr/controller/StorageController.java   | 17 ++++++++++-------
 src/main/java/org/yzr/storage/StorageUtil.java  | 17 ++++++++++++-----
 src/main/resources/templates/index.html         |  2 +-
 3 files changed, 23 insertions(+), 13 deletions(-)

diff --git a/src/main/java/org/yzr/controller/StorageController.java b/src/main/java/org/yzr/controller/StorageController.java
index 71bab8a..ba3600d 100644
--- a/src/main/java/org/yzr/controller/StorageController.java
+++ b/src/main/java/org/yzr/controller/StorageController.java
@@ -6,15 +6,13 @@ import org.springframework.http.HttpHeaders;
 import org.springframework.http.MediaType;
 import org.springframework.http.ResponseEntity;
 import org.springframework.stereotype.Controller;
-import org.springframework.web.bind.annotation.GetMapping;
-import org.springframework.web.bind.annotation.PathVariable;
-import org.springframework.web.bind.annotation.PostMapping;
-import org.springframework.web.bind.annotation.RequestParam;
+import org.springframework.web.bind.annotation.*;
 import org.springframework.web.multipart.MultipartFile;
 import org.yzr.model.Storage;
 import org.yzr.service.StorageService;
 import org.yzr.storage.StorageUtil;
 import org.yzr.utils.CharUtil;
+import org.yzr.utils.response.BaseResponse;
 import org.yzr.utils.response.ResponseUtil;
 
 import java.io.IOException;
@@ -44,10 +42,15 @@ public class StorageController {
     }
 
     @PostMapping("/upload")
-    public Object upload(@RequestParam("file") MultipartFile file) throws IOException {
+    @ResponseBody
+    public BaseResponse upload(@RequestParam("file") MultipartFile file) throws IOException {
         String originalFilename = file.getOriginalFilename();
-        Storage Storage = storageUtil.store(file.getInputStream(), file.getSize(), file.getContentType(), originalFilename);
-        return ResponseUtil.ok(Storage);
+        Storage storage = storageUtil.store(file.getInputStream(), file.getSize(), file.getContentType(), originalFilename);
+        if (storage != null) {
+            return ResponseUtil.ok(storage);
+        } else {
+            return ResponseUtil.fail(401, "不支持的文件类型");
+        }
     }
 
     /**
diff --git a/src/main/java/org/yzr/storage/StorageUtil.java b/src/main/java/org/yzr/storage/StorageUtil.java
index 58ae4f9..7c73585 100644
--- a/src/main/java/org/yzr/storage/StorageUtil.java
+++ b/src/main/java/org/yzr/storage/StorageUtil.java
@@ -46,18 +46,25 @@ public class StorageUtil {
      * @param fileName      文件索引名
      */
     public Storage store(InputStream inputStream, long contentLength, String contentType, String fileName) {
+        if (!(contentType != null && contentType.equalsIgnoreCase("application/octet-stream"))) {
+            return null;
+        }
         String key = generateKey(fileName);
+        int len = 28;
+        PushbackInputStream pushbackInputStream = new PushbackInputStream(inputStream, len);
         try {
-            int len = 28;
             byte[] b = new byte[len];
-            FileType type = FileUtil.getType(input);
-            input.unread(b);
-            inputStream = input;
+            FileType type = FileUtil.getType(pushbackInputStream);
+            if (type != FileType.ZIP) {
+                pushbackInputStream.close();
+                return null;
+            }
+            pushbackInputStream.unread(b);
         } catch (Exception e) {
             e.printStackTrace();
         }
 
-        storage.store(inputStream, contentLength, contentType, key);
+        storage.store(pushbackInputStream, contentLength, contentType, key);
 
         String url = generateUrl(key);
         Storage storageInfo = new Storage();
diff --git a/src/main/resources/templates/index.html b/src/main/resources/templates/index.html
index fd471b8..349502f 100644
--- a/src/main/resources/templates/index.html
+++ b/src/main/resources/templates/index.html
@@ -138,7 +138,7 @@
         var file = files[0]
         //上传
         var xhr = new XMLHttpRequest();
-        xhr.open("post", "/app/upload", true);
+        xhr.open("post", "/upload", true);
         xhr.setRequestHeader("X-Requested-With", "XMLHttpRequest");
         // 获取上传进度
         xhr.upload.onprogress = function (event) {