From 97d3b26c8b22cc8e18927ca278c79913b71e1688 Mon Sep 17 00:00:00 2001 From: Alexander Yastrebov Date: Tue, 2 May 2017 18:46:48 +0200 Subject: [PATCH] GH-146 Delegates getOutputStream and getWriter after withoutBody was called --- .../logbook/servlet/LocalResponse.java | 126 +++++++++++------- .../logbook/servlet/LocalResponseTest.java | 98 ++++++++++++-- 2 files changed, 167 insertions(+), 57 deletions(-) diff --git a/logbook-servlet/src/main/java/org/zalando/logbook/servlet/LocalResponse.java b/logbook-servlet/src/main/java/org/zalando/logbook/servlet/LocalResponse.java index 7818f9e2a..cf7722983 100644 --- a/logbook-servlet/src/main/java/org/zalando/logbook/servlet/LocalResponse.java +++ b/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) { + 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())); - } - } - } diff --git a/logbook-servlet/src/test/java/org/zalando/logbook/servlet/LocalResponseTest.java b/logbook-servlet/src/test/java/org/zalando/logbook/servlet/LocalResponseTest.java index e2838b666..e0e7b3935 100644 --- a/logbook-servlet/src/test/java/org/zalando/logbook/servlet/LocalResponseTest.java +++ b/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(); + } + }