Custom connection flow

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

@@ -1,5 +1,8 @@
 package org.littleshoot.proxy;
 
+import org.littleshoot.proxy.impl.ConnectionFlowStep; //Change: @AlmogBaku
+import org.littleshoot.proxy.impl.ProxyConnection; //Change: @AlmogBaku
+
 import io.netty.handler.codec.http.HttpObject;
 
 import java.net.InetSocketAddress;
@@ -49,6 +52,26 @@ public interface ChainedProxy extends SslEngineSource {
      */
     boolean requiresEncryption();
 
+    /**
+     * //Change: @AlmogBaku
+     * Implement this method to tell LittleProxy whether or not to use custom ConnectionFlow
+     * to the chained proxy for the given request. If true,
+     * LittleProxy will call {@link ChainedProxy#customConnectionFlow(ProxyConnection)} to obtain a
+     * ConnectionFlow.
+     *
+     * @return true of the connection to the chained proxy should be used a custom ConnectionFlow
+     */
+    boolean requiresCustomConnectionFlow();
+
+    /**
+     * //Change: @AlmogBaku
+     * Returns an {@link ConnectionFlowStep} to use for a server connection from
+     * LittleProxy to the client.
+     *
+     * @return
+     */
+    ConnectionFlowStep customConnectionFlow(ProxyConnection connection);
+
     /**
      * Filters requests on their way to the chained proxy.
      * 

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

@@ -1,5 +1,8 @@
 package org.littleshoot.proxy;
 
+import org.littleshoot.proxy.impl.ConnectionFlowStep;//Change: @AlmogBaku
+import org.littleshoot.proxy.impl.ProxyConnection;//Change: @AlmogBaku
+
 import io.netty.handler.codec.http.HttpObject;
 
 import java.net.InetSocketAddress;
@@ -36,6 +39,18 @@ public boolean requiresEncryption() {
         return false;
     }
 
+    //Change: @AlmogBaku
+    @Override
+    public boolean requiresCustomConnectionFlow() {
+        return false;
+    }
+
+    //Change: @AlmogBaku
+    @Override
+    public ConnectionFlowStep customConnectionFlow(ProxyConnection connection) {
+        return null;
+    }
+
     @Override
     public SSLEngine newSslEngine() {
         return null;

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

@@ -11,7 +11,8 @@
  * establishing a socket connection, SSL handshaking, HTTP CONNECT request
  * processing, and so on.
  */
-class ConnectionFlow {
+//Change(expose to public): @AlmogBaku
+public class ConnectionFlow {
     private Queue<ConnectionFlowStep> steps = new ConcurrentLinkedQueue<ConnectionFlowStep>();
 
     private final ClientToProxyConnection clientConnection;
@@ -77,12 +78,13 @@ void start() {
     }
 
     /**
+     * //Change(expose to public): @AlmogBaku
      * <p>
      * Advances the flow. {@link #advance()} will be called until we're either
      * out of steps, or a step has failed.
      * </p>
      */
-    void advance() {
+    public void advance() {
         currentStep = steps.poll();
         if (currentStep == null) {
             succeed();
@@ -171,12 +173,13 @@ void succeed() {
     }
 
     /**
+     * //Change(expose to public): @AlmogBaku
      * Called when the flow fails at some {@link ConnectionFlowStep}.
      * Disconnects the {@link ProxyToServerConnection} and informs the
      * {@link ClientToProxyConnection} that our connection failed.
      */
     @SuppressWarnings("unchecked")
-    void fail(final Throwable cause) {
+    public void fail(final Throwable cause) {
         final ConnectionState lastStateBeforeFailure = serverConnection
                 .getCurrentState();
         serverConnection.disconnect().addListener(
@@ -202,9 +205,10 @@ public void operationComplete(Future future)
     }
 
     /**
+     * //Change(expose to public): @AlmogBaku
      * Like {@link #fail(Throwable)} but with no cause.
      */
-    void fail() {
+    public void fail() {
         fail(null);
     }
 

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

@@ -3,14 +3,16 @@
 import io.netty.util.concurrent.Future;
 
 /**
+ * //Change(expose to public): @AlmogBaku
  * Represents a phase in a {@link ConnectionFlow}.
  */
-abstract class ConnectionFlowStep {
+public abstract class ConnectionFlowStep {
     private final ProxyConnectionLogger LOG;
     private final ProxyConnection connection;
     private final ConnectionState state;
 
     /**
+     * //Change(expose to protected): @AlmogBaku
      * Construct a new step in a connection flow.
      * 
      * @param connection
@@ -19,8 +21,8 @@ abstract class ConnectionFlowStep {
      *            the state that the connection will show while we're processing
      *            this step
      */
-    ConnectionFlowStep(ProxyConnection connection,
-            ConnectionState state) {
+    protected ConnectionFlowStep(ProxyConnection connection,
+                                 ConnectionState state) {
         super();
         this.connection = connection;
         this.state = state;
@@ -72,18 +74,20 @@ boolean shouldExecuteOnEventLoop() {
     protected abstract Future execute();
 
     /**
+     * //Change(expose to protected): @AlmogBaku
      * When the flow determines that this step was successful, it calls into
      * this method. The default implementation simply continues with the flow.
      * Other implementations may choose to not continue and instead wait for a
      * message or something like that.
      * 
      * @param flow
      */
-    void onSuccess(ConnectionFlow flow) {
+    protected void onSuccess(ConnectionFlow flow) {
         flow.advance();
     }
 
     /**
+     * //Change(expose to protected): @AlmogBaku
      * <p>
      * Any messages that are read from the underlying connection while we're at
      * this step of the connection flow are passed to this method.
@@ -104,7 +108,7 @@ void onSuccess(ConnectionFlow flow) {
      * @param msg
      *            the message read from the underlying connection
      */
-    void read(ConnectionFlow flow, Object msg) {
+    protected void read(ConnectionFlow flow, Object msg) {
         LOG.debug("Received message while in the middle of connecting: {}", msg);
     }
 

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

@@ -1,6 +1,7 @@
 package org.littleshoot.proxy.impl;
 
-enum ConnectionState {
+//Change(expose to public): @AlmogBaku
+public enum ConnectionState {
     /**
      * Connection attempting to connect.
      */

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

@@ -60,7 +60,8 @@
  *            the type of "initial" message. This will be either
  *            {@link HttpResponse} or {@link HttpRequest}.
  */
-abstract class ProxyConnection<I extends HttpObject> extends
+//Change(expose to public): @AlmogBaku
+public abstract class ProxyConnection<I extends HttpObject> extends
         SimpleChannelInboundHandler<Object> {
     protected final ProxyConnectionLogger LOG = new ProxyConnectionLogger(this);
 
@@ -262,7 +263,8 @@ protected void writeRaw(ByteBuf buf) {
         writeToChannel(buf);
     }
 
-    protected ChannelFuture writeToChannel(final Object msg) {
+    //Change(expose to public): @AlmogBaku
+    public ChannelFuture writeToChannel(final Object msg) {
         return channel.writeAndFlush(msg);
     }
 

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

@@ -559,6 +559,11 @@ private void initializeConnectionFlow() {
                 connectLock)
                 .then(ConnectChannel);
 
+        //Change: @AlmogBaku
+        if (chainedProxy != null && chainedProxy.requiresCustomConnectionFlow()) {
+            connectionFlow.then(chainedProxy.customConnectionFlow(this));
+        }
+
         if (chainedProxy != null && chainedProxy.requiresEncryption()) {
             connectionFlow.then(serverConnection.EncryptChannel(chainedProxy
                     .newSslEngine()));
@@ -686,12 +691,18 @@ public void operationComplete(ChannelFuture arg0) throws Exception {
             }
         }
 
-        void onSuccess(ConnectionFlow flow) {
+        //Change(expose to protected): @AlmogBaku
+        protected void onSuccess(ConnectionFlow flow) {
             // Do nothing, since we want to wait for the CONNECT response to
             // come back
         }
 
-        void read(ConnectionFlow flow, Object msg) {
+        protected void read(ConnectionFlow flow, Object msg) {
+            //@AlmogBaku: Ignore previous reads
+            if (msg == LastHttpContent.EMPTY_LAST_CONTENT) {
+                return;
+            }
+
             // Here we're handling the response from a chained proxy to our
             // earlier CONNECT request
             boolean connectOk = false;