From 294b2a241829855455f6fea3dca03a7d49e9d262 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?K=C3=A1lm=C3=A1n=20Vince?=
 <30904009+DRKV333@users.noreply.github.com>
Date: Tue, 20 Jun 2023 19:54:35 +0200
Subject: [PATCH] Add diagram indexing support for all formats

---
 .../plantuml/servlet/DiagramResponse.java     | 51 ++++++++++++++++---
 1 file changed, 45 insertions(+), 6 deletions(-)

diff --git a/src/main/java/net/sourceforge/plantuml/servlet/DiagramResponse.java b/src/main/java/net/sourceforge/plantuml/servlet/DiagramResponse.java
index e019192..05a63e3 100644
--- a/src/main/java/net/sourceforge/plantuml/servlet/DiagramResponse.java
+++ b/src/main/java/net/sourceforge/plantuml/servlet/DiagramResponse.java
@@ -29,6 +29,7 @@ import java.io.FileReader;
 import java.io.IOException;
 import java.io.PrintWriter;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.List;
 
 import jakarta.servlet.http.HttpServletRequest;
@@ -57,6 +58,16 @@ import net.sourceforge.plantuml.version.Version;
  */
 public class DiagramResponse {
 
+    private static class BlockSelection {
+        private final BlockUml block;
+        private final int systemIdx;
+
+        BlockSelection(BlockUml blk, int idx) {
+            block = blk;
+            systemIdx = idx;
+        }
+    }
+
     /**
      * X-Powered-By http header value included in every response by default.
      */
@@ -163,20 +174,48 @@ public class DiagramResponse {
             response.getOutputStream().write(encodedBytes.getBytes());
             return;
         }
-        final BlockUml blockUml = reader.getBlocks().get(0);
-        if (notModified(blockUml)) {
-            addHeaderForCache(blockUml);
+
+        final BlockSelection blockSelection = getOutputBlockSelection(reader, idx);
+        if (blockSelection == null) {
+            response.sendError(HttpServletResponse.SC_BAD_REQUEST);
+            return;
+        }
+
+        if (notModified(blockSelection.block)) {
+            addHeaderForCache(blockSelection.block);
             response.sendError(HttpServletResponse.SC_NOT_MODIFIED);
             return;
         }
         if (StringUtils.isDiagramCacheable(uml)) {
-            addHeaderForCache(blockUml);
+            addHeaderForCache(blockSelection.block);
         }
-        final Diagram diagram = blockUml.getDiagram();
+        final Diagram diagram = blockSelection.block.getDiagram();
         if (diagram instanceof PSystemError) {
             response.setStatus(HttpServletResponse.SC_BAD_REQUEST);
         }
-        diagram.exportDiagram(response.getOutputStream(), idx, new FileFormatOption(format));
+        diagram.exportDiagram(response.getOutputStream(), blockSelection.systemIdx, new FileFormatOption(format));
+    }
+
+    private BlockSelection getOutputBlockSelection(SourceStringReader reader, int numImage) {
+        if (numImage < 0) {
+            return null;
+        }
+
+        Collection<BlockUml> blocks = reader.getBlocks();
+        if (blocks.isEmpty()) {
+            return null;
+        }
+
+        for (BlockUml b : blocks) {
+            final Diagram system = b.getDiagram();
+            final int nbInSystem = system.getNbImages();
+            if (numImage < nbInSystem) {
+                return new BlockSelection(b, numImage);
+            }
+            numImage -= nbInSystem;
+        }
+
+        return null;
     }
 
     /**