001/**
002 * Logback: the reliable, generic, fast and flexible logging framework.
003 * Copyright (C) 1999-2015, QOS.ch. All rights reserved.
004 *
005 * This program and the accompanying materials are dual-licensed under
006 * either the terms of the Eclipse Public License v1.0 as published by
007 * the Eclipse Foundation
008 *
009 *   or (per the licensee's choosing)
010 *
011 * under the terms of the GNU Lesser General Public License version 2.1
012 * as published by the Free Software Foundation.
013 */
014package ch.qos.logback.access.servlet;
015
016import java.io.IOException;
017import java.io.OutputStreamWriter;
018import java.io.PrintWriter;
019
020import javax.servlet.ServletOutputStream;
021import javax.servlet.http.HttpServletResponse;
022import javax.servlet.http.HttpServletResponseWrapper;
023
024public class TeeHttpServletResponse extends HttpServletResponseWrapper {
025
026    TeeServletOutputStream teeServletOutputStream;
027    PrintWriter teeWriter;
028
029    public TeeHttpServletResponse(HttpServletResponse httpServletResponse) {
030        super(httpServletResponse);
031    }
032
033    @Override
034    public ServletOutputStream getOutputStream() throws IOException {
035        if (teeServletOutputStream == null) {
036            teeServletOutputStream = new TeeServletOutputStream(this.getResponse());
037        }
038        return teeServletOutputStream;
039    }
040
041    @Override
042    public PrintWriter getWriter() throws IOException {
043        if (this.teeWriter == null) {
044            this.teeWriter = new PrintWriter(
045                    new OutputStreamWriter(getOutputStream(), this.getResponse().getCharacterEncoding()), true);
046        }
047        return this.teeWriter;
048    }
049
050    @Override
051    public void flushBuffer() {
052        if (this.teeWriter != null) {
053            this.teeWriter.flush();
054        }
055    }
056
057    byte[] getOutputBuffer() {
058        // teeServletOutputStream can be null if the getOutputStream method is never
059        // called.
060        if (teeServletOutputStream != null) {
061            return teeServletOutputStream.getOutputStreamAsByteArray();
062        } else {
063            return null;
064        }
065    }
066
067    void finish() throws IOException {
068        if (this.teeWriter != null) {
069            this.teeWriter.close();
070        }
071        if (this.teeServletOutputStream != null) {
072            this.teeServletOutputStream.close();
073        }
074    }
075}