1 files changed, 31 insertions, 33 deletions

diff --git a/juick-common/src/main/java/com/juick/server/util/HttpUtils.java b/juick-common/src/main/java/com/juick/server/util/HttpUtils.java
index e348d6a2..dbdbc062 100644
--- a/juick-common/src/main/java/com/juick/server/util/HttpUtils.java
+++ b/juick-common/src/main/java/com/juick/server/util/HttpUtils.java
@@ -22,6 +22,7 @@ import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.springframework.web.multipart.MultipartFile;
 
+import java.io.BufferedInputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.net.URI;
@@ -38,29 +39,37 @@ import java.util.UUID;
 public class HttpUtils {
     private static final Logger logger = LoggerFactory.getLogger(HttpUtils.class);
 
-    public static URI receiveMultiPartFile(MultipartFile attach, String tmpDir) {
-        if (attach !=null && !attach.isEmpty()) {
-            String partname = attach.getOriginalFilename();
-            if (partname != null && partname.length() > 0) {
-                String attachmentType = partname.substring(partname.length() - 3).toLowerCase();
-                if (attachmentType.equals("jpg") || attachmentType.equals("peg") || attachmentType.equals("png")) {
-                    if (attachmentType.equals("peg")) {
-                        attachmentType = "jpg";
-                    }
-                    String attachmentFName = DigestUtils.md5Hex(UUID.randomUUID().toString()) + "." + attachmentType;
-                    try {
-                        Files.write(Paths.get(tmpDir, attachmentFName),
-                                attach.getBytes());
-                        return URI.create(String.format("juick://%s", attachmentFName));
-                    } catch (IOException e) {
-                        logger.warn("file receive error", e);
-                    }
+    public static URI receiveMultiPartFile(MultipartFile attach, String tmpDir) throws IOException {
+        if (attach != null && !attach.isEmpty()) {
+            InputStream attachmentStream = new BufferedInputStream(attach.getInputStream());
+            String guessedContentType = URLConnection.guessContentTypeFromStream(attachmentStream);
+            String attachmentType = attachmentTypeFromContentType(guessedContentType);
+            if (attachmentType.equals("jpg") || attachmentType.equals("png")) {
+                String attachmentFName = DigestUtils.md5Hex(UUID.randomUUID().toString()) + "." + attachmentType;
+                try {
+                    Files.write(Paths.get(tmpDir, attachmentFName),
+                            attach.getBytes());
+                    return URI.create(String.format("juick://%s", attachmentFName));
+                } catch (IOException e) {
+                    logger.warn("file receive error", e);
                 }
             }
+            logger.warn("file type is unknown: {}", attach.getOriginalFilename());
         }
         return URI.create(StringUtils.EMPTY);
     }
-    public static URI downloadImage(URL url, String tmpDir) throws Exception {
+
+    private static String attachmentTypeFromContentType(String mime) throws IOException {
+        if (mime != null && mime.equals("image/jpeg")) {
+            return "jpg";
+        } else if (mime != null && mime.equals("image/png")) {
+            return "png";
+        } else {
+            throw new IOException("Wrong file type: " + mime);
+        }
+    }
+
+    public static URI downloadImage(URL url, String tmpDir) throws IOException {
         URLConnection urlConn;
         try {
             urlConn = url.openConnection();
@@ -69,26 +78,15 @@ public class HttpUtils {
             throw e;
         }
 
-        try (InputStream is = urlConn.getInputStream()) {
-            String mime = urlConn.getContentType();
+        try (InputStream is = new BufferedInputStream(urlConn.getInputStream())) {
+            String mime = URLConnection.guessContentTypeFromStream(is);
 
-            String attachmentType;
-            if (mime != null && mime.equals("image/jpeg")) {
-                attachmentType = "jpg";
-            } else if (mime != null && mime.equals("image/png")) {
-                attachmentType = "png";
-            } else if (url.getFile().toLowerCase().endsWith("jpg")) {
-                attachmentType = "jpg";
-            } else if (url.getFile().toLowerCase().endsWith("png")) {
-                attachmentType = "png";
-            } else {
-                throw new Exception("Wrong file type: " + mime);
-            }
+            String attachmentType = attachmentTypeFromContentType(mime);
 
             String attachmentFName = DigestUtils.md5Hex(UUID.randomUUID().toString()) + "." + attachmentType;
             Files.copy(is, Paths.get(tmpDir, attachmentFName));
             return URI.create(String.format("juick://%s", attachmentFName));
-        } catch (Exception e) {
+        } catch (IOException e) {
             logger.error(String.format("Failed download image by url: %s", url.toString()), e);
             throw e;
         }