Ignore uploaded file extension

* guess content type from content stream

3 files changed, 63 insertions, 43 deletions

diff --git a/juick-common/src/main/java/com/juick/server/CommandsManager.java b/juick-common/src/main/java/com/juick/server/CommandsManager.java
index b9f27940..7dd372b1 100644
--- a/juick-common/src/main/java/com/juick/server/CommandsManager.java
+++ b/juick-common/src/main/java/com/juick/server/CommandsManager.java
@@ -106,12 +106,17 @@ public class CommandsManager {
         }
         // new message
         String body = TagUtils.toString(tags).equals(input) ? StringUtils.EMPTY : input.substring(TagUtils.toString(tags).length()).trim();
-        String attachmentType = StringUtils.isNotEmpty(attachment.toString()) ? attachment.toString().substring(attachment.toString().length() - 3) : null;
+        boolean haveAttachment = StringUtils.isNotEmpty(attachment.toString());
+        String attachmentFName = null;
+        String attachmentType = null;
+        if (haveAttachment) {
+             attachmentFName = attachment.getScheme().equals("juick") ? attachment.getHost()
+                    : HttpUtils.downloadImage(attachment.toURL(), tmpDir).getHost();
+             attachmentType = attachmentFName.substring(attachmentFName.length() - 3);
+        }
         int mid = messagesService.createMessage(user.getUid(), body, attachmentType, tags);
         subscriptionService.subscribeMessage(messagesService.getMessage(mid), user);
-        if (StringUtils.isNotEmpty(attachmentType)) {
-            String attachmentFName = attachment.getScheme().equals("juick") ? attachment.getHost()
-                    : HttpUtils.downloadImage(attachment.toURL(), tmpDir).getHost();
+        if (haveAttachment) {
             String fname = String.format("%d.%s", mid, attachmentType);
             imagesService.saveImageWithPreviews(attachmentFName, fname);
         }
@@ -494,13 +499,17 @@ public class CommandsManager {
                 return CommandResult.fromString("Tags are NOT updated (5 tags maximum?)");
             }
         } else {
-            String attachmentStr = attachment.toString();
-            String attachmentType = StringUtils.isNotEmpty(attachmentStr) ? attachmentStr.substring(attachmentStr.length() - 3) : null;
+            boolean haveAttachment = StringUtils.isNotEmpty(attachment.toString());
+            String attachmentFName = null;
+            String attachmentType = null;
+            if (haveAttachment) {
+                attachmentFName = attachment.getScheme().equals("juick") ? attachment.getHost()
+                        : HttpUtils.downloadImage(attachment.toURL(), tmpDir).getHost();
+                attachmentType = attachmentFName.substring(attachmentFName.length() - 3);
+            }
             int newrid = messagesService.createReply(mid, rid, user, txt, attachmentType);
             applicationEventPublisher.publishEvent(new MessageReadEvent(this, user, msg));
-            if (StringUtils.isNotEmpty(attachmentType)) {
-                String attachmentFName = attachment.getScheme().equals("juick") ? attachment.getHost()
-                        : HttpUtils.downloadImage(attachment.toURL(), tmpDir).getHost();
+            if (haveAttachment) {
                 String fname = String.format("%d-%d.%s", mid, newrid, attachmentType);
                 imagesService.saveImageWithPreviews(attachmentFName, fname);
             }
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;
         }
diff --git a/juick-server/src/test/java/com/juick/server/tests/ServerTests.java b/juick-server/src/test/java/com/juick/server/tests/ServerTests.java
index 2d45aae0..e4e5f298 100644
--- a/juick-server/src/test/java/com/juick/server/tests/ServerTests.java
+++ b/juick-server/src/test/java/com/juick/server/tests/ServerTests.java
@@ -56,7 +56,6 @@ import org.springframework.test.web.servlet.MvcResult;
 import org.springframework.util.FileSystemUtils;
 import org.springframework.util.LinkedMultiValueMap;
 import org.springframework.util.MultiValueMap;
-import org.springframework.web.client.RestTemplate;
 import org.springframework.web.context.WebApplicationContext;
 import org.springframework.web.util.UriComponents;
 import org.springframework.web.util.UriComponentsBuilder;
@@ -84,6 +83,7 @@ import java.io.*;
 import java.net.URI;
 import java.net.URISyntaxException;
 import java.nio.file.Files;
+import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.sql.Timestamp;
 import java.time.Instant;
@@ -1098,4 +1098,17 @@ public class ServerTests {
         assertThat(postJpgCmyk.getNewMessage().get().getAttachment().getMedium().getHeight(), is(1024));
         assertThat(postJpgCmyk.getNewMessage().get().getAttachment().getSmall().getHeight(), is(512));
     }
+    @Test
+    public void changeExtensionWhenReceiveFileWithWrongContentType() throws Exception {
+        Path pngOutput = Paths.get(tmpDir, "cmyk.png");
+        Files.delete(pngOutput);
+        Files.copy(getClass().getClassLoader().getResourceAsStream("cmyk.jpg"), pngOutput);
+        assertThat(pngOutput.toFile().exists(), is(true));
+        CommandResult postJpgCmyk = commandsManager.processCommand(ugnich, "YO", pngOutput.toUri());
+        assertThat(postJpgCmyk.getNewMessage().isPresent(), is(true));
+        assertThat(postJpgCmyk.getNewMessage().get().getAttachmentType(), is("jpg"));
+        CommandResult replyJpgCmyk = commandsManager.processCommand(ugnich, String.format("#%d YO", postJpgCmyk.getNewMessage().get().getMid()), pngOutput.toUri());
+        assertThat(replyJpgCmyk.getNewMessage().isPresent(), is(true));
+        assertThat(replyJpgCmyk.getNewMessage().get().getAttachmentType(), is("jpg"));
+    }
 }