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(new OutputStreamWriter(getOutputStream(), this.getResponse().getCharacterEncoding()), true);
045        }
046        return this.teeWriter;
047    }
048
049    @Override
050    public void flushBuffer() {
051        if (this.teeWriter != null) {
052            this.teeWriter.flush();
053        }
054    }
055
056    byte[] getOutputBuffer() {
057        // teeServletOutputStream can be null if the getOutputStream method is never
058        // called.
059        if (teeServletOutputStream != null) {
060            return teeServletOutputStream.getOutputStreamAsByteArray();
061        } else {
062            return null;
063        }
064    }
065
066    void finish() throws IOException {
067        if (this.teeWriter != null) {
068            this.teeWriter.close();
069        }
070        if (this.teeServletOutputStream != null) {
071            this.teeServletOutputStream.close();
072        }
073    }
074}