Implements custom proxy error handler

pom.xml

@@ -3,7 +3,7 @@
     <groupId>org.littleshoot</groupId>
     <artifactId>littleproxy</artifactId>
     <packaging>jar</packaging>
-    <version>1.1.3.2-VGS-SNAPSHOT</version>
+    <version>1.1.3.3-VGS-SNAPSHOT</version>
     <name>LittleProxy</name>
     <description>
         LittleProxy is a high performance HTTP proxy written in Java and using the Netty networking framework.

src/main/java/org/littleshoot/proxy/ExceptionHandler.java

@@ -0,0 +1,11 @@
+package org.littleshoot.proxy;
+
+public interface ExceptionHandler {
+
+  /**
+   * Handles proxy exceptions
+   *
+   * @param cause error cause
+   */
+  void handle(Throwable cause);
+}

src/main/java/org/littleshoot/proxy/HttpProxyServerBootstrap.java

@@ -181,6 +181,36 @@ HttpProxyServerBootstrap withChainProxyManager(
     HttpProxyServerBootstrap withManInTheMiddle(
             MitmManagerFactory mitmManager);
 
+    /**
+     * <p>
+     * Specify an {@link ExceptionHandler} to handle client to proxy errors
+     * </p>
+     *
+     * <p>
+     * Default = null
+     * </p>
+     *
+     * @param clientToProxyExHandler
+     * @return exception handler
+     */
+    HttpProxyServerBootstrap withClientToProxyExHandler(
+        ExceptionHandler clientToProxyExHandler);
+
+    /**
+     * <p>
+     * Specify an {@link ExceptionHandler} to handle proxy to server errors
+     * </p>
+     *
+     * <p>
+     * Default = null
+     * </p>
+     *
+     * @param proxyToServerExHandler
+     * @return exception handler
+     */
+    HttpProxyServerBootstrap withProxyToServerExHandler(
+        ExceptionHandler proxyToServerExHandler);
+
     /**
      * <p>
      * Specify a {@link HttpFiltersSource} to use for filtering requests and/or

src/main/java/org/littleshoot/proxy/impl/ClientToProxyConnection.java

@@ -26,6 +26,7 @@
 import org.apache.commons.lang3.StringUtils;
 import org.littleshoot.proxy.ActivityTracker;
 import org.littleshoot.proxy.BadGatewayFailureHttpResponseComposer;
+import org.littleshoot.proxy.ExceptionHandler;
 import org.littleshoot.proxy.FailureHttpResponseComposer;
 import org.littleshoot.proxy.FlowContext;
 import org.littleshoot.proxy.FullFlowContext;
@@ -755,7 +756,13 @@ protected void exceptionCaught(Throwable cause) {
                 LOG.info("An executor rejected a read or write operation on the ClientToProxyConnection (this is normal if the proxy is shutting down). Message: " + cause.getMessage());
                 LOG.debug("A RejectedExecutionException occurred on ClientToProxyConnection", cause);
             } else {
-                LOG.error("Caught an exception on ClientToProxyConnection", cause);
+                ExceptionHandler exHandler = proxyServer.getClientToProxyExHandler();
+                if (exHandler != null) {
+                    LOG.debug("Custom exception handler '" + exHandler.toString() + "' invoked", cause);
+                    exHandler.handle(cause);
+                } else {
+                    LOG.error("Caught an exception on ClientToProxyConnection", cause);
+                }
             }
         } finally {
             // always disconnect the client when an exception occurs on the channel

src/main/java/org/littleshoot/proxy/impl/DefaultHttpProxyServer.java

@@ -31,6 +31,7 @@
 import org.littleshoot.proxy.MitmManager;
 import org.littleshoot.proxy.MitmManagerFactory;
 import org.littleshoot.proxy.ProxyAuthenticator;
+import org.littleshoot.proxy.ExceptionHandler;
 import org.littleshoot.proxy.SslEngineSource;
 import org.littleshoot.proxy.TransportProtocol;
 import org.littleshoot.proxy.UnknownTransportProtocolException;
@@ -110,6 +111,8 @@ public class DefaultHttpProxyServer implements HttpProxyServer {
     private final ProxyAuthenticator proxyAuthenticator;
     private final ChainedProxyManager chainProxyManager;
     private final MitmManagerFactory mitmManagerFactory;
+    private final ExceptionHandler clientToProxyExHandler;
+    private final ExceptionHandler proxyToServerExHandler;
     private final HttpFiltersSource filtersSource;
     private final FailureHttpResponseComposer unrecoverableFailureHttpResponseComposer;
     private final boolean transparent;
@@ -243,6 +246,8 @@ private DefaultHttpProxyServer(ServerGroup serverGroup,
             ProxyAuthenticator proxyAuthenticator,
             ChainedProxyManager chainProxyManager,
             MitmManagerFactory mitmManagerFactory,
+            ExceptionHandler clientToProxyExHandler,
+            ExceptionHandler proxyToServerExHandler,
             HttpFiltersSource filtersSource,
             FailureHttpResponseComposer unrecoverableFailureHttpResponseComposer,
             boolean transparent,
@@ -266,6 +271,8 @@ private DefaultHttpProxyServer(ServerGroup serverGroup,
         this.proxyAuthenticator = proxyAuthenticator;
         this.chainProxyManager = chainProxyManager;
         this.mitmManagerFactory = mitmManagerFactory;
+        this.clientToProxyExHandler = clientToProxyExHandler;
+        this.proxyToServerExHandler = proxyToServerExHandler;
         this.filtersSource = filtersSource;
         this.unrecoverableFailureHttpResponseComposer = unrecoverableFailureHttpResponseComposer;
         this.transparent = transparent;
@@ -401,6 +408,8 @@ public HttpProxyServerBootstrap clone() {
                     proxyAuthenticator,
                     chainProxyManager,
                     mitmManagerFactory,
+                    clientToProxyExHandler,
+                    proxyToServerExHandler,
                     filtersSource,
                     unrecoverableFailureHttpResponseComposer,
                     transparent,
@@ -579,6 +588,14 @@ protected MitmManager getMitmManager(Channel channel) {
         return null;
     }
 
+    protected ExceptionHandler getClientToProxyExHandler() {
+        return clientToProxyExHandler;
+    }
+
+    protected ExceptionHandler getProxyToServerExHandler() {
+        return proxyToServerExHandler;
+    }
+
     protected SslEngineSource getSslEngineSource() {
         return sslEngineSource;
     }
@@ -621,6 +638,8 @@ private static class DefaultHttpProxyServerBootstrap implements HttpProxyServerB
         private ProxyAuthenticator proxyAuthenticator = null;
         private ChainedProxyManager chainProxyManager = null;
         private MitmManagerFactory mitmManagerFactory = null;
+        private ExceptionHandler clientToProxyExHandler = null;
+        private ExceptionHandler proxyToServerExHandler = null;
         private HttpFiltersSource filtersSource = new HttpFiltersSourceAdapter();
         private FailureHttpResponseComposer unrecoverableFailureHttpResponseComposer = new BadGatewayFailureHttpResponseComposer();
         private boolean transparent = false;
@@ -652,6 +671,8 @@ private DefaultHttpProxyServerBootstrap(
                 ProxyAuthenticator proxyAuthenticator,
                 ChainedProxyManager chainProxyManager,
                 MitmManagerFactory mitmManagerFactory,
+                ExceptionHandler clientToProxyExHandler,
+                ExceptionHandler proxyToServerExHandler,
                 HttpFiltersSource filtersSource,
                 FailureHttpResponseComposer unrecoverableFailureHttpResponseComposer,
                 boolean transparent, int idleConnectionTimeout,
@@ -674,6 +695,8 @@ private DefaultHttpProxyServerBootstrap(
             this.proxyAuthenticator = proxyAuthenticator;
             this.chainProxyManager = chainProxyManager;
             this.mitmManagerFactory = mitmManagerFactory;
+            this.clientToProxyExHandler = clientToProxyExHandler;
+            this.proxyToServerExHandler = proxyToServerExHandler;
             this.filtersSource = filtersSource;
             this.unrecoverableFailureHttpResponseComposer = unrecoverableFailureHttpResponseComposer;
             this.transparent = transparent;
@@ -807,6 +830,20 @@ public HttpProxyServerBootstrap withManInTheMiddle(
             return this;
         }
 
+        @Override
+        public HttpProxyServerBootstrap withProxyToServerExHandler(
+            ExceptionHandler proxyToServerExHandler) {
+            this.proxyToServerExHandler = proxyToServerExHandler;
+            return this;
+        }
+
+        @Override
+        public HttpProxyServerBootstrap withClientToProxyExHandler(
+            ExceptionHandler clientToProxyExHandler) {
+            this.clientToProxyExHandler = clientToProxyExHandler;
+            return this;
+        }
+
         @Override
         public HttpProxyServerBootstrap withFiltersSource(
                 HttpFiltersSource filtersSource) {
@@ -923,6 +960,7 @@ private DefaultHttpProxyServer build() {
                     transportProtocol, determineListenAddress(),
                     sslEngineSource, authenticateSslClients,
                     proxyAuthenticator, chainProxyManager, mitmManagerFactory,
+                    clientToProxyExHandler, proxyToServerExHandler,
                     filtersSource, unrecoverableFailureHttpResponseComposer, transparent,
                     idleConnectionTimeout, activityTrackers, connectTimeout,
                     serverResolver, readThrottleBytesPerSecond, writeThrottleBytesPerSecond,

src/main/java/org/littleshoot/proxy/impl/ProxyToServerConnection.java

@@ -39,6 +39,7 @@
 import org.littleshoot.proxy.FullFlowContext;
 import org.littleshoot.proxy.HttpFilters;
 import org.littleshoot.proxy.MitmManager;
+import org.littleshoot.proxy.ExceptionHandler;
 import org.littleshoot.proxy.TransportProtocol;
 import org.littleshoot.proxy.UnknownTransportProtocolException;
 
@@ -439,7 +440,13 @@ protected void exceptionCaught(Throwable cause) {
                 LOG.info("An executor rejected a read or write operation on the ProxyToServerConnection (this is normal if the proxy is shutting down). Message: " + cause.getMessage());
                 LOG.debug("A RejectedExecutionException occurred on ProxyToServerConnection", cause);
             } else {
-                LOG.error("Caught an exception on ProxyToServerConnection", cause);
+                ExceptionHandler exHandler = proxyServer.getProxyToServerExHandler();
+                if (exHandler != null) {
+                    LOG.debug("Custom exception handler '" + exHandler.toString() + "' invoked", cause);
+                    exHandler.handle(cause);
+                } else {
+                    LOG.error("Caught an exception on ProxyToServerConnection", cause);
+                }
             }
         } finally {
             if (!is(DISCONNECTED)) {

src/test/java/org/littleshoot/proxy/CustomClientToProxyExHandlerTest.java

@@ -0,0 +1,51 @@
+package org.littleshoot.proxy;
+
+import io.netty.channel.ChannelHandlerContext;
+import io.netty.handler.codec.http.HttpRequest;
+import org.apache.http.NoHttpResponseException;
+import org.junit.Assert;
+import org.junit.Test;
+import org.littleshoot.proxy.extras.SelfSignedMitmManagerFactory;
+
+import java.util.ArrayList;
+import java.util.List;
+
+public class CustomClientToProxyExHandlerTest extends AbstractProxyTest {
+
+  private final List<Throwable> customExHandlerEntered = new ArrayList<>();
+
+  private static final String EXCEPTION_MESSAGE = "Error occurred in client to proxy connection";
+
+  @Override
+  protected void setUp() {
+    this.proxyServer = bootstrapProxy()
+        .withPort(0)
+        .withManInTheMiddle(new SelfSignedMitmManagerFactory())
+        .withClientToProxyExHandler(new ExceptionHandler() {
+          @Override
+          public void handle(Throwable cause) {
+            customExHandlerEntered.add(cause);
+          }
+        })
+        .withFiltersSource(new HttpFiltersSourceAdapter() {
+          @Override
+          public HttpFilters filterRequest(HttpRequest originalRequest,
+                                           ChannelHandlerContext ctx) {
+            throw new RuntimeException(EXCEPTION_MESSAGE);
+          }
+        })
+        .start();
+  }
+
+  @Test
+  public void testCustomClientToProxyExHandler() throws Exception {
+    try {
+      httpGetWithApacheClient(webHost, DEFAULT_RESOURCE, true, true);
+    } catch (NoHttpResponseException e) {
+      // expected
+    }
+    Assert.assertFalse("Custom ex handler was not called", customExHandlerEntered.isEmpty());
+    Assert.assertEquals("Incorrect exception was passed to custom ex handles",
+        customExHandlerEntered.get(0).getMessage(), EXCEPTION_MESSAGE);
+  }
+}

src/test/java/org/littleshoot/proxy/CustomProxyToServerExHandlerTest.java

@@ -0,0 +1,44 @@
+package org.littleshoot.proxy;
+
+import org.junit.Assert;
+import org.junit.Test;
+import org.littleshoot.proxy.extras.SelfSignedMitmManagerFactory;
+
+import java.util.ArrayList;
+import java.util.List;
+
+public class CustomProxyToServerExHandlerTest extends MitmWithBadServerAuthenticationTCPChainedProxyTest {
+
+  private final List<Throwable> customExHandlerEntered = new ArrayList<>();
+
+  @Override
+  protected void setUp() {
+    this.upstreamProxy = upstreamProxy().start();
+
+    this.proxyServer = bootstrapProxy()
+        .withPort(0)
+        .withChainProxyManager(chainedProxyManager())
+        .plusActivityTracker(DOWNSTREAM_TRACKER)
+        .withManInTheMiddle(new SelfSignedMitmManagerFactory())
+        .withProxyToServerExHandler(new ExceptionHandler() {
+          @Override
+          public void handle(Throwable cause) {
+            customExHandlerEntered.add(cause);
+          }
+        })
+        .start();
+  }
+
+  @Override
+  protected void tearDown() throws Exception {
+    this.upstreamProxy.abort();
+  }
+
+  @Test
+  public void testCustomProxyToServerExHandler() throws Exception {
+    super.testSimpleGetRequestOverHTTPS();
+    Assert.assertFalse("Custom ex handler was not called", customExHandlerEntered.isEmpty());
+    Assert.assertEquals("Incorrect exception was passed to custom ex handles",
+        customExHandlerEntered.get(0).getMessage(), "javax.net.ssl.SSLHandshakeException: General SSLEngine problem");
+  }
+}