From 8254187bf316e4cdda32e1be0e07e73bc2446243 Mon Sep 17 00:00:00 2001
From: Aldo Cortesi <aldo@nullcube.com>
Date: Sun, 10 Jun 2012 10:10:46 +1200
Subject: [PATCH] Add proxy.should_connection_close, and strip out unused code.

---
 libmproxy/proxy.py | 50 ++++++++++++++++------------------------------
 test/test_proxy.py | 36 +++++++++------------------------
 2 files changed, 26 insertions(+), 60 deletions(-)

diff --git a/libmproxy/proxy.py b/libmproxy/proxy.py
index 938b5d5cf..0c032bb00 100644
--- a/libmproxy/proxy.py
+++ b/libmproxy/proxy.py
@@ -183,38 +183,22 @@ def parse_init_http(line):
     return method, url, mm[0], mm[1]
 
 
-#FIXME: Return full HTTP version specification from here. Allow non-HTTP
-#protocol specs, and make it all editable.
-def parse_request_line(request):
+def should_connection_close(httpmajor, httpminor, headers):
     """
-        Parse a proxy request line. Return (method, scheme, host, port, path, minor).
-        Raise ProxyError on error.
+        Checks the HTTP version and headers to see if this connection should be
+        closed.
     """
-    try:
-        method, url, protocol = string.split(request)
-    except ValueError:
-        raise ProxyError(400, "Can't parse request")
-    if method == 'CONNECT':
-        scheme = None
-        path = None
-        try:
-            host, port = url.split(":")
-        except ValueError:
-            raise ProxyError(400, "Can't parse request")
-        port = int(port)
-    elif url.startswith("/") or url == "*":
-            scheme, port, host, path = None, None, None, url
-    else:
-        parts = utils.parse_url(url)
-        if not parts:
-            raise ProxyError(400, "Invalid url: %s"%url)
-        scheme, host, port, path = parts
-    if not protocol.startswith("HTTP/"):
-        raise ProxyError(400, "Unsupported protocol")
-    major,minor = protocol.split('/')[1].split('.')
-    major = int(major)
-    minor = int(minor)
-    return method, scheme, host, port, path, minor
+    if "connection" in headers:
+        for value in ",".join(headers['connection']).split(","):
+            value = value.strip()
+            if value == "close":
+                return True
+            elif value == "keep-alive":
+                return False
+    # HTTP 1.1 connections are assumed to be persistent
+    if httpmajor == 1 and httpminor == 1:
+        return False
+    return True
 
 
 class FileLike:
@@ -302,7 +286,6 @@ class ServerConnection:
         self.rfile, self.wfile = server.makefile('rb'), server.makefile('wb')
 
     def send(self, request):
-        request.close = self.close
         try:
             d = request._assemble()
             if not d:
@@ -350,6 +333,7 @@ class ProxyHandler(SocketServer.StreamRequestHandler):
     def __init__(self, config, request, client_address, server, q):
         self.config = config
         self.mqueue = q
+        self.server_conn = None
         SocketServer.StreamRequestHandler.__init__(self, request, client_address, server)
 
     def handle(self):
@@ -389,9 +373,9 @@ class ProxyHandler(SocketServer.StreamRequestHandler):
                 else:
                     if self.config.reverse_proxy:
                         scheme, host, port = self.config.reverse_proxy
-                        server_conn = ServerConnection(self.config, scheme, host, port)
                     else:
-                        server_conn = ServerConnection(self.config, request.scheme, request.host, request.port)
+                        scheme, host, port = request.scheme, request.host, request.port
+                    server_conn = ServerConnection(self.config, scheme, host, port)
                     server_conn.send(request)
                     try:
                         response = server_conn.read_response(request)
diff --git a/test/test_proxy.py b/test/test_proxy.py
index a5dc666ac..3e1d32fea 100644
--- a/test/test_proxy.py
+++ b/test/test_proxy.py
@@ -24,6 +24,15 @@ def test_read_chunked():
     tutils.raises(proxy.ProxyError, proxy.read_chunked, s, None)
 
 
+def test_should_connection_close():
+    h = flow.ODictCaseless()
+    assert proxy.should_connection_close(1, 0, h)
+    assert not proxy.should_connection_close(1, 1, h)
+
+    h["connection"] = ["keep-alive"]
+    assert not proxy.should_connection_close(1, 1, h)
+
+
 def test_read_http_body():
     d = Dummy()
     h = flow.ODict()
@@ -47,33 +56,6 @@ def test_read_http_body():
     assert len(proxy.read_http_body(s, d, h, True, 100)) == 7
 
 
-class TestParseRequestLine:
-    def test_simple(self):
-        tutils.raises(proxy.ProxyError, proxy.parse_request_line, "")
-
-        u = "GET ... HTTP/1.1"
-        tutils.raises("invalid url", proxy.parse_request_line, u)
-
-        u = "GET http://foo.com:8888/test HTTP/1.1"
-        m, s, h, po, pa, minor = proxy.parse_request_line(u)
-        assert m == "GET"
-        assert s == "http"
-        assert h == "foo.com"
-        assert po == 8888
-        assert pa == "/test"
-        assert minor == 1
-
-    def test_connect(self):
-        u = "CONNECT host.com:443 HTTP/1.0"
-        expected = ('CONNECT', None, 'host.com', 443, None, 0)
-        ret = proxy.parse_request_line(u)
-        assert expected == ret
-
-    def test_inner(self):
-        u = "GET / HTTP/1.1"
-        assert proxy.parse_request_line(u) == ('GET', None, None, None, '/', 1)
-
-
 class TestFileLike:
     def test_wrap(self):
         s = cStringIO.StringIO("foobar\nfoobar")