ResponseCapturingWrapper.java

/* 
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */

package org.esigate.servlet.impl;

import java.io.IOException;
import java.io.PrintWriter;
import java.io.Writer;
import java.util.Locale;

import javax.servlet.ServletOutputStream;
import javax.servlet.http.Cookie;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;

import org.apache.commons.io.output.ByteArrayOutputStream;
import org.apache.commons.io.output.StringBuilderWriter;
import org.apache.http.Header;
import org.apache.http.HttpHeaders;
import org.apache.http.HttpStatus;
import org.apache.http.HttpVersion;
import org.apache.http.client.methods.CloseableHttpResponse;
import org.apache.http.entity.ByteArrayEntity;
import org.apache.http.entity.ContentType;
import org.apache.http.entity.StringEntity;
import org.apache.http.message.BasicHttpResponse;
import org.apache.http.message.BasicStatusLine;
import org.esigate.Parameters;
import org.esigate.http.BasicCloseableHttpResponse;
import org.esigate.http.ContentTypeHelper;
import org.esigate.http.DateUtils;
import org.esigate.http.HttpResponseUtils;
import org.esigate.http.IncomingRequest;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Wrapper to the HttpServletResponse that intercepts the content written in order to build an
 * {@link org.apache.http.HttpResponse}.
 * <ul>
 * <li>If the content of the response is required for transformation (parseable content-type or proxy=false) or smaller
 * than the buffer size the {@link org.apache.http.HttpResponse} will contain the entire response</li>
 * <li>If the content of the response is not required for transformation, the {@link org.apache.http.HttpResponse} will
 * contain only an abstract of the response truncated to the bufer size. The complete response will have already been
 * written to the original {@link HttpServletResponse}</li>
 * </ul>
 * 
 * @author Francois-Xavier Bonnet
 * 
 */
public class ResponseCapturingWrapper extends HttpServletResponseWrapper {
    private static final Logger LOG = LoggerFactory.getLogger(ResponseCapturingWrapper.class);

    // OutputStream and Writer exposed
    private ServletOutputStream outputStream;
    private PrintWriter writer;

    // OutputStream and Writer of the wrapped HttpServletResponse
    private ServletOutputStream responseOutputStream;
    private PrintWriter responseWriter;

    // OutputStream and Writer buffers
    private ByteArrayOutputStream internalOutputStream;
    private StringBuilderWriter internalWriter;

    private HttpServletResponse response;

    private CloseableHttpResponse httpClientResponse = BasicCloseableHttpResponse.adapt(new BasicHttpResponse(
            new BasicStatusLine(HttpVersion.HTTP_1_1, HttpStatus.SC_OK, "OK")));

    private String contentType;
    private String characterEncoding;
    private int bufferSize;
    private int bytesWritten = 0;
    private boolean committed = false;
    private ContentTypeHelper contentTypeHelper;
    private final boolean proxy;
    private boolean capture = true;
    private final ResponseSender responseSender;
    private final IncomingRequest incomingRequest;

    public ResponseCapturingWrapper(HttpServletResponse response, ContentTypeHelper contentTypeHelper, boolean proxy,
            int bufferSize, ResponseSender responseSender, IncomingRequest incomingRequest) {
        super(response);
        this.response = response;
        this.bufferSize = bufferSize;
        this.contentTypeHelper = contentTypeHelper;
        this.proxy = proxy;
        this.responseSender = responseSender;
        this.incomingRequest = incomingRequest;
    }

    @Override
    public void setStatus(int sc) {
        httpClientResponse.setStatusLine(new BasicStatusLine(HttpVersion.HTTP_1_1, sc, ""));
    }

    @Override
    public void setStatus(int sc, String sm) {
        httpClientResponse.setStatusLine(new BasicStatusLine(HttpVersion.HTTP_1_1, sc, sm));
    }

    public int getStatus() {
        return httpClientResponse.getStatusLine().getStatusCode();
    }

    @Override
    public void sendError(int sc, String msg) {
        httpClientResponse.setStatusLine(new BasicStatusLine(HttpVersion.HTTP_1_1, sc, msg));
    }

    @Override
    public void sendError(int sc) {
        httpClientResponse.setStatusLine(new BasicStatusLine(HttpVersion.HTTP_1_1, sc, ""));
    }

