package com.atlassian.stash.internal.web.filters;

import com.google.common.base.MoreObjects;
import com.google.common.base.Throwables;
import com.google.common.collect.ImmutableSet;
import com.google.common.io.ByteStreams;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nonnull;
import javax.servlet.*;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;
import java.io.IOException;
import java.io.OutputStream;
import java.util.Set;

/**
 * A simple {@code Filter} which wraps an {@code HttpServletResponse}'s {@code ServletOutputStream} with a version
 * which protects against {@code ClientAbortException}s and broken pipes which happen when HTTP clients disconnect
 * or navigate away while output is being streamed.
 * <p>
 * This is most common in two scenarios:
 * <ul>
 *     <li>HTTP hosting operations: Killing (Ctrl+C) a {@code git clone} causes the hosting process on the
 *     server side to fail with a SIGPIPE (signal 13) while trying to stream pack output</li>
 *     <li>Streaming pull request changes: When a pull request merge, used to display the diff, takes more
 *     than a minute, the client abandons the request. When the server then runs {@code git diff-tree} and
 *     tries to stream the changes the command fails due to the closed {@code OutputStream}</li>
 * </ul>
 * To avoid these types of failures, when the guard class detects a broken pipe or client abort, it masks
 * that exception and instead simply discards subsequent writes.
 */
public class StreamGuardFilter implements Filter {

    private static final Logger log = LoggerFactory.getLogger(StreamGuardFilter.class);

    @Override
    public void destroy() {
    }

    @Override
    public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain)
            throws IOException, ServletException {
        if (response instanceof HttpServletResponse) {
            response = new GuardedHttpServletResponse((HttpServletResponse) response);
        }

        chain.doFilter(request, response);
    }

    @Override
    public void init(FilterConfig filterConfig) {
    }

    private static class GuardedHttpServletResponse extends HttpServletResponseWrapper {

        public GuardedHttpServletResponse(HttpServletResponse response) {
            super(response);
        }

        @Override
        public ServletOutputStream getOutputStream() throws IOException {
            ServletOutputStream raw = super.getOutputStream();

            return new GuardedServletOutputStream(raw);
        }
    }

    //Note: The methods overloaded here are strongly inspired by the methods overridden in CoyoteOutputStream
    private static class GuardedServletOutputStream extends ServletOutputStream {

        private static final Set<String> ABORT_CLASSES = ImmutableSet.of(
                "org.apache.catalina.connector.ClientAbortException"
        );
        private static final Set<String> ABORT_MESSAGES = ImmutableSet.of(
                "An existing connection was forcibly closed by the remote host", //Windows, kill git clone via taskmgr
                "Broken pipe",                                                   //Linux/MacOS
                "Connection reset by peer: socket write error"                   //Windows, Ctrl+C git clone
        );

        private OutputStream out;

        public GuardedServletOutputStream(OutputStream out) {
            this.out = out;
        }

        @Override
        public void close() throws IOException {
            try {
                out.close();
            } catch (IOException e) {
                handleException(e);
            }
        }

        @Override
        public void flush() throws IOException {
            try {
                out.flush();
            } catch (IOException e) {
                handleException(e);
            }
        }

        @Override
        public void write(@Nonnull byte[] b) throws IOException {
            try {
                out.write(b);
            } catch (IOException e) {
                handleException(e);
            }
        }

        @Override
        public void write(@Nonnull byte[] b, int off, int len) throws IOException {
            try {
                out.write(b, off, len);
            } catch (IOException e) {
                handleException(e);
            }
        }

        @Override
        public void write(int b) throws IOException {
            try {
                out.write(b);
            } catch (IOException e) {
                handleException(e);
            }
        }

        private static boolean isClientAbort(IOException e) {
            //noinspection ThrowableResultOfMethodCallIgnored
            return ABORT_CLASSES.contains(e.getClass().getName()) ||
                    ABORT_MESSAGES.contains(MoreObjects.firstNonNull(Throwables.getRootCause(e), e).getMessage());
        }

        private void handleException(IOException e) throws IOException {
            if (isClientAbort(e)) {
                //If the client aborts the connection, rather than propagating that exception up (which can
                //result in forked git processes failing with SIGPIPE and/or ugly exceptions in the logs),
                //swallow the exception and instead discard all further output
                log.info("The remote client has aborted the connection");
                out = ByteStreams.nullOutputStream();
            } else {
                throw e;
            }
        }
    }
}
