package com.nokia.tsl_data.util.logging.entity;

import java.io.ByteArrayOutputStream;
import java.io.IOException;

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

public class RepeatableHttpServletResponseWrapper extends HttpServletResponseWrapper {

    private final ByteArrayOutputStream byteArrayOutputStream;
    private final ServletOutputStream servletOutputStream;

    public RepeatableHttpServletResponseWrapper(HttpServletResponse response) throws IOException {
        super(response);
        // 从当前的response获取ServletOutputStream
        ServletOutputStream outputStream = response.getOutputStream();
        byteArrayOutputStream = new ByteArrayOutputStream();
        // 匿名内部类实现ServletOutputStream
        servletOutputStream = new ServletOutputStream() {

            @Override
            public boolean isReady() {
                return outputStream.isReady();
            }

            @Override
            public void setWriteListener(WriteListener writeListener) {
                outputStream.setWriteListener(writeListener);
            }

            @Override
            public void write(int b) throws IOException {
                // 写入原有的基础上同时写入ByteArrayOutputStream
                outputStream.write(b);
                byteArrayOutputStream.write(b);
            }
        };
    }

    @Override
    public ServletOutputStream getOutputStream() {
        return servletOutputStream;
    }

    public byte[] toByteArray() {
        return byteArrayOutputStream.toByteArray();
    }
}