    @Override
    public void sendRedirect(String location) {
        httpClientResponse.setStatusLine(new BasicStatusLine(HttpVersion.HTTP_1_1, HttpStatus.SC_MOVED_TEMPORARILY,
                "Temporary redirect"));
        httpClientResponse.setHeader(HttpHeaders.LOCATION, location);
    }

    @Override
    public boolean containsHeader(String name) {
        return httpClientResponse.containsHeader(name);
    }

    @Override
    public void setHeader(String name, String value) {
        if (HttpHeaders.CONTENT_TYPE.equalsIgnoreCase(name)) {
            setContentType(value);
        } else {
            httpClientResponse.setHeader(name, value);
        }
    }

    @Override
    public void addHeader(String name, String value) {
        if (HttpHeaders.CONTENT_TYPE.equalsIgnoreCase(name)) {
            setContentType(value);
        } else {
            httpClientResponse.addHeader(name, value);
        }
    }

    @Override
    public void setDateHeader(String name, long date) {
        setHeader(name, DateUtils.formatDate(date));
    }

    @Override
    public void addDateHeader(String name, long date) {
        addHeader(name, DateUtils.formatDate(date));
    }

    @Override
    public void setIntHeader(String name, int value) {
        setHeader(name, Integer.toString(value));
    }

    @Override
    public void addIntHeader(String name, int value) {
        addHeader(name, Integer.toString(value));
    }

    @Override
    public void setContentLength(int len) {
        setHeader(HttpHeaders.CONTENT_LENGTH, Integer.toString(len));
    }

    @Override
    public void setCharacterEncoding(String charset) {
        this.characterEncoding = charset;
        updateContentTypeHeader();
    }

    @Override
    public String getCharacterEncoding() {
        return this.characterEncoding;
    }

    @Override
    public String getContentType() {
        Header contentTypeHeader = httpClientResponse.getFirstHeader(HttpHeaders.CONTENT_TYPE);
        if (contentTypeHeader != null) {
            return contentTypeHeader.getValue();
        } else {
            return null;
        }
    }

    @Override
    public void setContentType(String type) {
        ContentType parsedContentType = ContentType.parse(type);
        this.contentType = parsedContentType.getMimeType();
        if (parsedContentType.getCharset() != null) {
            this.characterEncoding = parsedContentType.getCharset().name();
        }
        updateContentTypeHeader();
    }

    private void updateContentTypeHeader() {
        if (contentType != null) {
            if (characterEncoding == null) {
                httpClientResponse.setHeader(HttpHeaders.CONTENT_TYPE, contentType);
            } else {
                httpClientResponse.setHeader(HttpHeaders.CONTENT_TYPE, contentType + ";charset=" + characterEncoding);
            }
        }
    }

    @Override
    public void setLocale(Locale loc) {
        response.setLocale(loc);
        if (characterEncoding == null) {
            characterEncoding = response.getCharacterEncoding();
            updateContentTypeHeader();
        }
    }

    @Override
    public Locale getLocale() {
        return response.getLocale();
    }

    @Override
    public void setBufferSize(int size) {
        this.bufferSize = size;
    }

    @Override
    public int getBufferSize() {
        return bufferSize;
    }

    @Override
    public void addCookie(Cookie cookie) {
        response.addCookie(cookie);
    }

    @Override
    public String encodeURL(String url) {
        return response.encodeURL(url);
    }

    @Override
    public String encodeRedirectURL(String url) {
        return response.encodeRedirectURL(url);
    }

    @SuppressWarnings("deprecation")
    @Override
    public String encodeUrl(String url) {
        return response.encodeUrl(url);
    }

    @SuppressWarnings("deprecation")
    @Override
    public String encodeRedirectUrl(String url) {
        return response.encodeRedirectUrl(url);
    }

    @Override
    public void reset() {
        if (isCommitted()) {
            throw new IllegalStateException("Response is already committed");
        }
        httpClientResponse =
                BasicCloseableHttpResponse.adapt(new BasicHttpResponse(new BasicStatusLine(HttpVersion.HTTP_1_1,
                        HttpStatus.SC_OK, "OK")));
    }

