Write JSON response to a buffer instead of directly to the output stream (#11424, #11156)

* Allows safe abortion of JSON output and writing a critical notification instead
* Unifies json writing for init request and other UIDL requests
* Ensures headers are written before the response

Change-Id: Idd8acb672aac8716b727701d6c057bbe58f50993

2 files changed, 52 insertions, 33 deletions

diff --git a/server/src/com/vaadin/server/communication/UIInitHandler.java b/server/src/com/vaadin/server/communication/UIInitHandler.java
index 0e18209986..c3e7119d3f 100644
--- a/server/src/com/vaadin/server/communication/UIInitHandler.java
+++ b/server/src/com/vaadin/server/communication/UIInitHandler.java
@@ -16,11 +16,8 @@
 
 package com.vaadin.server.communication;
 
-import java.io.BufferedWriter;
 import java.io.IOException;
-import java.io.OutputStream;
 import java.io.OutputStreamWriter;
-import java.io.PrintWriter;
 import java.io.StringWriter;
 import java.util.List;
 import java.util.Map;
@@ -62,12 +59,7 @@ public abstract class UIInitHandler extends SynchronizedRequestHandler {
             return false;
         }
 
-        // NOTE! GateIn requires, for some weird reason, getOutputStream
-        // to be used instead of getWriter() (it seems to interpret
-        // application/json as a binary content type)
-        final OutputStream out = response.getOutputStream();
-        final PrintWriter outWriter = new PrintWriter(new BufferedWriter(
-                new OutputStreamWriter(out, "UTF-8")));
+        StringWriter stringWriter = new StringWriter();
 
         try {
             assert UI.getCurrent() == null;
@@ -75,8 +67,6 @@ public abstract class UIInitHandler extends SynchronizedRequestHandler {
             // Set browser information from the request
             session.getBrowser().updateRequestDetails(request);
 
-            response.setContentType("application/json; charset=UTF-8");
-
             UI uI = getBrowserDetailsUI(request, session);
 
             session.getCommunicationManager().repaintAll(uI);
@@ -86,15 +76,51 @@ public abstract class UIInitHandler extends SynchronizedRequestHandler {
             String initialUIDL = getInitialUidl(request, uI);
             params.put("uidl", initialUIDL);
 
-            outWriter.write(params.toString());
-            // NOTE GateIn requires the buffers to be flushed to work
-            outWriter.flush();
-            out.flush();
+            stringWriter.write(params.toString());
         } catch (JSONException e) {
-            // TODO PUSH handle
-            e.printStackTrace();
+            throw new IOException("Error producing initial UIDL", e);
+        } finally {
+            stringWriter.close();
+        }
+
+        return commitJsonResponse(request, response, stringWriter.toString());
+    }
+
+    /**
+     * Commit the JSON response. We can't write immediately to the output stream
+     * as we want to write only a critical notification if something goes wrong
+     * during the response handling.
+     * 
+     * @param request
+     *            The request that resulted in this response
+     * @param response
+     *            The response to write to
+     * @param json
+     *            The JSON to write
+     * @return true if the JSON was written successfully, false otherwise
+     * @throws IOException
+     *             If there was an exception while writing to the output
+     */
+    static boolean commitJsonResponse(VaadinRequest request,
+            VaadinResponse response, String json) throws IOException {
+        // The response was produced without errors so write it to the client
+        response.setContentType("application/json; charset=UTF-8");
+
+        // Ensure that the browser does not cache UIDL responses.
+        // iOS 6 Safari requires this (#9732)
+        response.setHeader("Cache-Control", "no-cache");
+
+        // NOTE! GateIn requires, for some weird reason, getOutputStream
+        // to be used instead of getWriter() (it seems to interpret
+        // application/json as a binary content type)
+        OutputStreamWriter outputWriter = new OutputStreamWriter(
+                response.getOutputStream(), "UTF-8");
+        try {
+            outputWriter.write(json);
+            // NOTE GateIn requires the buffers to be flushed to work
+            outputWriter.flush();
         } finally {
-            outWriter.close();
+            outputWriter.close();
         }
 
         return true;
diff --git a/server/src/com/vaadin/server/communication/UidlRequestHandler.java b/server/src/com/vaadin/server/communication/UidlRequestHandler.java
index e626bbe58d..0de9029063 100644
--- a/server/src/com/vaadin/server/communication/UidlRequestHandler.java
+++ b/server/src/com/vaadin/server/communication/UidlRequestHandler.java
@@ -16,10 +16,8 @@
 
 package com.vaadin.server.communication;
 
-import java.io.BufferedWriter;
 import java.io.IOException;
-import java.io.OutputStream;
-import java.io.OutputStreamWriter;
+import java.io.StringWriter;
 import java.io.Writer;
 import java.util.LinkedList;
 import java.util.logging.Level;
@@ -77,7 +75,6 @@ public class UidlRequestHandler extends SynchronizedRequestHandler {
         ClientConnector highlightedConnector;
         // repaint requested or session has timed out and new one is created
         boolean repaintAll;
-        final OutputStream out = response.getOutputStream();
 
         // TODO PUSH repaintAll, analyzeLayouts, highlightConnector should be
         // part of the message payload to make the functionality transport
@@ -101,8 +98,7 @@ public class UidlRequestHandler extends SynchronizedRequestHandler {
             }
         }
 
-        final Writer outWriter = new BufferedWriter(new OutputStreamWriter(out,
-                "UTF-8"));
+        StringWriter stringWriter = new StringWriter();
 
         try {
             rpcHandler.handleRpc(uI, request.getReader(), request);
@@ -111,7 +107,7 @@ public class UidlRequestHandler extends SynchronizedRequestHandler {
                 session.getCommunicationManager().repaintAll(uI);
             }
 
-            writeUidl(request, response, uI, outWriter, repaintAll,
+            writeUidl(request, response, uI, stringWriter, repaintAll,
                     analyzeLayouts);
             postHandleRequest(uI);
         } catch (JSONException e) {
@@ -119,6 +115,7 @@ public class UidlRequestHandler extends SynchronizedRequestHandler {
             // Refresh on client side
             criticalNotifier.criticalNotification(request, response, null,
                     null, null, null);
+            return true;
         } catch (InvalidUIDLSecurityKeyException e) {
             getLogger().log(Level.WARNING,
                     "Invalid security key received from {}",
@@ -126,16 +123,14 @@ public class UidlRequestHandler extends SynchronizedRequestHandler {
             // Refresh on client side
             criticalNotifier.criticalNotification(request, response, null,
                     null, null, null);
+            return true;
         } finally {
-            outWriter.close();
+            stringWriter.close();
             requestThemeName = null;
         }
 
-        // Ensure that the browser does not cache UIDL responses.
-        // iOS 6 Safari requires this (#9732)
-        response.setHeader("Cache-Control", "no-cache");
-
-        return true;
+        return UIInitHandler.commitJsonResponse(request, response,
+                stringWriter.toString());
     }
 
     /**
@@ -205,8 +200,6 @@ public class UidlRequestHandler extends SynchronizedRequestHandler {
      */
     protected void openJsonMessage(Writer outWriter, VaadinResponse response)
             throws IOException {
-        // Sets the response type
-        response.setContentType("application/json; charset=UTF-8");
         // some dirt to prevent cross site scripting
         outWriter.write("for(;;);[{");
     }