refactor: Determine content-type dynamically

PolinaPolupan · Mar 4, 2025 · 2bd039b · 2bd039b
1 parent ce1b562
commit 2bd039b
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 5 deletions.
diff --git a/src/main/java/com/example/mypixel/controller/ImageUploadController.java b/src/main/java/com/example/mypixel/controller/ImageUploadController.java
@@ -3,11 +3,15 @@
 
 import java.io.IOException;
 import java.net.URI;
+import java.net.URLConnection;
+import java.nio.file.Files;
+import java.nio.file.Paths;
 import java.util.List;
 import java.util.stream.Collectors;
 
 import com.example.mypixel.exception.InvalidImageFormat;
 import com.example.mypixel.storage.StorageService;
+import jakarta.servlet.ServletContext;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.core.io.Resource;
 import org.springframework.http.HttpHeaders;
@@ -23,6 +27,9 @@ public class ImageUploadController {
 
     private final StorageService storageService;
 
+    @Autowired
+    private ServletContext servletContext;
+
     @Autowired
     public ImageUploadController(StorageService storageService) {
         this.storageService = storageService;
@@ -46,13 +53,19 @@ public ResponseEntity<Resource> serveFile(@PathVariable String filename) throws
         if (file == null)
             return ResponseEntity.notFound().build();
 
-        String location = ServletUriComponentsBuilder
-                .fromCurrentRequest()
-                .toUriString();
+        String contentType = null;
+        try {
+            contentType = Files.probeContentType(Paths.get(file.getFile().getAbsolutePath()));
+        } catch (IOException e) {
+            contentType = URLConnection.guessContentTypeFromName(file.getFilename());
+        }
+
+        if (contentType == null) {
+            contentType = "application/octet-stream";
+        }
 
         return ResponseEntity.ok()
-                .contentType(MediaType.IMAGE_JPEG)
-                .header(HttpHeaders.LOCATION, location)
+                .contentType(MediaType.valueOf(contentType))
                 .header(HttpHeaders.CONTENT_DISPOSITION,
                 "attachment; filename=\"" + file.getFilename() + "\"").body(file);
     }

diff --git a/src/test/java/com/example/mypixel/ImageUploadControllerTests.java b/src/test/java/com/example/mypixel/ImageUploadControllerTests.java
@@ -1,5 +1,6 @@
 package com.example.mypixel;
 
+import java.io.File;
 import java.nio.file.Paths;
 import java.util.stream.Stream;
 
@@ -66,6 +67,7 @@ public void shouldServeImage() throws Exception {
         String filename = "test.jpg";
         Resource mockResource = mock(Resource.class);
         given(mockResource.getFilename()).willReturn(filename);
+        given(mockResource.getFile()).willReturn(new File(filename));
         given(mockResource.exists()).willReturn(true);
         given(mockResource.isReadable()).willReturn(true);
 
@@ -74,6 +76,7 @@ public void shouldServeImage() throws Exception {
         mockMvc.perform(get("/images/{filename}", filename))
                 .andExpect(status().isOk())
                 .andExpect(header().string("Content-Disposition", containsString("attachment; filename=\"test.jpg\"")))
+                .andExpect(header().string("Content-Type", containsString("image/jpeg")))
                 .andExpect(content().contentType(MediaType.IMAGE_JPEG));
     }