OnCommittedResponseWrapper triggers commit when content length is reached

Fixes #26

spring-session/src/main/java/org/springframework/session/web/http/OnCommittedResponseWrapper.java

@@ -34,6 +34,17 @@ abstract class OnCommittedResponseWrapper extends HttpServletResponseWrapper {
 
     private boolean disableOnCommitted;
 
+    /**
+     * The Content-Length response header. If this is greater than 0, then once {@link #contentWritten} is larger than
+     * or equal the response is considered committed.
+     */
+    private long contentLength;
+
+    /**
+     * The size of data written to the response body.
+     */
+    private long contentWritten;
+
     /**
      * @param response the response to be wrapped
      */
@@ -41,6 +52,25 @@ abstract class OnCommittedResponseWrapper extends HttpServletResponseWrapper {
         super(response);
     }
 
+    @Override
+    public void addHeader(String name, String value) {
+        if("Content-Length".equalsIgnoreCase(name)) {
+            setContentLength(Long.parseLong(value));
+        }
+        super.addHeader(name, value);
+    }
+
+    @Override
+    public void setContentLength(int len) {
+        setContentLength((long) len);
+        super.setContentLength(len);
+    }
+
+    private void setContentLength(long len) {
+        this.contentLength = len;
+        checkContentLength(0);
+    }
+
     /**
      * Invoke this method to disable invoking {@link OnCommittedResponseWrapper#onResponseCommitted()} when the {@link javax.servlet.http.HttpServletResponse} is
      * committed. This can be useful in the event that Async Web Requests are
@@ -113,6 +143,55 @@ abstract class OnCommittedResponseWrapper extends HttpServletResponseWrapper {
         super.flushBuffer();
     }
 
+    private void trackContentLength(boolean content) {
+        checkContentLength(content ? 4 : 5); // TODO Localization
+    }
+
+    private void trackContentLength(char content) {
+        checkContentLength(1);
+    }
+
+    private void trackContentLength(Object content) {
+        trackContentLength(String.valueOf(content));
+    }
+
+    private void trackContentLength(char[] content) {
+        checkContentLength(content == null ? 0 : content.length);
+    }
+
+    private void trackContentLength(int content) {
+        trackContentLength(String.valueOf(content));
+    }
+
+    private void trackContentLength(float content) {
+        trackContentLength(String.valueOf(content));
+    }
+
+    private void trackContentLength(double content) {
+        trackContentLength(String.valueOf(content));
+    }
+
+    private void trackContentLengthLn() {
+        trackContentLength("\r\n");
+    }
+
+    private void trackContentLength(String content) {
+        checkContentLength(content.length());
+    }
+
+    /**
+     * Adds the contentLengthToWrite to the total contentWritten size and checks to see if the response should be
+     * written.
+     *
+     * @param contentLengthToWrite the size of the content that is about to be written.
+     */
+    private void checkContentLength(long contentLengthToWrite) {
+        contentWritten += contentLengthToWrite;
+        if(contentLength > 0  && contentWritten >= contentLength) {
+            doOnResponseCommitted();
+        }
+    }
+
     /**
      * Calls <code>onResponseCommmitted()</code> with the current contents as long as
      * {@link #disableOnResponseCommitted()()} was not invoked.
@@ -166,98 +245,131 @@ abstract class OnCommittedResponseWrapper extends HttpServletResponseWrapper {
         }
 
         public void write(int c) {
+            trackContentLength(c);
             delegate.write(c);
         }
 
         public void write(char[] buf, int off, int len) {
+            checkContentLength(len);
             delegate.write(buf, off, len);
         }
 
         public void write(char[] buf) {
+            trackContentLength(buf);
             delegate.write(buf);
         }
 
         public void write(String s, int off, int len) {
+            checkContentLength(len);
             delegate.write(s, off, len);
         }
 
         public void write(String s) {
+            trackContentLength(s);
             delegate.write(s);
         }
 
         public void print(boolean b) {
+            trackContentLength(b);
             delegate.print(b);
         }
 
         public void print(char c) {
+            trackContentLength(c);
             delegate.print(c);
         }
 
         public void print(int i) {
+            trackContentLength(i);
             delegate.print(i);
         }
 
         public void print(long l) {
+            trackContentLength(l);
             delegate.print(l);
         }
 
         public void print(float f) {
+            trackContentLength(f);
             delegate.print(f);
         }
 
         public void print(double d) {
+            trackContentLength(d);
             delegate.print(d);
         }
 
         public void print(char[] s) {
+            trackContentLength(s);
             delegate.print(s);
         }
 
         public void print(String s) {
+            trackContentLength(s);
             delegate.print(s);
         }
 
         public void print(Object obj) {
+            trackContentLength(obj);
             delegate.print(obj);
         }
 
         public void println() {
+            trackContentLengthLn();
             delegate.println();
         }
 
         public void println(boolean x) {
+            trackContentLength(x);
+            trackContentLengthLn();
             delegate.println(x);
         }
 
         public void println(char x) {
+            trackContentLength(x);
+            trackContentLengthLn();
             delegate.println(x);
         }
 
         public void println(int x) {
+            trackContentLength(x);
+            trackContentLengthLn();
             delegate.println(x);
         }
 
         public void println(long x) {
+            trackContentLength(x);
+            trackContentLengthLn();
             delegate.println(x);
         }
 
         public void println(float x) {
+            trackContentLength(x);
+            trackContentLengthLn();
             delegate.println(x);
         }
 
         public void println(double x) {
+            trackContentLength(x);
+            trackContentLengthLn();
             delegate.println(x);
         }
 
         public void println(char[] x) {
+            trackContentLength(x);
+            trackContentLengthLn();
             delegate.println(x);
         }
 
         public void println(String x) {
+            trackContentLength(x);
+            trackContentLengthLn();
             delegate.println(x);
         }
 
         public void println(Object x) {
+            trackContentLength(x);
+            trackContentLengthLn();
             delegate.println(x);
         }
 
@@ -278,14 +390,17 @@ abstract class OnCommittedResponseWrapper extends HttpServletResponseWrapper {
         }
 
         public PrintWriter append(CharSequence csq) {
+            checkContentLength(csq.length());
             return delegate.append(csq);
         }
 
         public PrintWriter append(CharSequence csq, int start, int end) {
+            checkContentLength(end - start);
             return delegate.append(csq, start, end);
         }
 
         public PrintWriter append(char c) {
+            trackContentLength(c);
             return delegate.append(c);
         }
     }
@@ -305,6 +420,7 @@ abstract class OnCommittedResponseWrapper extends HttpServletResponseWrapper {
         }
 
         public void write(int b) throws IOException {
+            trackContentLength(b);
             this.delegate.write(b);
         }
 
@@ -327,70 +443,94 @@ abstract class OnCommittedResponseWrapper extends HttpServletResponseWrapper {
         }
 
         public void print(boolean b) throws IOException {
+            trackContentLength(b);
             delegate.print(b);
         }
 
         public void print(char c) throws IOException {
+            trackContentLength(c);
             delegate.print(c);
         }
 
         public void print(double d) throws IOException {
+            trackContentLength(d);
             delegate.print(d);
         }
 
         public void print(float f) throws IOException {
+            trackContentLength(f);
             delegate.print(f);
         }
 
         public void print(int i) throws IOException {
+            trackContentLength(i);
             delegate.print(i);
         }
 
         public void print(long l) throws IOException {
+            trackContentLength(l);
             delegate.print(l);
         }
 
-        public void print(String arg0) throws IOException {
-            delegate.print(arg0);
+        public void print(String s) throws IOException {
+            trackContentLength(s);
+            delegate.print(s);
         }
 
         public void println() throws IOException {
+            trackContentLengthLn();
             delegate.println();
         }
 
         public void println(boolean b) throws IOException {
+            trackContentLength(b);
+            trackContentLengthLn();
             delegate.println(b);
         }
 
         public void println(char c) throws IOException {
+            trackContentLength(c);
+            trackContentLengthLn();
             delegate.println(c);
         }
 
         public void println(double d) throws IOException {
+            trackContentLength(d);
+            trackContentLengthLn();
             delegate.println(d);
         }
 
         public void println(float f) throws IOException {
+            trackContentLength(f);
+            trackContentLengthLn();
             delegate.println(f);
         }
 
         public void println(int i) throws IOException {
+            trackContentLength(i);
+            trackContentLengthLn();
             delegate.println(i);
         }
 
         public void println(long l) throws IOException {
+            trackContentLength(l);
+            trackContentLengthLn();
             delegate.println(l);
         }
 
         public void println(String s) throws IOException {
+            trackContentLength(s);
+            trackContentLengthLn();
             delegate.println(s);
         }
 
         public void write(byte[] b) throws IOException {
+            trackContentLength(b);
             delegate.write(b);
         }
 
         public void write(byte[] b, int off, int len) throws IOException {
+            checkContentLength(len);
             delegate.write(b, off, len);
         }
 

spring-session/src/test/java/org/springframework/session/web/http/OnCommittedResponseWrapperTests.java

@@ -1,5 +1,6 @@
 package org.springframework.session.web.http;
 
+import java.io.IOException;
 import java.io.PrintWriter;
 import java.util.Locale;
 
@@ -8,7 +9,6 @@ import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.runners.MockitoJUnitRunner;
-import org.springframework.session.web.http.OnCommittedResponseWrapper;
 
 import javax.servlet.ServletOutputStream;
 import javax.servlet.http.HttpServletResponse;
@@ -19,6 +19,8 @@ import static org.mockito.Mockito.when;
 
 @RunWith(MockitoJUnitRunner.class)
 public class OnCommittedResponseWrapperTests {
+    private static final String NL = "\r\n";
+
     @Mock
     HttpServletResponse delegate;
     @Mock
@@ -522,4 +524,539 @@ public class OnCommittedResponseWrapperTests {
 
         verify(out).println(x);
     }
+
+    // The amount of content specified in the setContentLength method of the response
+    // has been greater than zero and has been written to the response.
+
+    @Test
+    public void contentLengthPrintWriterWriteIntCommits() throws Exception {
+        int expected = 1;
+        response.setContentLength(String.valueOf(expected).length());
+
+        response.getWriter().write(expected);
+
+        assertThat(committed).isTrue();
+    }
+
+    @Test
+    public void contentLengthPrintWriterWriteIntMultiDigitCommits() throws Exception {
+        int expected = 10000;
+        response.setContentLength(String.valueOf(expected).length());
+
+        response.getWriter().write(expected);
+
+        assertThat(committed).isTrue();
+    }
+
+    @Test
+    public void contentLengthPlus1PrintWriterWriteIntMultiDigitCommits() throws Exception {
+        int expected = 10000;
+        response.setContentLength(String.valueOf(expected).length() + 1);
+
+        response.getWriter().write(expected);
+
+        assertThat(committed).isFalse();
+
+        response.getWriter().write(1);
+
+        assertThat(committed).isTrue();
+    }
+
+
+    @Test
+    public void contentLengthPrintWriterWriteCharIntIntCommits() throws Exception {
+        char[] buff = new char[0];
+        int off = 2;
+        int len = 3;
+        response.setContentLength(3);
+
+        response.getWriter().write(buff,off,len);
+
+        assertThat(committed).isTrue();
+    }
+
+    @Test
+    public void contentLengthPrintWriterWriteCharCommits() throws Exception {
+        char[] buff = new char[4];
+        response.setContentLength(buff.length);
+
+        response.getWriter().write(buff);
+
+        assertThat(committed).isTrue();
+    }
+
+    @Test
+    public void contentLengthPrintWriterWriteStringIntIntCommits() throws Exception {
+        String s = "";
+        int off = 2;
+        int len = 3;
+        response.setContentLength(3);
+
+        response.getWriter().write(s,off,len);
+
+        assertThat(committed).isTrue();
+    }
+
+
+    @Test
+    public void contentLengthPrintWriterWriteStringCommits() throws IOException {
+        String body = "something";
+        response.setContentLength(body.length());
+
+        response.getWriter().write(body);
+
+        assertThat(committed).isTrue();
+    }
+
+    @Test
+    public void printWriterWriteStringContentLengthCommits() throws IOException {
+        String body = "something";
+        response.getWriter().write(body);
+
+        response.setContentLength(body.length());
+
+        assertThat(committed).isTrue();
+    }
+
+    @Test
+    public void printWriterWriteStringDoesNotCommit() throws IOException {
+        String body = "something";
+
+        response.getWriter().write(body);
+
+        assertThat(committed).isFalse();
+    }
+
+    @Test
+    public void contentLengthPrintWriterPrintBooleanCommits() throws Exception {
+        boolean b = true;
+        response.setContentLength(1);
+
+        response.getWriter().print(b);
+
+        assertThat(committed).isTrue();
+    }
+
+    @Test
+    public void contentLengthPrintWriterPrintCharCommits() throws Exception {
+        char c = 1;
+        response.setContentLength(1);
+
+        response.getWriter().print(c);
+
+        assertThat(committed).isTrue();
+    }
+
+    @Test
+    public void contentLengthPrintWriterPrintIntCommits() throws Exception {
+        int i = 1234;
+        response.setContentLength(String.valueOf(i).length());
+
+        response.getWriter().print(i);
+
+        assertThat(committed).isTrue();
+    }
+
+    @Test
+    public void contentLengthPrintWriterPrintLongCommits() throws Exception {
+        long l = 12345;
+        response.setContentLength(String.valueOf(l).length());
+
+        response.getWriter().print(l);
+
+        assertThat(committed).isTrue();
+    }
+
+    @Test
+    public void contentLengthPrintWriterPrintFloatCommits() throws Exception {
+        float f = 12345;
+        response.setContentLength(String.valueOf(f).length());
+
+        response.getWriter().print(f);
+
+        assertThat(committed).isTrue();
+    }
+
+    @Test
+    public void contentLengthPrintWriterPrintDoubleCommits() throws Exception {
+        double x = 1.2345;
+        response.setContentLength(String.valueOf(x).length());
+
+        response.getWriter().print(x);
+
+        assertThat(committed).isTrue();
+    }
+
+    @Test
+    public void contentLengthPrintWriterPrintCharArrayCommits() throws Exception {
+        char[] x = new char[10];
+        response.setContentLength(x.length);
+
+        response.getWriter().print(x);
+
+        assertThat(committed).isTrue();
+    }
+
+    @Test
+    public void contentLengthPrintWriterPrintStringCommits() throws Exception {
+        String x = "12345";
+        response.setContentLength(x.length());
+
+        response.getWriter().print(x);
+
+        assertThat(committed).isTrue();
+    }
+
+    @Test
+    public void contentLengthPrintWriterPrintObjectCommits() throws Exception {
+        Object x = "12345";
+        response.setContentLength(String.valueOf(x).length());
+
+        response.getWriter().print(x);
+
+        assertThat(committed).isTrue();
+    }
+
+    @Test
+    public void contentLengthPrintWriterPrintlnCommits() throws Exception {
+        response.setContentLength(NL.length());
+
+        response.getWriter().println();
+
+        assertThat(committed).isTrue();
+    }
+
+    @Test
+    public void contentLengthPrintWriterPrintlnBooleanCommits() throws Exception {
+        boolean b = true;
+        response.setContentLength(1);
+
+        response.getWriter().println(b);
+
+        assertThat(committed).isTrue();
+    }
+
+    @Test
+    public void contentLengthPrintWriterPrintlnCharCommits() throws Exception {
+        char c = 1;
+        response.setContentLength(1);
+
+        response.getWriter().println(c);
+
+        assertThat(committed).isTrue();
+    }
+
+    @Test
+    public void contentLengthPrintWriterPrintlnIntCommits() throws Exception {
+        int i = 12345;
+        response.setContentLength(String.valueOf(i).length());
+
+        response.getWriter().println(i);
+
+        assertThat(committed).isTrue();
+    }
+
+    @Test
+    public void contentLengthPrintWriterPrintlnLongCommits() throws Exception {
+        long l = 12345678;
+        response.setContentLength(String.valueOf(l).length());
+
+        response.getWriter().println(l);
+
+        assertThat(committed).isTrue();
+    }
+
+    @Test
+    public void contentLengthPrintWriterPrintlnFloatCommits() throws Exception {
+        float f = 1234;
+        response.setContentLength(String.valueOf(f).length());
+
+        response.getWriter().println(f);
+
+        assertThat(committed).isTrue();
+    }
+
+    @Test
+    public void contentLengthPrintWriterPrintlnDoubleCommits() throws Exception {
+        double x = 1;
+        response.setContentLength(String.valueOf(x).length());
+
+        response.getWriter().println(x);
+
+        assertThat(committed).isTrue();
+    }
+
+    @Test
+    public void contentLengthPrintWriterPrintlnCharArrayCommits() throws Exception {
+        char[] x = new char[20];
+        response.setContentLength(x.length);
+
+        response.getWriter().println(x);
+
+        assertThat(committed).isTrue();
+    }
+
+    @Test
+    public void contentLengthPrintWriterPrintlnStringCommits() throws Exception {
+        String x = "1";
+        response.setContentLength(String.valueOf(x).length());
+
+        response.getWriter().println(x);
+
+        assertThat(committed).isTrue();
+    }
+
+    @Test
+    public void contentLengthPrintWriterPrintlnObjectCommits() throws Exception {
+        Object x = "1";
+        response.setContentLength(String.valueOf(x).length());
+
+        response.getWriter().println(x);
+
+        assertThat(committed).isTrue();
+    }
+
+    @Test
+    public void contentLengthPrintWriterAppendCharSequenceCommits() throws Exception {
+        String x = "a";
+        response.setContentLength(String.valueOf(x).length());
+
+        response.getWriter().append(x);
+
+        assertThat(committed).isTrue();
+    }
+
+    @Test
+    public void contentLengthPrintWriterAppendCharSequenceIntIntCommits() throws Exception {
+        String x = "abcdef";
+        int start = 1;
+        int end = 3;
+        response.setContentLength(end - start);
+
+        response.getWriter().append(x, start, end);
+
+        assertThat(committed).isTrue();
+    }
+
+    @Test
+    public void contentLengthPrintWriterAppendCharCommits() throws Exception {
+        char x = 1;
+        response.setContentLength(1);
+
+        response.getWriter().append(x);
+
+        assertThat(committed).isTrue();
+    }
+
+    @Test
+    public void contentLengthOutputStreamWriteIntCommits() throws Exception {
+        int expected = 1;
+        response.setContentLength(String.valueOf(expected).length());
+
+        response.getOutputStream().write(expected);
+
+        assertThat(committed).isTrue();
+    }
+
+    @Test
+    public void contentLengthOutputStreamWriteIntMultiDigitCommits() throws Exception {
+        int expected = 10000;
+        response.setContentLength(String.valueOf(expected).length());
+
+        response.getOutputStream().write(expected);
+
+        assertThat(committed).isTrue();
+    }
+
+    @Test
+    public void contentLengthPlus1OutputStreamWriteIntMultiDigitCommits() throws Exception {
+        int expected = 10000;
+        response.setContentLength(String.valueOf(expected).length() + 1);
+
+        response.getOutputStream().write(expected);
+
+        assertThat(committed).isFalse();
+
+        response.getOutputStream().write(1);
+
+        assertThat(committed).isTrue();
+    }
+
+    @Test
+    public void contentLengthOutputStreamPrintBooleanCommits() throws Exception {
+        boolean b = true;
+        response.setContentLength(1);
+
+        response.getOutputStream().print(b);
+
+        assertThat(committed).isTrue();
+    }
+
+    @Test
+    public void contentLengthOutputStreamPrintCharCommits() throws Exception {
+        char c = 1;
+        response.setContentLength(1);
+
+        response.getOutputStream().print(c);
+
+        assertThat(committed).isTrue();
+    }
+
+    @Test
+    public void contentLengthOutputStreamPrintIntCommits() throws Exception {
+        int i = 1234;
+        response.setContentLength(String.valueOf(i).length());
+
+        response.getOutputStream().print(i);
+
+        assertThat(committed).isTrue();
+    }
+
+    @Test
+    public void contentLengthOutputStreamPrintLongCommits() throws Exception {
+        long l = 12345;
+        response.setContentLength(String.valueOf(l).length());
+
+        response.getOutputStream().print(l);
+
+        assertThat(committed).isTrue();
+    }
+
+    @Test
+    public void contentLengthOutputStreamPrintFloatCommits() throws Exception {
+        float f = 12345;
+        response.setContentLength(String.valueOf(f).length());
+
+        response.getOutputStream().print(f);
+
+        assertThat(committed).isTrue();
+    }
+
+    @Test
+    public void contentLengthOutputStreamPrintDoubleCommits() throws Exception {
+        double x = 1.2345;
+        response.setContentLength(String.valueOf(x).length());
+
+        response.getOutputStream().print(x);
+
+        assertThat(committed).isTrue();
+    }
+
+    @Test
+    public void contentLengthOutputStreamPrintStringCommits() throws Exception {
+        String x = "12345";
+        response.setContentLength(x.length());
+
+        response.getOutputStream().print(x);
+
+        assertThat(committed).isTrue();
+    }
+
+    @Test
+    public void contentLengthOutputStreamPrintlnCommits() throws Exception {
+        response.setContentLength(NL.length());
+
+        response.getOutputStream().println();
+
+        assertThat(committed).isTrue();
+    }
+
+    @Test
+    public void contentLengthOutputStreamPrintlnBooleanCommits() throws Exception {
+        boolean b = true;
+        response.setContentLength(1);
+
+        response.getOutputStream().println(b);
+
+        assertThat(committed).isTrue();
+    }
+
+    @Test
+    public void contentLengthOutputStreamPrintlnCharCommits() throws Exception {
+        char c = 1;
+        response.setContentLength(1);
+
+        response.getOutputStream().println(c);
+
+        assertThat(committed).isTrue();
+    }
+
+    @Test
+    public void contentLengthOutputStreamPrintlnIntCommits() throws Exception {
+        int i = 12345;
+        response.setContentLength(String.valueOf(i).length());
+
+        response.getOutputStream().println(i);
+
+        assertThat(committed).isTrue();
+    }
+
+    @Test
+    public void contentLengthOutputStreamPrintlnLongCommits() throws Exception {
+        long l = 12345678;
+        response.setContentLength(String.valueOf(l).length());
+
+        response.getOutputStream().println(l);
+
+        assertThat(committed).isTrue();
+    }
+
+    @Test
+    public void contentLengthOutputStreamPrintlnFloatCommits() throws Exception {
+        float f = 1234;
+        response.setContentLength(String.valueOf(f).length());
+
+        response.getOutputStream().println(f);
+
+        assertThat(committed).isTrue();
+    }
+
+    @Test
+    public void contentLengthOutputStreamPrintlnDoubleCommits() throws Exception {
+        double x = 1;
+        response.setContentLength(String.valueOf(x).length());
+
+        response.getOutputStream().println(x);
+
+        assertThat(committed).isTrue();
+    }
+
+    @Test
+    public void contentLengthOutputStreamPrintlnStringCommits() throws Exception {
+        String x = "1";
+        response.setContentLength(String.valueOf(x).length());
+
+        response.getOutputStream().println(x);
+
+        assertThat(committed).isTrue();
+    }
+
+    @Test
+    public void contentLengthDoesNotCommit() throws IOException {
+        String body = "something";
+
+        response.setContentLength(body.length());
+
+        assertThat(committed).isFalse();
+    }
+
+    @Test
+    public void contentLengthOutputStreamWriteStringCommits() throws IOException {
+        String body = "something";
+        response.setContentLength(body.length());
+
+        response.getOutputStream().print(body);
+
+        assertThat(committed).isTrue();
+    }
+
+    @Test
+    public void addHeaderContentLengthPrintWriterWriteStringCommits() throws Exception {
+        int expected = 1234;
+        response.addHeader("Content-Length",String.valueOf(String.valueOf(expected).length()));
+
+        response.getWriter().write(expected);
+
+        assertThat(committed).isTrue();
+    }
 }