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 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