#96 If server times out, send 504 to client

adamfisk · Oct 8, 2013 · 10faf9d · 10faf9d
1 parent 4d31a8a
commit 10faf9d
Show file tree

Hide file tree

Showing 5 changed files with 64 additions and 737 deletions.
diff --git a/src/main/java/org/littleshoot/proxy/impl/ClientToProxyConnection.java b/src/main/java/org/littleshoot/proxy/impl/ClientToProxyConnection.java
@@ -112,9 +112,9 @@ public class ClientToProxyConnection extends ProxyConnection<HttpRequest> {
     private volatile HttpFilters currentFilters;
 
     private volatile SSLSession clientSslSession;
-    
+
     /**
-     * Tracks whether or not this ClientToProxyConnection is current doing MITM. 
+     * Tracks whether or not this ClientToProxyConnection is current doing MITM.
      */
     private volatile boolean mitming = false;
 
@@ -129,7 +129,8 @@ public class ClientToProxyConnection extends ProxyConnection<HttpRequest> {
 
         if (sslEngineSource != null) {
             LOG.debug("Enabling encryption of traffic from client to proxy");
-            encrypt(pipeline, sslEngineSource.newSslEngine(), authenticateClients)
+            encrypt(pipeline, sslEngineSource.newSslEngine(),
+                    authenticateClients)
                     .addListener(
                             new GenericFutureListener<Future<? super Channel>>() {
                                 @Override
@@ -374,6 +375,18 @@ protected void connected() {
         recordClientConnected();
     }
 
+    @Override
+    protected void timedOut() {
+        boolean clientReadMoreRecentlyThanServer =
+                currentServerConnection == null
+                        || this.lastReadTime > currentServerConnection.lastReadTime;
+        if (clientReadMoreRecentlyThanServer) {
+            LOG.debug("Server timed out: {}", currentServerConnection);
+            writeGatewayTimeout();
+        }
+        super.timedOut();
+    }
+
     /**
      * On disconnect of the client, disconnect all server connections.
      */
@@ -607,7 +620,7 @@ private void initChannelPipeline(ChannelPipeline pipeline) {
         pipeline.addLast(
                 "idle",
                 new IdleStateHandler(0, 0, proxyServer
-                        .getIdleConnectionTimeout()));
+                        .getIdleConnectionTimeout() + 2));
 
         pipeline.addLast("handler", this);
     }
@@ -1008,6 +1021,19 @@ private void writeBadGateway(HttpRequest request) {
         disconnect();
     }
 
+    /**
+     * Tells the client that the connection to the server timed out.
+     * 
+     * @param request
+     */
+    private void writeGatewayTimeout() {
+        String body = "Gateway Timeout";
+        DefaultFullHttpResponse response = responseFor(HttpVersion.HTTP_1_1,
+                HttpResponseStatus.GATEWAY_TIMEOUT, body);
+        response.headers().set(HttpHeaders.Names.CONNECTION, "close");
+        write(response);
+    }
+
     /**
      * Factory for {@link DefaultFullHttpResponse}s.
      * 
@@ -1094,11 +1120,11 @@ private String identifyHostAndPort(HttpRequest httpRequest) {
     private void writeEmptyBuffer() {
         write(Unpooled.EMPTY_BUFFER);
     }
-    
+
     public boolean isMitming() {
         return mitming;
     }
-    
+
     protected void setMitming(boolean isMitming) {
         this.mitming = isMitming;
     }
@@ -1203,4 +1229,5 @@ private FlowContext flowContext() {
             return new FlowContext(this);
         }
     }
+
 }
diff --git a/src/main/java/org/littleshoot/proxy/impl/ProxyConnection.java b/src/main/java/org/littleshoot/proxy/impl/ProxyConnection.java
@@ -82,6 +82,7 @@ abstract class ProxyConnection<I extends HttpObject> extends
 
     private volatile ConnectionState currentState;
     private volatile boolean tunneling = false;
+    protected volatile long lastReadTime = 0;
 
     /**
      * If using encryption, this holds our {@link SSLEngine}.
@@ -119,6 +120,8 @@ protected ProxyConnection(ConnectionState initialState,
     protected void read(Object msg) {
         LOG.debug("Reading: {}", msg);
 
+        lastReadTime = System.currentTimeMillis();
+
         if (tunneling) {
             // In tunneling mode, this connection is simply shoveling bytes
             readRaw((ByteBuf) msg);
@@ -288,6 +291,14 @@ protected void disconnected() {
         become(DISCONNECTED);
         LOG.debug("Disconnected");
     }
+
+    /**
+     * This method is called when the underlying {@link Channel} times out due
+     * to an idle timeout.
+     */
+    protected void timedOut() {
+        disconnect();
+    }
 
     /**
      * <p>
@@ -627,8 +638,8 @@ public final void userEventTriggered(ChannelHandlerContext ctx, Object evt)
             throws Exception {
         try {
             if (evt instanceof IdleStateEvent) {
-                LOG.debug("Got idle, disconnecting");
-                disconnect();
+                LOG.debug("Got idle");
+                timedOut();
             }
         } finally {
             super.userEventTriggered(ctx, evt);

diff --git a/src/main/java/org/littleshoot/proxy/impl/ProxyToServerConnection.java b/src/main/java/org/littleshoot/proxy/impl/ProxyToServerConnection.java
@@ -307,6 +307,12 @@ protected void becameWritable() {
         super.becameWritable();
         this.clientConnection.serverBecameWriteable(this);
     }
+
+    @Override
+    protected void timedOut() {
+        super.timedOut();
+        clientConnection.timedOut();
+    }
 
     @Override
     protected void disconnected() {