Initial support for #165.

structurizr · Oct 9, 2024 · d22785e · d22785e
1 parent bc802b0
commit d22785e
Show file tree

Hide file tree

Showing 5 changed files with 87 additions and 55 deletions.
diff --git a/src/main/java/com/structurizr/lite/component/workspace/FileSystemWorkspaceComponentImpl.java b/src/main/java/com/structurizr/lite/component/workspace/FileSystemWorkspaceComponentImpl.java
@@ -6,7 +6,7 @@
 import com.structurizr.lite.component.search.SearchComponent;
 import com.structurizr.lite.domain.WorkspaceMetaData;
 import com.structurizr.lite.util.DateUtils;
-import com.structurizr.lite.util.InputStreamAndContentLength;
+import com.structurizr.lite.util.Image;
 import com.structurizr.util.DslTemplate;
 import com.structurizr.util.WorkspaceUtils;
 import com.structurizr.validation.WorkspaceScopeValidatorFactory;
@@ -18,7 +18,6 @@
 
 import java.io.BufferedWriter;
 import java.io.File;
-import java.io.FileInputStream;
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
@@ -31,6 +30,8 @@ class FileSystemWorkspaceComponentImpl implements WorkspaceComponent {
 
     private static final Log log = LogFactory.getLog(FileSystemWorkspaceComponentImpl.class);
 
+    private static final String IMAGES_DIRECTORY = "images";
+
     private File dataDirectory;
     private String filename;
 
@@ -233,15 +234,21 @@ public String getError() {
     }
 
     @Override
-    public InputStreamAndContentLength getImage(long workspaceId, String diagramKey) throws WorkspaceComponentException {
+    public Image getImage(long workspaceId, String filename) throws WorkspaceComponentException {
         try {
-            File imagesDirectory = getPathToWorkspaceImages(workspaceId);
-            File file = new File(imagesDirectory, diagramKey);
+            // first try .structurizr/{workspaceId}/images
+            File file = new File(getPathToWorkspaceWorkDirectoryImages(workspaceId), filename);
             if (file.exists()) {
-                return new InputStreamAndContentLength(new FileInputStream(file), file.length());
+                return new Image(file);
+            } else {
+                // otherwise try {workspaceId}/images
+                file = new File(getPathToWorkspaceImages(workspaceId), filename);
+                if (file.exists()) {
+                    return new Image(file);
+                }
             }
         } catch (Exception e) {
-            String message = "Could not get " + diagramKey + " for workspace";
+            String message = "Could not get image \"" + filename + "\" for workspace";
             log.warn(e.getMessage() + " - " + message);
         }
 
@@ -254,7 +261,7 @@ public boolean putImage(long workspaceId, String filename, String imageAsBase64D
         byte[] decodedImage = Base64.getDecoder().decode(base64Image.getBytes(StandardCharsets.UTF_8));
 
         try {
-            File imagesDirectory = getPathToWorkspaceImages(workspaceId);
+            File imagesDirectory = getPathToWorkspaceWorkDirectoryImages(workspaceId);
             File file = new File(imagesDirectory, filename);
             Files.write(file.toPath(), decodedImage);
 
@@ -265,8 +272,21 @@ public boolean putImage(long workspaceId, String filename, String imageAsBase64D
         }
     }
 
+    private File getPathToWorkspaceWorkDirectoryImages(long workspaceId) {
+        File path = new File(new File(Configuration.getInstance().getWorkDirectory(), "" + workspaceId), IMAGES_DIRECTORY);
+        if (!path.exists()) {
+            try {
+                Files.createDirectories(path.toPath());
+            } catch (IOException e) {
+                log.error(e);
+            }
+        }
+
+        return path;
+    }
+
     private File getPathToWorkspaceImages(long workspaceId) {
-        File path = new File(new File(Configuration.getInstance().getWorkDirectory(), "" + workspaceId), "images");
+        File path = new File(getDataDirectory(workspaceId), IMAGES_DIRECTORY);
         if (!path.exists()) {
             try {
                 Files.createDirectories(path.toPath());

diff --git a/src/main/java/com/structurizr/lite/component/workspace/WorkspaceComponent.java b/src/main/java/com/structurizr/lite/component/workspace/WorkspaceComponent.java
@@ -2,7 +2,7 @@
 
 import com.structurizr.Workspace;
 import com.structurizr.lite.domain.WorkspaceMetaData;
-import com.structurizr.lite.util.InputStreamAndContentLength;
+import com.structurizr.lite.util.Image;
 
 import java.util.List;
 
@@ -21,7 +21,7 @@ public interface WorkspaceComponent {
 
     String getError();
 
-    InputStreamAndContentLength getImage(long workspaceId, String diagramKey) throws WorkspaceComponentException;
+    Image getImage(long workspaceId, String filename) throws WorkspaceComponentException;
 
     boolean putImage(long workspaceId, String filename, String imageAsBase64) throws WorkspaceComponentException;
 

diff --git a/src/main/java/com/structurizr/lite/util/Image.java b/src/main/java/com/structurizr/lite/util/Image.java
@@ -0,0 +1,28 @@
+package com.structurizr.lite.util;
+
+import java.io.File;
+import java.io.FileInputStream;
+import java.io.InputStream;
+import java.net.URLConnection;
+
+public final class Image {
+
+    private final File file;
+
+    public Image(File file) {
+        this.file = file;
+    }
+
+    public InputStream getInputStream() throws Exception {
+        return new FileInputStream(file);
+    }
+
+    public long getContentLength() {
+        return file.length();
+    }
+
+    public String getContentType() {
+        return URLConnection.guessContentTypeFromName(file.getName());
+    }
+
+}
diff --git a/src/main/java/com/structurizr/lite/util/InputStreamAndContentLength.java b/src/main/java/com/structurizr/lite/util/InputStreamAndContentLength.java
diff --git a/src/main/java/com/structurizr/lite/web/DiagramsController.java b/src/main/java/com/structurizr/lite/web/DiagramsController.java
@@ -4,14 +4,14 @@
 import com.structurizr.lite.component.workspace.WorkspaceComponentException;
 import com.structurizr.lite.domain.WorkspaceMetaData;
 import com.structurizr.lite.util.HtmlUtils;
-import com.structurizr.lite.util.InputStreamAndContentLength;
+import com.structurizr.lite.util.Image;
 import com.structurizr.view.PaperSize;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.springframework.core.io.InputStreamResource;
-import org.springframework.core.io.Resource;
 import org.springframework.http.HttpHeaders;
 import org.springframework.http.MediaType;
+import org.springframework.http.ResponseEntity;
 import org.springframework.stereotype.Controller;
 import org.springframework.ui.ModelMap;
 import org.springframework.web.bind.annotation.*;
@@ -97,29 +97,36 @@ private void addAccessControlAllowHeaders(HttpServletResponse response) {
     }
 
     @ResponseBody
-    @RequestMapping(value = "/workspace/{workspaceId}/images/{diagramKey}.png", method = RequestMethod.GET, produces = MediaType.IMAGE_PNG_VALUE)
-    public Resource getImage(@PathVariable("workspaceId") long workspaceId,
-                             @PathVariable("diagramKey") String diagramKey,
+    @RequestMapping(value = "/workspace/images/{filename}", method = RequestMethod.GET)
+    public ResponseEntity getImage(@PathVariable("filename") String filename,
+                                   HttpServletResponse response) {
+        return getImage(1, filename, response);
+    }
+
+    @ResponseBody
+    @RequestMapping(value = "/workspace/{workspaceId}/images/{filename}", method = RequestMethod.GET)
+    public ResponseEntity getImage(@PathVariable("workspaceId") long workspaceId,
+                             @PathVariable("filename") String filename,
                              HttpServletResponse response) {
-        diagramKey = HtmlUtils.filterHtml(diagramKey);
-
-        if (diagramKey.equals("thumbnail") || diagramKey.endsWith("-thumbnail")) {
-            try {
-                InputStreamAndContentLength inputStreamAndContentLength = workspaceComponent.getImage(workspaceId, diagramKey + ".png");
-                if (inputStreamAndContentLength != null) {
-                    return new InputStreamResource(inputStreamAndContentLength.getInputStream()) {
-                        @Override
-                        public long contentLength() {
-                            return inputStreamAndContentLength.getContentLength();
-                        }
-                    };
-                }
-            } catch (Exception e) {
-                log.error(e);
+        filename = HtmlUtils.filterHtml(filename);
+
+        try {
+            Image image = workspaceComponent.getImage(workspaceId, filename);
+            if (image != null) {
+                return ResponseEntity
+                        .ok()
+                        .contentType(MediaType.valueOf(image.getContentType()))
+                        .body(new InputStreamResource(image.getInputStream()) {
+                            @Override
+                            public long contentLength() {
+                                return image.getContentLength();
+                            }
+                        });
             }
+        } catch (Exception e) {
+            log.error(e);
         }
 
-
         response.setStatus(404);
         return null;
     }