diff --git a/src/main/java/org/littleshoot/proxy/AcceptHandler.java b/src/main/java/org/littleshoot/proxy/AcceptHandler.java new file mode 100644 index 000000000..327fa941a --- /dev/null +++ b/src/main/java/org/littleshoot/proxy/AcceptHandler.java @@ -0,0 +1,23 @@ +package org.littleshoot.proxy; + +import io.netty.buffer.ByteBuf; +import io.netty.channel.ChannelHandlerContext; + +/** + *

+ * Interface to manage the bytes read the first time after a client + * connection has been accepted. + *

+ */ +public interface AcceptHandler { + + /** + * Process the bytes coming from the first read performed on the + * underlying channel after the connection has been accepted. + * @param bytes + * the bytes read from the underlying channel + * @return the bytes that will be processed by the proxy + */ + ByteBuf process(ChannelHandlerContext ctx, ByteBuf bytes); + +} diff --git a/src/main/java/org/littleshoot/proxy/HttpProxyServerBootstrap.java b/src/main/java/org/littleshoot/proxy/HttpProxyServerBootstrap.java index b3d186031..241a343d0 100644 --- a/src/main/java/org/littleshoot/proxy/HttpProxyServerBootstrap.java +++ b/src/main/java/org/littleshoot/proxy/HttpProxyServerBootstrap.java @@ -268,6 +268,17 @@ HttpProxyServerBootstrap withConnectTimeout( */ HttpProxyServerBootstrap withServerResolver(HostResolver serverResolver); + /** + *

+ * Specify an {@link AcceptHandler} to manage the bytes read the first + * time after a client connection has been accepted. + *

+ * + * @param acceptHandler + * @return + */ + HttpProxyServerBootstrap withAcceptHandler(AcceptHandler acceptHandler); + /** *

* Add an {@link ActivityTracker} for tracking activity in this proxy. diff --git a/src/main/java/org/littleshoot/proxy/impl/ClientToProxyConnection.java b/src/main/java/org/littleshoot/proxy/impl/ClientToProxyConnection.java index 70d799dc5..7852b888e 100644 --- a/src/main/java/org/littleshoot/proxy/impl/ClientToProxyConnection.java +++ b/src/main/java/org/littleshoot/proxy/impl/ClientToProxyConnection.java @@ -3,6 +3,7 @@ import io.netty.buffer.ByteBuf; import io.netty.buffer.Unpooled; import io.netty.channel.Channel; +import io.netty.channel.ChannelHandlerContext; import io.netty.channel.ChannelPipeline; import io.netty.handler.codec.http.DefaultFullHttpResponse; import io.netty.handler.codec.http.DefaultHttpRequest; @@ -25,6 +26,7 @@ import io.netty.util.concurrent.GenericFutureListener; import org.apache.commons.codec.binary.Base64; import org.apache.commons.lang3.StringUtils; +import org.littleshoot.proxy.AcceptHandler; import org.littleshoot.proxy.ActivityTracker; import org.littleshoot.proxy.FlowContext; import org.littleshoot.proxy.FullFlowContext; @@ -763,6 +765,7 @@ protected void exceptionCaught(Throwable cause) { private void initChannelPipeline(ChannelPipeline pipeline) { LOG.debug("Configuring ChannelPipeline"); + pipeline.addLast("acceptMonitor", acceptMonitor); pipeline.addLast("bytesReadMonitor", bytesReadMonitor); pipeline.addLast("bytesWrittenMonitor", bytesWrittenMonitor); @@ -1393,6 +1396,27 @@ protected void setMitming(boolean isMitming) { this.mitming = isMitming; } + /*************************************************************************** + * Accept tracker + * + * Used to have access to the bytes read from the underlying channel after + * the client connection has been accepted, it is added as the first + * handler of the pipeline and immediately removed after the first read. + **************************************************************************/ + private final ByteStreamMonitor acceptMonitor = new ByteStreamMonitor() { + @Override + protected ByteBuf processBytes(ChannelHandlerContext ctx, ByteBuf bytes) { + AcceptHandler acceptHandler = proxyServer.getAcceptHandler(); + if (acceptHandler != null) { + ByteBuf processedBytes = acceptHandler.process(ctx, bytes); + ctx.pipeline().remove("acceptMonitor"); + return processedBytes; + } else { + return bytes; + } + } + }; + /*************************************************************************** * Activity Tracking/Statistics * diff --git a/src/main/java/org/littleshoot/proxy/impl/DefaultHttpProxyServer.java b/src/main/java/org/littleshoot/proxy/impl/DefaultHttpProxyServer.java index 6d972e4a4..ddca72f9a 100644 --- a/src/main/java/org/littleshoot/proxy/impl/DefaultHttpProxyServer.java +++ b/src/main/java/org/littleshoot/proxy/impl/DefaultHttpProxyServer.java @@ -17,6 +17,7 @@ import io.netty.handler.traffic.GlobalTrafficShapingHandler; import io.netty.util.concurrent.GlobalEventExecutor; import org.apache.commons.io.IOUtils; +import org.littleshoot.proxy.AcceptHandler; import org.littleshoot.proxy.ActivityTracker; import org.littleshoot.proxy.ChainedProxyManager; import org.littleshoot.proxy.DefaultHostResolver; @@ -121,6 +122,11 @@ public class DefaultHttpProxyServer implements HttpProxyServer { */ private final AtomicBoolean stopped = new AtomicBoolean(false); + /** + * Track the bytes read after the client connection has been accepted. + */ + private final AcceptHandler acceptHandler; + /** * Track all ActivityTrackers for tracking proxying activity. */ @@ -234,6 +240,7 @@ private DefaultHttpProxyServer(ServerGroup serverGroup, HttpFiltersSource filtersSource, boolean transparent, int idleConnectionTimeout, + AcceptHandler acceptHandler, Collection activityTrackers, int connectTimeout, HostResolver serverResolver, @@ -252,6 +259,7 @@ private DefaultHttpProxyServer(ServerGroup serverGroup, this.filtersSource = filtersSource; this.transparent = transparent; this.idleConnectionTimeout = idleConnectionTimeout; + this.acceptHandler = acceptHandler; if (activityTrackers != null) { this.activityTrackers.addAll(activityTrackers); } @@ -366,6 +374,7 @@ public HttpProxyServerBootstrap clone() { filtersSource, transparent, idleConnectionTimeout, + acceptHandler, activityTrackers, connectTimeout, serverResolver, @@ -549,6 +558,10 @@ protected Collection getActivityTrackers() { return activityTrackers; } + protected AcceptHandler getAcceptHandler() { + return acceptHandler; + } + public String getProxyAlias() { return proxyAlias; } @@ -574,6 +587,7 @@ private static class DefaultHttpProxyServerBootstrap implements HttpProxyServerB private HttpFiltersSource filtersSource = new HttpFiltersSourceAdapter(); private boolean transparent = false; private int idleConnectionTimeout = 70; + private AcceptHandler acceptHandler; private Collection activityTrackers = new ConcurrentLinkedQueue(); private int connectTimeout = 40000; private HostResolver serverResolver = new DefaultHostResolver(); @@ -599,6 +613,7 @@ private DefaultHttpProxyServerBootstrap( MitmManager mitmManager, HttpFiltersSource filtersSource, boolean transparent, int idleConnectionTimeout, + AcceptHandler acceptHandler, Collection activityTrackers, int connectTimeout, HostResolver serverResolver, long readThrottleBytesPerSecond, @@ -617,6 +632,7 @@ private DefaultHttpProxyServerBootstrap( this.filtersSource = filtersSource; this.transparent = transparent; this.idleConnectionTimeout = idleConnectionTimeout; + this.acceptHandler = acceptHandler; if (activityTrackers != null) { this.activityTrackers.addAll(activityTrackers); } @@ -781,6 +797,12 @@ public HttpProxyServerBootstrap withServerResolver( return this; } + @Override + public HttpProxyServerBootstrap withAcceptHandler(AcceptHandler acceptHandler) { + this.acceptHandler = acceptHandler; + return this; + } + @Override public HttpProxyServerBootstrap plusActivityTracker( ActivityTracker activityTracker) { @@ -823,7 +845,7 @@ transportProtocol, determineListenAddress(), sslEngineSource, authenticateSslClients, proxyAuthenticator, chainProxyManager, mitmManager, filtersSource, transparent, - idleConnectionTimeout, activityTrackers, connectTimeout, + idleConnectionTimeout, acceptHandler, activityTrackers, connectTimeout, serverResolver, readThrottleBytesPerSecond, writeThrottleBytesPerSecond, localAddress, proxyAlias); } diff --git a/src/main/java/org/littleshoot/proxy/impl/ProxyConnection.java b/src/main/java/org/littleshoot/proxy/impl/ProxyConnection.java index 58c3eb240..4c27def1d 100644 --- a/src/main/java/org/littleshoot/proxy/impl/ProxyConnection.java +++ b/src/main/java/org/littleshoot/proxy/impl/ProxyConnection.java @@ -696,6 +696,31 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) protected abstract void bytesRead(int numberOfBytes); } + /** + * Utility handler for monitoring byte streams on this connection. + */ + @Sharable + protected abstract class ByteStreamMonitor extends + ChannelInboundHandlerAdapter { + @Override + public void channelRead(ChannelHandlerContext ctx, Object msg) + throws Exception { + Object processedMsg = msg; + + try { + if (msg instanceof ByteBuf) { + processedMsg = processBytes(ctx, (ByteBuf) msg); + } + } catch (Throwable t) { + LOG.warn("Unable to call processBytes", t); + } finally { + super.channelRead(ctx, processedMsg); + } + } + + protected abstract ByteBuf processBytes(ChannelHandlerContext ctx, ByteBuf bytes); + } + /** * Utility handler for monitoring requests read on this connection. */