FetchAPI/src/main/java/com/minicircuits/BufferedResponseWrapper.java

package com.minicircuits;

import java.io.BufferedInputStream;
import java.io.BufferedOutputStream;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.io.OutputStreamWriter;
import java.io.PrintWriter;
import java.nio.file.Files;
import java.nio.file.Path;

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

public class BufferedResponseWrapper extends HttpServletResponseWrapper {
    private final Path tempFile;
    private final FileOutputStream fileOutputStream;
    private final ServletOutputStream servletOutputStream;
    private PrintWriter writer;

    public BufferedResponseWrapper(HttpServletResponse response) throws IOException {
        super(response);
        this.tempFile = Files.createTempFile("fetchapi_", ".tmp");
        this.fileOutputStream = new FileOutputStream(tempFile.toFile());
        this.servletOutputStream = new ServletOutputStream() {
            @Override
            public void write(int b) throws IOException {
                fileOutputStream.write(b);
            }

            @Override
            public void write(byte[] b, int off, int len) throws IOException {
                fileOutputStream.write(b, off, len);
            }

            @Override
            public void flush() throws IOException {
                fileOutputStream.flush();
            }

            @Override
            public void close() throws IOException {
                fileOutputStream.close();
            }

            @Override
            public boolean isReady() {
                return true;
            }

            @Override
            public void setWriteListener(WriteListener writeListener) {
                // Not used for blocking I/O
            }
        };
    }

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

    @Override
    public PrintWriter getWriter() throws IOException {
        if (writer == null) {
            String charset = getCharacterEncoding();
            writer = new PrintWriter(new OutputStreamWriter(fileOutputStream,
                charset != null ? charset : "UTF-8"));
        }
        return writer;
    }

    @Override
    public void flushBuffer() {
        // Don't flush to the real response — we're buffering
    }

    @Override
    public void reset() {
        // Don't reset the real response during buffering
    }

    @Override
    public void resetBuffer() {
        // Don't reset during buffering
    }

    /**
     * Flush any remaining data, then stream the temp file to the real response.
     * Sets Content-Length. Deletes the temp file when done.
     */
    public void writeToResponse() throws IOException {
        // Flush and close writer/stream
        if (writer != null) {
            writer.close();
        } else {
            fileOutputStream.close();
        }

        // Set content length from temp file size
        long fileSize = Files.size(tempFile);
        if (fileSize <= Integer.MAX_VALUE) {
            getResponse().setContentLength((int) fileSize);
        } else {
            ((HttpServletResponse) getResponse()).setHeader("Content-Length", Long.toString(fileSize));
        }

        // Stream temp file to real response
        HttpServletResponse realResponse = (HttpServletResponse) getResponse();
        try (InputStream in = new BufferedInputStream(new FileInputStream(tempFile.toFile()), 32768);
             OutputStream out = new BufferedOutputStream(realResponse.getOutputStream(), 32768)) {
            byte[] buffer = new byte[32768];
            int bytesRead;
            while ((bytesRead = in.read(buffer)) != -1) {
                out.write(buffer, 0, bytesRead);
            }
            out.flush();
        } finally {
            // Always delete the temp file
            Files.deleteIfExists(tempFile);
        }
    }

    /**
     * Clean up temp file if writeToResponse was never called (e.g., on error).
     */
    public void cleanup() {
        try {
            if (fileOutputStream != null) {
                fileOutputStream.close();
            }
        } catch (IOException ignored) {
        }
        try {
            Files.deleteIfExists(tempFile);
        } catch (IOException ignored) {
        }
    }
}