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