io.github.gilbertoowen.request.logging.wrapper.SpringResponseWrapper Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of request-logging-spring-boot-starter Show documentation

HTTP logging with Spring Boot

There is a newer version: 1.4.1.RELEASE

Show newest version

package io.github.gilbertoowen.request.logging.wrapper;

import javax.servlet.ServletOutputStream;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;
import java.io.IOException;
import java.io.OutputStreamWriter;
import java.io.PrintWriter;
import java.util.HashMap;
import java.util.Map;

public class SpringResponseWrapper extends HttpServletResponseWrapper {

	private ServletOutputStream outputStream;
	private PrintWriter writer;
	private ServletOutputStreamWrapper copier;

	public SpringResponseWrapper(HttpServletResponse response) throws IOException {
		super(response);
	}

	@Override
	public ServletOutputStream getOutputStream() throws IOException {
		if (writer != null) {
			throw new IllegalStateException("getWriter() has already been called on this response.");
		}

		if (outputStream == null) {
			outputStream = getResponse().getOutputStream();
			copier = new ServletOutputStreamWrapper(outputStream);
		}

		return copier;
	}

	@Override
	public PrintWriter getWriter() throws IOException {
		if (outputStream != null) {
			throw new IllegalStateException("getOutputStream() has already been called on this response.");
		}

		if (writer == null) {
			copier = new ServletOutputStreamWrapper(getResponse().getOutputStream());
			writer = new PrintWriter(new OutputStreamWriter(copier, getResponse().getCharacterEncoding()), true);
		}

		return writer;
	}

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

	public byte[] getContentAsByteArray() {
		if (copier != null) {
			return copier.getCopy();
		}
		else {
			return new byte[0];
		}
	}

	public Map getAllHeaders() {
		final Map headers = new HashMap<>();
		getHeaderNames().forEach(it -> headers.put(it, getHeader(it)));
		return  headers;
	}

}