From 81da08b8adf780d15f09e380063e372bc6db5195 Mon Sep 17 00:00:00 2001 From: Francesco Locunto Date: Wed, 7 Sep 2016 16:10:09 +0200 Subject: [PATCH] Added the AcceptHandler feature to the proxy. Using the withAcceptHandler() method on the proxy bootstrap is now possible to get access to the byte stream of the first read performed on the underlying channel after a connection has been accepted and before the message is http-parsed inside the channel pipeline. An implementation of the interface AcceptHandler can be able to read and/or modify the first bytes sent by a client before they are parsed, as well as close the connection if some requirements are not met and so on. --- .../org/littleshoot/proxy/AcceptHandler.java | 23 +++++++++++++++++ .../proxy/HttpProxyServerBootstrap.java | 11 ++++++++ .../proxy/impl/ClientToProxyConnection.java | 24 ++++++++++++++++++ .../proxy/impl/DefaultHttpProxyServer.java | 24 +++++++++++++++++- .../proxy/impl/ProxyConnection.java | 25 +++++++++++++++++++ 5 files changed, 106 insertions(+), 1 deletion(-) create mode 100644 src/main/java/org/littleshoot/proxy/AcceptHandler.java 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. */