    @Override
    public ServletOutputStream getOutputStream() {
        LOG.debug("getOutputStream");
        if (writer != null) {
            throw new IllegalStateException("Writer already obtained");
        }
        if (outputStream == null) {
            internalOutputStream = new ByteArrayOutputStream(Parameters.DEFAULT_BUFFER_SIZE);
            outputStream = new ServletOutputStream() {

                @Override
                public void write(int b) throws IOException {
                    if (capture || bytesWritten < bufferSize) {
                        internalOutputStream.write(b);
                    } else {
                        responseOutputStream.write(b);
                    }
                    bytesWritten++;
                    if (bytesWritten == bufferSize) {
                        commit();
                    }
                }

                @Override
                public void flush() throws IOException {
                    commit();
                }

                @Override
                public void close() throws IOException {
                    commit();
                }

                private void commit() throws IOException {
                    if (!committed) {
                        capture = hasToCaptureOutput();
                        if (!capture) {
                            responseSender.sendHeaders(httpClientResponse, incomingRequest, response);
                            responseOutputStream = response.getOutputStream();
                            internalOutputStream.writeTo(responseOutputStream);
                        }
                        committed = true;
                    }

                }

            };
        }
        return outputStream;
    }

    @Override
    public PrintWriter getWriter() {
        LOG.debug("getWriter");
        if (outputStream != null) {
            throw new IllegalStateException("OutputStream already obtained");
        }
        if (writer == null) {
            internalWriter = new StringBuilderWriter(Parameters.DEFAULT_BUFFER_SIZE);
            writer = new PrintWriter(new Writer() {

                @Override
                public void write(char[] cbuf, int off, int len) throws IOException {
                    if (capture || bytesWritten < bufferSize) {
                        internalWriter.write(cbuf, off, len);
                    } else {
                        responseWriter.write(cbuf, off, len);
                    }
                    bytesWritten++;
                    if (bytesWritten == bufferSize) {
                        commit();
                    }
                }

                @Override
                public void flush() throws IOException {
                    commit();
                }

                @Override
                public void close() throws IOException {
                    commit();
                }

                private void commit() throws IOException {
                    if (!committed) {
                        capture = hasToCaptureOutput();
                        if (!capture) {
                            responseSender.sendHeaders(httpClientResponse, incomingRequest, response);
                            responseWriter = response.getWriter();
                            responseWriter.write(internalWriter.toString());
                        }
                        committed = true;
                    }

                }

            });
        }
        return writer;
    }

    @Override
    public void flushBuffer() throws IOException {
        if (outputStream != null) {
            outputStream.flush();
        }
        if (writer != null) {
            writer.flush();
        }
    }

    @Override
    public void resetBuffer() {
        if (isCommitted()) {
            throw new IllegalStateException("Response is already committed");
        }
        if (internalOutputStream != null) {
            internalOutputStream.reset();
        }
        if (internalWriter != null) {
            internalWriter = new StringBuilderWriter(Parameters.DEFAULT_BUFFER_SIZE);
        }
        bytesWritten = 0;
    }

    @Override
    public boolean isCommitted() {
        return committed;
    }

    /**
     * We have to capture the output of the response in 2 cases :
     * <ol>
     * <li>the content type is text and may have to be transformed</li>
     * <li>we are inside an include, in this case the content type must be considered as text</li>
     * </ol>
     * 
     * @return true if we have to capture the output of the response
     */
    private boolean hasToCaptureOutput() {
        return !proxy || contentTypeHelper.isTextContentType(httpClientResponse)
                || HttpResponseUtils.getFirstHeader(HttpHeaders.CONTENT_TYPE, httpClientResponse) == null;
    }

    /**
     * Returns the response. If the response has not been captured and has been written directly to the
     * {@link HttpServletResponse}, calling this method closes the HttpServletResponse writer OutputStream
     * 
     * @return the response
     */
    public CloseableHttpResponse getCloseableHttpResponse() {
        ContentType resultContentType = null;
        if (this.contentType != null) {
            resultContentType = ContentType.create(this.contentType, characterEncoding);
        }
        if (internalWriter != null) {
            writer.flush();
            httpClientResponse.setEntity(new StringEntity(internalWriter.toString(), resultContentType));
        } else if (internalOutputStream != null) {
            try {
                outputStream.flush();
            } catch (IOException e) {
                // Nothing to do;
            }
            httpClientResponse.setEntity(new ByteArrayEntity(internalOutputStream.toByteArray(), resultContentType));
        }
        if (!capture) {
            // The result has already been written to the response, let's close
            // the response stream
            if (responseWriter != null) {
                responseWriter.close();
            }
            if (responseOutputStream != null) {
                try {
                    responseOutputStream.close();
                } catch (IOException e) {
                    LOG.warn("Could not close servlet output stream: " + e.getMessage());
                }
            }
        }
        return httpClientResponse;
    }

}