From 4c845a5dcdae09b1ffe4fa96d5422c828b10bfc2 Mon Sep 17 00:00:00 2001
From: Maxime Sinclair <maxime.sinclair@gmail.com>
Date: Mon, 2 Dec 2013 12:45:11 +0100
Subject: [PATCH] [FEATURE] Proxy returns the correct Content-type

---
 .../plantuml/servlet/ProxyServlet.java        | 24 ++++++++++++++++---
 .../plantuml/servlet/TestProxy.java           |  2 +-
 2 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/src/main/java/net/sourceforge/plantuml/servlet/ProxyServlet.java b/src/main/java/net/sourceforge/plantuml/servlet/ProxyServlet.java
index 566be37..ff4b590 100644
--- a/src/main/java/net/sourceforge/plantuml/servlet/ProxyServlet.java
+++ b/src/main/java/net/sourceforge/plantuml/servlet/ProxyServlet.java
@@ -35,12 +35,16 @@ import javax.servlet.http.HttpServlet;
 import javax.servlet.http.HttpServletRequest;
 import javax.servlet.http.HttpServletResponse;
 
+import net.sourceforge.plantuml.BlockUml;
 import net.sourceforge.plantuml.FileFormat;
-import net.sourceforge.plantuml.FileFormatOption;
 import net.sourceforge.plantuml.SourceStringReader;
+import net.sourceforge.plantuml.core.Diagram;
+import net.sourceforge.plantuml.core.UmlSource;
 
 import java.security.cert.Certificate;
+import java.util.List;
 
+import javax.imageio.IIOException;
 import javax.net.ssl.HttpsURLConnection;
 import javax.net.ssl.SSLPeerUnverifiedException;
 
@@ -73,8 +77,22 @@ public class ProxyServlet extends HttpServlet {
         System.out.println("getSource=>" + diagmarkup);
         SourceStringReader reader = new SourceStringReader(diagmarkup);
         int n = index == null ? 0 : Integer.parseInt(index);
-
-        reader.generateImage(response.getOutputStream(), n, new FileFormatOption(getOutputFormat(), false));
+        List<BlockUml> blocks = reader.getBlocks();
+        BlockUml block = blocks.get(n);
+        Diagram diagram = block.getDiagram();
+        UmlSource umlSrc = diagram.getSource();
+        String uml = umlSrc.getPlainString();
+        System.out.println("uml="+uml);
+        
+        // generate the response
+        DiagramResponse dr = new DiagramResponse(response, getOutputFormat());
+        try {
+            dr.sendDiagram(uml);
+        } catch (IIOException iioe) {
+            // Browser has closed the connection, so the HTTP OutputStream is closed
+            // Silently catch the exception to avoid annoying log 
+        }
+        dr = null;        
     }
 
     private String getSource(URL url) throws IOException {
diff --git a/src/test/java/net/sourceforge/plantuml/servlet/TestProxy.java b/src/test/java/net/sourceforge/plantuml/servlet/TestProxy.java
index e1804ec..1640723 100644
--- a/src/test/java/net/sourceforge/plantuml/servlet/TestProxy.java
+++ b/src/test/java/net/sourceforge/plantuml/servlet/TestProxy.java
@@ -19,7 +19,7 @@ public class TestProxy extends WebappTestCase {
         WebResponse response = conversation.getResource(request);
         // Analyze response
         // Verifies the Content-Type header
-        // assertEquals( "Response content type is not PNG", "image/png", response.getContentType());
+        assertEquals( "Response content type is not PNG", "image/png", response.getContentType());
         // Get the image and verify its size (~2000 bytes)
         InputStream responseStream = response.getInputStream();
         ByteArrayOutputStream imageStream = new ByteArrayOutputStream();