diff --git a/src/main/java/org/littleshoot/proxy/ChainedProxy.java b/src/main/java/org/littleshoot/proxy/ChainedProxy.java index 3292c4253..97caa7aa8 100644 --- a/src/main/java/org/littleshoot/proxy/ChainedProxy.java +++ b/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. * diff --git a/src/main/java/org/littleshoot/proxy/ChainedProxyAdapter.java b/src/main/java/org/littleshoot/proxy/ChainedProxyAdapter.java index 61f42881f..0eca7370c 100644 --- a/src/main/java/org/littleshoot/proxy/ChainedProxyAdapter.java +++ b/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; diff --git a/src/main/java/org/littleshoot/proxy/impl/ConnectionFlow.java b/src/main/java/org/littleshoot/proxy/impl/ConnectionFlow.java index 23a65f08e..ad3e41a72 100644 --- a/src/main/java/org/littleshoot/proxy/impl/ConnectionFlow.java +++ b/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 steps = new ConcurrentLinkedQueue(); private final ClientToProxyConnection clientConnection; @@ -77,12 +78,13 @@ void start() { } /** + * //Change(expose to public): @AlmogBaku *

* Advances the flow. {@link #advance()} will be called until we're either * out of steps, or a step has failed. *

*/ - 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); } diff --git a/src/main/java/org/littleshoot/proxy/impl/ConnectionFlowStep.java b/src/main/java/org/littleshoot/proxy/impl/ConnectionFlowStep.java index c60910d59..caaf355b4 100644 --- a/src/main/java/org/littleshoot/proxy/impl/ConnectionFlowStep.java +++ b/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,6 +74,7 @@ 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 @@ -79,11 +82,12 @@ boolean shouldExecuteOnEventLoop() { * * @param flow */ - void onSuccess(ConnectionFlow flow) { + protected void onSuccess(ConnectionFlow flow) { flow.advance(); } /** + * //Change(expose to protected): @AlmogBaku *

* 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); } diff --git a/src/main/java/org/littleshoot/proxy/impl/ConnectionState.java b/src/main/java/org/littleshoot/proxy/impl/ConnectionState.java index 371fcd586..f3598c7f6 100644 --- a/src/main/java/org/littleshoot/proxy/impl/ConnectionState.java +++ b/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. */ diff --git a/src/main/java/org/littleshoot/proxy/impl/ProxyConnection.java b/src/main/java/org/littleshoot/proxy/impl/ProxyConnection.java index 58c3eb240..b3db3c7dd 100644 --- a/src/main/java/org/littleshoot/proxy/impl/ProxyConnection.java +++ b/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 extends +//Change(expose to public): @AlmogBaku +public abstract class ProxyConnection extends SimpleChannelInboundHandler { 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); } diff --git a/src/main/java/org/littleshoot/proxy/impl/ProxyToServerConnection.java b/src/main/java/org/littleshoot/proxy/impl/ProxyToServerConnection.java index c9cd0f2d4..8645cdba2 100644 --- a/src/main/java/org/littleshoot/proxy/impl/ProxyToServerConnection.java +++ b/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;