From d0639e8925541bd6f6f386386c982d23b3828d3d Mon Sep 17 00:00:00 2001
From: Aldo Cortesi <aldo@nullcube.com>
Date: Sun, 24 Feb 2013 14:04:56 +1300
Subject: [PATCH] Handle server disconnects better.

Server connections can be closed for legitimate reasons, like timeouts. If
we've already pumped data over a server connection, we reconnect on error. If
not, we treat it as a legitimate error and pass it on to the client.

Fixes #85
---
 libmproxy/proxy.py  | 39 +++++++++++++++++++++++++++++----------
 test/test_server.py | 14 +++++++++++++-
 2 files changed, 42 insertions(+), 11 deletions(-)

diff --git a/libmproxy/proxy.py b/libmproxy/proxy.py
index c8fac5f46..088fe94c8 100644
--- a/libmproxy/proxy.py
+++ b/libmproxy/proxy.py
@@ -117,8 +117,8 @@ class ServerConnectionPool:
     def get_connection(self, scheme, host, port):
         sc = self.conn
         if self.conn and (host, port) != (sc.host, sc.port):
-            sc.terminate()
-            self.conn = None
+                sc.terminate()
+                self.conn = None
         if not self.conn:
             try:
                 self.conn = ServerConnection(self.config, host, port)
@@ -127,6 +127,9 @@ class ServerConnectionPool:
                 raise ProxyError(502, v)
         return self.conn
 
+    def del_connection(self, scheme, host, port):
+        self.conn = None
+
 
 class ProxyHandler(tcp.BaseHandler):
     def __init__(self, config, connection, client_address, server, channel, server_version):
@@ -181,14 +184,30 @@ class ProxyHandler(tcp.BaseHandler):
                         scheme, host, port = self.config.reverse_proxy
                     else:
                         scheme, host, port = request.scheme, request.host, request.port
-                    sc = self.server_conn_pool.get_connection(scheme, host, port)
-                    sc.send(request)
-                    sc.rfile.reset_timestamps()
-                    httpversion, code, msg, headers, content = http.read_response(
-                        sc.rfile,
-                        request.method,
-                        self.config.body_size_limit
-                    )
+
+                    # If we've already pumped a request over this connection,
+                    # it's possible that the server has timed out. If this is
+                    # the case, we want to reconnect without sending an error
+                    # to the client.
+                    while 1:
+                        try:
+                            sc = self.server_conn_pool.get_connection(scheme, host, port)
+                            sc.send(request)
+                            sc.rfile.reset_timestamps()
+                            httpversion, code, msg, headers, content = http.read_response(
+                                sc.rfile,
+                                request.method,
+                                self.config.body_size_limit
+                            )
+                        except http.HttpErrorConnClosed, v:
+                            if sc.requestcount > 1:
+                                self.server_conn_pool.del_connection(scheme, host, port)
+                                continue
+                            else:
+                                raise
+                        else:
+                            break
+
                     response = flow.Response(
                         request, httpversion, code, msg, headers, content, sc.cert,
                         sc.rfile.first_byte_timestamp, utils.timestamp()
diff --git a/test/test_server.py b/test/test_server.py
index 9df88400d..924b63b79 100644
--- a/test/test_server.py
+++ b/test/test_server.py
@@ -98,6 +98,19 @@ class TestHTTP(tservers.HTTPProxTest, SanityMixin):
         assert p.request("get:'%s':h'Connection'='close'"%response)
         tutils.raises("disconnect", p.request, "get:'%s'"%response)
 
+    def test_reconnect(self):
+        req = "get:'%s/p/200:b@1:da'"%self.urlbase
+        p = self.pathoc()
+        assert p.request(req)
+        # Server has disconnected. Mitmproxy should detect this, and reconnect.
+        assert p.request(req)
+        assert p.request(req)
+
+        # However, if the server disconnects on our first try, it's an error.
+        req = "get:'%s/p/200:b@1:d0'"%self.urlbase
+        p = self.pathoc()
+        tutils.raises("server disconnect", p.request, req)
+
     def test_proxy_ioerror(self):
         # Tests a difficult-to-trigger condition, where an IOError is raised
         # within our read loop.
@@ -106,7 +119,6 @@ class TestHTTP(tservers.HTTPProxTest, SanityMixin):
             tutils.raises("empty reply", self.pathod, "304")
 
 
-
 class TestHTTPS(tservers.HTTPProxTest, SanityMixin):
     ssl = True
     clientcerts = True