Merge pull request #154 from zalando/GH-146-fix-get-output-stream-alr…

…eady-called

GH-146 Delegates getOutputStream and getWriter after withoutBody was called

logbook-servlet/src/main/java/org/zalando/logbook/servlet/LocalResponse.java

@@ -1,16 +1,6 @@
 package org.zalando.logbook.servlet;
 
-import org.zalando.logbook.HttpResponse;
-import org.zalando.logbook.Origin;
-import org.zalando.logbook.RawHttpResponse;
-
-import javax.annotation.Nullable;
-import javax.servlet.ServletOutputStream;
-import javax.servlet.http.HttpServletResponse;
-import javax.servlet.http.HttpServletResponseWrapper;
 import java.io.ByteArrayOutputStream;
-import java.io.DataOutput;
-import java.io.DataOutputStream;
 import java.io.IOException;
 import java.io.OutputStream;
 import java.io.OutputStreamWriter;
@@ -19,28 +9,27 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.function.Supplier;
+import javax.servlet.ServletOutputStream;
+import javax.servlet.http.HttpServletResponse;
+import javax.servlet.http.HttpServletResponseWrapper;
+import org.zalando.logbook.HttpResponse;
+import org.zalando.logbook.Origin;
+import org.zalando.logbook.RawHttpResponse;
 
 import static java.nio.charset.StandardCharsets.UTF_8;
