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