com.catchpoint.trace.instrument.integrations.spring.web.advice.CapturingHttpServletResponse Maven / Gradle / Ivy
package com.catchpoint.trace.instrument.integrations.spring.web.advice;
import org.springframework.util.FastByteArrayOutputStream;
import org.springframework.web.util.WebUtils;
import javax.servlet.ServletOutputStream;
import javax.servlet.WriteListener;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;
import java.io.IOException;
import java.io.OutputStreamWriter;
import java.io.PrintWriter;
import java.io.UnsupportedEncodingException;
/**
* @author serkan
*/
public class CapturingHttpServletResponse extends HttpServletResponseWrapper {
private final int maxSize;
private final FastByteArrayOutputStream capturedContent;
private ServletOutputStream outputStream;
private PrintWriter writer;
private boolean overflowed = false;
public CapturingHttpServletResponse(HttpServletResponse response, int maxSize) {
super(response);
this.maxSize = maxSize;
this.capturedContent = new FastByteArrayOutputStream(maxSize);
}
@Override
public ServletOutputStream getOutputStream() throws IOException {
if (outputStream == null) {
outputStream = new ResponseServletOutputStream(getResponse().getOutputStream());
}
return outputStream;
}
@Override
public PrintWriter getWriter() throws IOException {
if (writer == null) {
ServletOutputStream sos = getOutputStream();
String characterEncoding = getCharacterEncoding();
writer = characterEncoding != null
? new ResponsePrintWriter(sos, characterEncoding)
: new ResponsePrintWriter(sos, WebUtils.DEFAULT_CHARACTER_ENCODING);
}
return writer;
}
public boolean isOverflowed() {
return overflowed;
}
public byte[] getCapturedResponse() {
if (overflowed) {
return null;
}
return capturedContent.toByteArray();
}
private class ResponseServletOutputStream extends ServletOutputStream {
private final ServletOutputStream os;
public ResponseServletOutputStream(ServletOutputStream os) {
this.os = os;
}
private boolean willOverflow(int count) {
return capturedContent.size() + count > maxSize;
}
@Override
public void write(int b) throws IOException {
os.write(b);
if (!overflowed) {
if (willOverflow(1)) {
overflowed = true;
capturedContent.reset();
return;
}
try {
capturedContent.write(b);
} catch (IndexOutOfBoundsException e) {
overflowed = true;
capturedContent.reset();
}
}
}
@Override
public void write(byte[] b, int off, int len) throws IOException {
os.write(b, off, len);
if (!overflowed) {
if (willOverflow(len)) {
overflowed = true;
capturedContent.reset();
return;
}
try {
capturedContent.write(b, off, len);
} catch (IndexOutOfBoundsException e) {
overflowed = true;
capturedContent.reset();
}
}
}
@Override
public boolean isReady() {
return os.isReady();
}
@Override
public void setWriteListener(WriteListener writeListener) {
os.setWriteListener(writeListener);
}
}
private class ResponsePrintWriter extends PrintWriter {
public ResponsePrintWriter(ServletOutputStream servletOutputStream,
String characterEncoding) throws UnsupportedEncodingException {
super(new OutputStreamWriter(servletOutputStream, characterEncoding));
}
@Override
public void write(char[] buf, int off, int len) {
super.write(buf, off, len);
super.flush();
}
@Override
public void write(String s, int off, int len) {
super.write(s, off, len);
super.flush();
}
@Override
public void write(int c) {
super.write(c);
super.flush();
}
}
}