From 5859c3b4271291c003f8170a2491debdaa512e0f Mon Sep 17 00:00:00 2001
From: HeinrichAD <heinrichad@users.noreply.github.com>
Date: Wed, 19 Jul 2023 22:22:35 +0200
Subject: [PATCH] support diagrams without startuml, enduml

Add support for diagram without @startuml and @enduml.
---
 .../plantuml/servlet/DiagramResponse.java     | 29 ++++++++++++++++---
 .../plantuml/servlet/ProxyServlet.java        | 21 ++++----------
 2 files changed, 30 insertions(+), 20 deletions(-)

diff --git a/src/main/java/net/sourceforge/plantuml/servlet/DiagramResponse.java b/src/main/java/net/sourceforge/plantuml/servlet/DiagramResponse.java
index 05a63e3..c034c1b 100644
--- a/src/main/java/net/sourceforge/plantuml/servlet/DiagramResponse.java
+++ b/src/main/java/net/sourceforge/plantuml/servlet/DiagramResponse.java
@@ -157,10 +157,10 @@ public class DiagramResponse {
         response.addHeader("Access-Control-Allow-Origin", "*");
         response.setContentType(getContentType());
 
-        final Defines defines = getPreProcDefines();
-        SourceStringReader reader = new SourceStringReader(defines, uml, CONFIG);
-        if (CONFIG.size() > 0 && reader.getBlocks().get(0).getDiagram().getWarningOrError() != null) {
-            reader = new SourceStringReader(uml);
+        final SourceStringReader reader = getSourceStringReader(uml);
+        if (reader == null) {
+            response.sendError(HttpServletResponse.SC_BAD_REQUEST, "No UML diagram found");
+            return;
         }
 
         if (format == FileFormat.BASE64) {
@@ -218,6 +218,27 @@ public class DiagramResponse {
         return null;
     }
 
+    private SourceStringReader getSourceStringReader(String uml) {
+        SourceStringReader reader = getSourceStringReaderWithConfig(uml);
+        if (reader.getBlocks().isEmpty()) {
+            uml = "@startuml\n" + uml + "\n@enduml";
+            reader = getSourceStringReaderWithConfig(uml);
+            if (reader.getBlocks().isEmpty()) {
+                return null;
+            }
+        }
+        return reader;
+    }
+
+    private SourceStringReader getSourceStringReaderWithConfig(String uml) {
+        final Defines defines = getPreProcDefines();
+        SourceStringReader reader = new SourceStringReader(defines, uml, CONFIG);
+        if (!CONFIG.isEmpty() && reader.getBlocks().get(0).getDiagram().getWarningOrError() != null) {
+            reader = new SourceStringReader(defines, uml);
+        }
+        return reader;
+    }
+
     /**
      * Get PlantUML preprocessor defines.
      *
diff --git a/src/main/java/net/sourceforge/plantuml/servlet/ProxyServlet.java b/src/main/java/net/sourceforge/plantuml/servlet/ProxyServlet.java
index 345d042..04766cf 100644
--- a/src/main/java/net/sourceforge/plantuml/servlet/ProxyServlet.java
+++ b/src/main/java/net/sourceforge/plantuml/servlet/ProxyServlet.java
@@ -29,7 +29,6 @@ import java.io.InputStreamReader;
 import java.net.HttpURLConnection;
 import java.net.MalformedURLException;
 import java.net.URL;
-import java.util.List;
 import java.util.stream.Collectors;
 
 import javax.imageio.IIOException;
@@ -39,11 +38,7 @@ import jakarta.servlet.http.HttpServlet;
 import jakarta.servlet.http.HttpServletRequest;
 import jakarta.servlet.http.HttpServletResponse;
 
-import net.sourceforge.plantuml.BlockUml;
 import net.sourceforge.plantuml.FileFormat;
-import net.sourceforge.plantuml.SourceStringReader;
-import net.sourceforge.plantuml.core.Diagram;
-import net.sourceforge.plantuml.core.UmlSource;
 
 /**
  * Proxy servlet of the webapp.
@@ -108,29 +103,23 @@ public class ProxyServlet extends HttpServlet {
         final String source = request.getParameter("src");
         final String index = request.getParameter("idx");
 
+        final int idx = index == null ? 0 : Integer.parseInt(index);
         final URL srcUrl = validateURL(source, response);
         if (srcUrl == null) {
             return;  // error is already set/handled inside `validateURL`
         }
 
-        // generate the response
-        String diagmarkup = getSource(srcUrl);
-        SourceStringReader reader = new SourceStringReader(diagmarkup);
-        int n = index == null ? 0 : Integer.parseInt(index);
-        List<BlockUml> blocks = reader.getBlocks();
-        BlockUml block = blocks.get(n);
-        Diagram diagram = block.getDiagram();
-        UmlSource umlSrc = diagram.getSource();
-        String uml = umlSrc.getPlainString("\n");
+        // fetch diagram from URL
+        final String uml = getSource(srcUrl);
 
         // generate the response
         DiagramResponse dr = new DiagramResponse(response, getOutputFormat(fmt), request);
         try {
             // special handling for the MAP since it's not using "#sendDiagram()" like the other types
             if ("map".equals(fmt)) {
-                dr.sendMap(uml, 0);
+                dr.sendMap(uml, idx);
             } else {
-                dr.sendDiagram(uml, 0);
+                dr.sendDiagram(uml, idx);
             }
         } catch (IIOException e) {
             // Browser has closed the connection, so the HTTP OutputStream is closed