001/*
002 * Copyright 2002-2015 the original author or authors.
003 *
004 * Licensed under the Apache License, Version 2.0 (the "License");
005 * you may not use this file except in compliance with the License.
006 * You may obtain a copy of the License at
007 *
008 *      https://www.apache.org/licenses/LICENSE-2.0
009 *
010 * Unless required by applicable law or agreed to in writing, software
011 * distributed under the License is distributed on an "AS IS" BASIS,
012 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013 * See the License for the specific language governing permissions and
014 * limitations under the License.
015 */
016
017package org.springframework.web.util;
018
019import java.io.IOException;
020import java.io.InputStream;
021import java.io.OutputStreamWriter;
022import java.io.PrintWriter;
023import java.io.UnsupportedEncodingException;
024
025import javax.servlet.ServletOutputStream;
026import javax.servlet.http.HttpServletResponse;
027import javax.servlet.http.HttpServletResponseWrapper;
028
029import org.springframework.http.HttpHeaders;
030import org.springframework.util.ClassUtils;
031import org.springframework.util.FastByteArrayOutputStream;
032
033/**
034 * {@link javax.servlet.http.HttpServletResponse} wrapper that caches all content written to
035 * the {@linkplain #getOutputStream() output stream} and {@linkplain #getWriter() writer},
036 * and allows this content to be retrieved via a {@link #getContentAsByteArray() byte array}.
037 *
038 * <p>Used e.g. by {@link org.springframework.web.filter.ShallowEtagHeaderFilter}.
039 *
040 * @author Juergen Hoeller
041 * @since 4.1.3
042 * @see ContentCachingRequestWrapper
043 */
044public class ContentCachingResponseWrapper extends HttpServletResponseWrapper {
045
046        /** Checking for Servlet 3.0+ HttpServletResponse.getHeader(String) */
047        private static final boolean servlet3Present =
048                        ClassUtils.hasMethod(HttpServletResponse.class, "getHeader", String.class);
049
050
051        private final FastByteArrayOutputStream content = new FastByteArrayOutputStream(1024);
052
053        private final ServletOutputStream outputStream = new ResponseServletOutputStream();
054
055        private PrintWriter writer;
056
057        private int statusCode = HttpServletResponse.SC_OK;
058
059        private Integer contentLength;
060
061
062        /**
063         * Create a new ContentCachingResponseWrapper for the given servlet response.
064         * @param response the original servlet response
065         */
066        public ContentCachingResponseWrapper(HttpServletResponse response) {
067                super(response);
068        }
069
070
071        @Override
072        public void setStatus(int sc) {
073                super.setStatus(sc);
074                this.statusCode = sc;
075        }
076
077        @SuppressWarnings("deprecation")
078        @Override
079        public void setStatus(int sc, String sm) {
080                super.setStatus(sc, sm);
081                this.statusCode = sc;
082        }
083
084        @Override
085        public void sendError(int sc) throws IOException {
086                copyBodyToResponse(false);
087                try {
088                        super.sendError(sc);
089                }
090                catch (IllegalStateException ex) {
091                        // Possibly on Tomcat when called too late: fall back to silent setStatus
092                        super.setStatus(sc);
093                }
094                this.statusCode = sc;
095        }
096
097        @Override
098        @SuppressWarnings("deprecation")
099        public void sendError(int sc, String msg) throws IOException {
100                copyBodyToResponse(false);
101                try {
102                        super.sendError(sc, msg);
103                }
104                catch (IllegalStateException ex) {
105                        // Possibly on Tomcat when called too late: fall back to silent setStatus
106                        super.setStatus(sc, msg);
107                }
108                this.statusCode = sc;
109        }
110
111        @Override
112        public void sendRedirect(String location) throws IOException {
113                copyBodyToResponse(false);
114                super.sendRedirect(location);
115        }
116
117        @Override
118        public ServletOutputStream getOutputStream() throws IOException {
119                return this.outputStream;
120        }
121
122        @Override
123        public PrintWriter getWriter() throws IOException {
124                if (this.writer == null) {
125                        String characterEncoding = getCharacterEncoding();
126                        this.writer = (characterEncoding != null ? new ResponsePrintWriter(characterEncoding) :
127                                        new ResponsePrintWriter(WebUtils.DEFAULT_CHARACTER_ENCODING));
128                }
129                return this.writer;
130        }
131
132        @Override
133        public void flushBuffer() throws IOException {
134                // do not flush the underlying response as the content as not been copied to it yet
135        }
136
137        @Override
138        public void setContentLength(int len) {
139                if (len > this.content.size()) {
140                        this.content.resize(len);
141                }
142                this.contentLength = len;
143        }
144
145        // Overrides Servlet 3.1 setContentLengthLong(long) at runtime
146        public void setContentLengthLong(long len) {
147                if (len > Integer.MAX_VALUE) {
148                        throw new IllegalArgumentException("Content-Length exceeds ContentCachingResponseWrapper's maximum (" +
149                                        Integer.MAX_VALUE + "): " + len);
150                }
151                int lenInt = (int) len;
152                if (lenInt > this.content.size()) {
153                        this.content.resize(lenInt);
154                }
155                this.contentLength = lenInt;
156        }
157
158        @Override
159        public void setBufferSize(int size) {
160                if (size > this.content.size()) {
161                        this.content.resize(size);
162                }
163        }
164
165        @Override
166        public void resetBuffer() {
167                this.content.reset();
168        }
169
170        @Override
171        public void reset() {
172                super.reset();
173                this.content.reset();
174        }
175
176        /**
177         * Return the status code as specified on the response.
178         */
179        public int getStatusCode() {
180                return this.statusCode;
181        }
182
183        /**
184         * Return the cached response content as a byte array.
185         */
186        public byte[] getContentAsByteArray() {
187                return this.content.toByteArray();
188        }
189
190        /**
191         * Return an {@link InputStream} to the cached content.
192         * @since 4.2
193         */
194        public InputStream getContentInputStream() {
195                return this.content.getInputStream();
196        }
197
198        /**
199         * Return the current size of the cached content.
200         * @since 4.2
201         */
202        public int getContentSize() {
203                return this.content.size();
204        }
205
206        /**
207         * Copy the complete cached body content to the response.
208         * @since 4.2
209         */
210        public void copyBodyToResponse() throws IOException {
211                copyBodyToResponse(true);
212        }
213
214        /**
215         * Copy the cached body content to the response.
216         * @param complete whether to set a corresponding content length
217         * for the complete cached body content
218         * @since 4.2
219         */
220        protected void copyBodyToResponse(boolean complete) throws IOException {
221                if (this.content.size() > 0) {
222                        HttpServletResponse rawResponse = (HttpServletResponse) getResponse();
223                        if ((complete || this.contentLength != null) && !rawResponse.isCommitted()) {
224                                boolean hasTransferEncoding = false;
225                                if (servlet3Present) {
226                                        hasTransferEncoding = (rawResponse.getHeader(HttpHeaders.TRANSFER_ENCODING) != null);
227                                }
228                                if (!hasTransferEncoding) {
229                                        rawResponse.setContentLength(complete ? this.content.size() : this.contentLength);
230                                }
231                                this.contentLength = null;
232                        }
233                        this.content.writeTo(rawResponse.getOutputStream());
234                        this.content.reset();
235                        if (complete) {
236                                super.flushBuffer();
237                        }
238                }
239        }
240
241
242        private class ResponseServletOutputStream extends ServletOutputStream {
243
244                @Override
245                public void write(int b) throws IOException {
246                        content.write(b);
247                }
248
249                @Override
250                public void write(byte[] b, int off, int len) throws IOException {
251                        content.write(b, off, len);
252                }
253        }
254
255
256        private class ResponsePrintWriter extends PrintWriter {
257
258                public ResponsePrintWriter(String characterEncoding) throws UnsupportedEncodingException {
259                        super(new OutputStreamWriter(content, characterEncoding));
260                }
261
262                @Override
263                public void write(char[] buf, int off, int len) {
264                        super.write(buf, off, len);
265                        super.flush();
266                }
267
268                @Override
269                public void write(String s, int off, int len) {
270                        super.write(s, off, len);
271                        super.flush();
272                }
273
274                @Override
275                public void write(int c) {
276                        super.write(c);
277                        super.flush();
278                }
279        }
280
281}