-import static org.zalando.logbook.servlet.NullOutputStream.NULL;
 
 final class LocalResponse extends HttpServletResponseWrapper implements RawHttpResponse, HttpResponse {
 
-    private ByteArrayOutputStream bytes = new ByteArrayOutputStream();
-    private DataOutput output = new DataOutputStream(bytes);
-
-    private ServletOutputStream stream = new TeeServletOutputStream();
-    private PrintWriter writer = new TeePrintWriter();
     private final String protocolVersion;
 
-    /**
-     * Null until we successfully intercepted it.
-     */
-    @Nullable
-    private byte[] body;
+    private boolean withBody;
+    private Tee tee;
 
     LocalResponse(final HttpServletResponse response, final String protocolVersion) throws IOException {
         super(response);
         this.protocolVersion = protocolVersion;
+        this.withBody = true;
     }
 
     @Override
@@ -76,74 +65,117 @@ public Charset getCharset() {
 
     @Override
     public HttpResponse withBody() {
+        assertWithBody();
         return this;
     }
 
     @Override
     public void withoutBody() throws IOException {
-        this.body = null;
-        this.bytes = new ByteArrayOutputStream(0);
-        this.output = new DataOutputStream(NULL);
-        this.stream = super.getOutputStream();
-        this.writer = super.getWriter();
+        assertWithBody();
+        withBody = false;
     }
 
     @Override
     public ServletOutputStream getOutputStream() throws IOException {
-        return stream;
+        if (withBody) {
+            return tee().getOutputStream();
+        } else {
+            return super.getOutputStream();
+        }
     }
 
     @Override
     public PrintWriter getWriter() throws IOException {
-        return writer;
+        if (withBody) {
+            return tee().getWriter(this::getCharset);
+        } else {
+            return super.getWriter();
+        }
     }
 
     @Override
-    public byte[] getBody() {
-        if (body == null) {
-            body = bytes.toByteArray();
+    public byte[] getBody() throws IOException {
+        assertWithBody();
+        return tee().getBytes();
+    }
+
+    private void assertWithBody() {
+        if (!withBody) {
+            throw new IllegalStateException("Response is without body");
+        }
+    }
+
+    private Tee tee() throws IOException {
+        if (tee == null) {
+            tee = new Tee(super.getOutputStream());
+        }
+        return tee;
+    }
+
+    private static class Tee {
+
+        private final ByteArrayOutputStream branch;
+        private final TeeServletOutputStream output;
+
+        private PrintWriter writer;
+        private byte[] bytes;
+
+        private Tee(final OutputStream original) throws IOException {
+            this.branch = new ByteArrayOutputStream();
+            this.output = new TeeServletOutputStream(original, branch);
+        }
+
+        ServletOutputStream getOutputStream() {
+            return output;
+        }
+
+        PrintWriter getWriter(final Supplier<Charset> charset) {
+            if (writer == null) {
+                writer = new PrintWriter(new OutputStreamWriter(output, charset.get()));
+            }
+            return writer;
+        }
+
+        byte[] getBytes() {
+            if (bytes == null || bytes.length != branch.size()) {
+                bytes = branch.toByteArray();
+            }
+            return bytes;
         }
-        return body;
     }
 
-    private final class TeeServletOutputStream extends ServletOutputStream {
+    private static class TeeServletOutputStream extends ServletOutputStream {
 
         private final OutputStream original;
+        private final OutputStream branch;
 
-        private TeeServletOutputStream() throws IOException {
-            this.original = LocalResponse.super.getOutputStream();
+        private TeeServletOutputStream(final OutputStream original, final OutputStream branch) {
+            this.original = original;
+            this.branch = branch;
         }
 
         @Override
         public void write(final int b) throws IOException {
-            output.write(b);
             original.write(b);
+            branch.write(b);
         }
 
         @Override
         public void write(final byte[] b, final int off, final int len) throws IOException {
-            output.write(b, off, len);
             original.write(b, off, len);
+            branch.write(b, off, len);
         }
 
         @Override
         public void flush() throws IOException {
             original.flush();
+            branch.flush();
         }
 
         @Override
         public void close() throws IOException {
             original.close();
+            branch.close();
         }
-
-    }
-
-    private final class TeePrintWriter extends PrintWriter {
-
-        public TeePrintWriter() {
-            super(new OutputStreamWriter(stream, getCharset()));
-        }
-
     }
-
 }

logbook-servlet/src/test/java/org/zalando/logbook/servlet/LocalResponseTest.java

@@ -1,32 +1,110 @@
 package org.zalando.logbook.servlet;
 
-import org.junit.Test;
-
+import java.io.IOException;
+import java.io.PrintWriter;
 import javax.servlet.ServletOutputStream;
 import javax.servlet.http.HttpServletResponse;
-import java.io.IOException;
+import org.junit.Before;
+import org.junit.Test;
 
+import static org.junit.Assert.assertNotSame;
 import static org.junit.Assert.assertSame;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
 import static org.mockito.Mockito.when;
 
 public class LocalResponseTest {
 
-    @Test
-    public void shouldUseSameBody() throws IOException {
-        final HttpServletResponse mock = mock(HttpServletResponse.class);
+    private HttpServletResponse mock;
+    private LocalResponse unit;
+
+    @Before
+    public void setUp() throws IOException {
+        mock = mock(HttpServletResponse.class);
         when(mock.getOutputStream()).thenReturn(new ServletOutputStream() {
             @Override
             public void write(final int b) throws IOException {
             }
         });
+        unit = new LocalResponse(mock, "1");
+    }
 
-        final LocalResponse response = new LocalResponse(mock, "1");
-        response.getOutputStream().write("test".getBytes());
+    @Test
+    public void shouldUseSameBody() throws IOException {
+        unit.getOutputStream().write("test".getBytes());
 
-        final byte[] body1 = response.getBody();
-        final byte[] body2 = response.getBody();
+        final byte[] body1 = unit.getBody();
+        final byte[] body2 = unit.getBody();
 
         assertSame(body1, body2);
     }
+
+    @Test
+    public void shouldUseDifferentBodyAfterWrite() throws IOException {
+        unit.getOutputStream().write("Hello".getBytes());
+        final byte[] body1 = unit.getBody();
+
+        unit.getOutputStream().write("World".getBytes());
+        final byte[] body2 = unit.getBody();
+
+        assertNotSame(body1, body2);
+    }
+
+    @Test
+    public void shouldTeeGetOutputStream() throws IOException {
+        unit.withBody();
+
+        final ServletOutputStream os1 = unit.getOutputStream();
+        final ServletOutputStream os2 = unit.getOutputStream();
+
+        assertSame(os1, os2);
+
+        verify(mock).getOutputStream();
+        verifyNoMoreInteractions(mock);
+    }
+
+    @Test
+    public void shouldDelegateGetOutputStream() throws IOException {
+        unit.withoutBody();
+
+        unit.getOutputStream();
+        unit.getOutputStream();
+
+        verify(mock, times(2)).getOutputStream();
+        verifyNoMoreInteractions(mock);
+    }
+
+    @Test
+    public void shouldTeeGetWriter() throws IOException {
+        unit.withBody();
+
+        final PrintWriter writer1 = unit.getWriter();
+        final PrintWriter writer2 = unit.getWriter();
+
+        assertSame(writer1, writer2);
+
+        verify(mock).getOutputStream();
+        verify(mock).getCharacterEncoding();
+        verifyNoMoreInteractions(mock);
+    }
+
+    @Test
+    public void shouldDelegateGetWriter() throws IOException {
+        unit.withoutBody();
+
+        unit.getWriter();
+        unit.getWriter();
+
+        verify(mock, times(2)).getWriter();
+        verifyNoMoreInteractions(mock);
+    }
+
+    @Test(expected = IllegalStateException.class)
+    public void shouldNotAllowWithBodyAfterWithoutBody() throws IOException {
+        unit.withoutBody();
+        unit.withBody();
+    }
+
 }