Merge pull request #164 from fredthomsen/add_proxy_support

Lukasa · Lukasa · commit 8c4a13bac35a · 2015-10-11T15:30:16.000+01:00
Add proxy support
diff --git a/CONTRIBUTORS.rst b/CONTRIBUTORS.rst
@@ -36,4 +36,4 @@ In chronological order:
 - Fred Thomsen (@fredthomsen)
 
   - Added support for upgrade of plaintext HTTP/1.1 to plaintext HTTP/2.
-
+  - Added proxy support.
diff --git a/hyper/common/connection.py b/hyper/common/connection.py
@@ -40,6 +40,10 @@ class HTTPConnection(object):
         :meth:`get_pushes() <hyper.HTTP20Connection.get_pushes>`).
     :param ssl_context: (optional) A class with custom certificate settings.
         If not provided then hyper's default ``SSLContext`` is used instead.
+    :param proxy_host: (optional) The proxy to connect to.  This can be an IP address
+        or a host name and may include a port.
+    :param proxy_port: (optional) The proxy port to connect to. If not provided 
+        and one also isn't provided in the ``proxy`` parameter, defaults to 8080.
     """
     def __init__(self,
                  host,
@@ -48,14 +52,20 @@ def __init__(self,
                  window_manager=None,
                  enable_push=False,
                  ssl_context=None,
+                 proxy_host=None,
+                 proxy_port=None,
                  **kwargs):
 
         self._host = host
         self._port = port
-        self._h1_kwargs = {'secure': secure, 'ssl_context': ssl_context}
+        self._h1_kwargs = {
+            'secure': secure, 'ssl_context': ssl_context, 
+            'proxy_host': proxy_host, 'proxy_port': proxy_port 
+        }
         self._h2_kwargs = {
             'window_manager': window_manager, 'enable_push': enable_push,
-            'secure': secure, 'ssl_context': ssl_context
+            'secure': secure, 'ssl_context': ssl_context, 
+            'proxy_host': proxy_host, 'proxy_port': proxy_port
         }
 
         # Add any unexpected kwargs to both dictionaries.
diff --git a/hyper/http11/connection.py b/hyper/http11/connection.py
@@ -47,9 +47,14 @@ class HTTP11Connection(object):
         port 443.
     :param ssl_context: (optional) A class with custom certificate settings.
         If not provided then hyper's default ``SSLContext`` is used instead.
+    :param proxy_host: (optional) The proxy to connect to.  This can be an IP 
+        address or a host name and may include a port.
+    :param proxy_port: (optional) The proxy port to connect to. If not provided 
+        and one also isn't provided in the ``proxy`` parameter, 
+        defaults to 8080.
     """
-    def __init__(self, host, port=None, secure=None, ssl_context=None,
-                 **kwargs):
+    def __init__(self, host, port=None, secure=None, ssl_context=None, 
+                 proxy_host=None, proxy_port=None, **kwargs):
         if port is None:
             try:
                 self.host, self.port = host.split(':')
@@ -75,6 +80,21 @@ def __init__(self, host, port=None, secure=None, ssl_context=None,
         self.ssl_context = ssl_context
         self._sock = None
 
+        # Setup proxy details if applicable.
+        if proxy_host:
+            if proxy_port is None:
+                try:
+                    self.proxy_host, self.proxy_port = proxy_host.split(':')
+                except ValueError:
+                    self.proxy_host, self.proxy_port = proxy_host, 8080
+                else:
+                    self.proxy_port = int(self.proxy_port)
+            else:
+                self.proxy_host, self.proxy_port = proxy_host, proxy_port
+        else:
+            self.proxy_host = None
+            self.proxy_port = None
+
         #: The size of the in-memory buffer used to store data from the
         #: network. This is used as a performance optimisation. Increase buffer
         #: size to improve performance: decrease it to conserve memory.
@@ -93,11 +113,19 @@ def connect(self):
         :returns: Nothing.
         """
         if self._sock is None:
-            sock = socket.create_connection((self.host, self.port), 5)
+            if not self.proxy_host:
+                host = self.host
+                port = self.port
+            else:
+                host = self.proxy_host
+                port = self.proxy_port
+                
+            sock = socket.create_connection((host, port), 5)
             proto = None
 
             if self.secure:
-                sock, proto = wrap_socket(sock, self.host, self.ssl_context)
+                assert not self.proxy_host, "Using a proxy with HTTPS not yet supported."
+                sock, proto = wrap_socket(sock, host, self.ssl_context)
 
             log.debug("Selected protocol: %s", proto)
             sock = BufferedSocket(sock, self.network_buffer_size)
@@ -176,7 +204,8 @@ def get_response(self):
         self._sock.advance_buffer(response.consumed)
 
         if (response.status == 101 and 
-           b'upgrade' in headers['connection'] and H2C_PROTOCOL.encode('utf-8') in headers['upgrade']):
+           b'upgrade' in headers['connection'] and 
+           H2C_PROTOCOL.encode('utf-8') in headers['upgrade']):
             raise HTTPUpgrade(H2C_PROTOCOL, self._sock)
 
         return HTTP11Response(
diff --git a/hyper/http20/connection.py b/hyper/http20/connection.py
@@ -56,9 +56,13 @@ class HTTP20Connection(object):
         :meth:`get_pushes() <hyper.HTTP20Connection.get_pushes>`).
     :param ssl_context: (optional) A class with custom certificate settings.
         If not provided then hyper's default ``SSLContext`` is used instead.
+    :param proxy_host: (optional) The proxy to connect to.  This can be an IP address
+        or a host name and may include a port.
+    :param proxy_port: (optional) The proxy port to connect to. If not provided 
+        and one also isn't provided in the ``proxy`` parameter, defaults to 8080.
     """
     def __init__(self, host, port=None, secure=None, window_manager=None, enable_push=False,
-                 ssl_context=None, **kwargs):
+                 ssl_context=None, proxy_host=None, proxy_port=None, **kwargs):
         """
         Creates an HTTP/2 connection to a specific server.
         """
@@ -81,6 +85,21 @@ def __init__(self, host, port=None, secure=None, window_manager=None, enable_pus
         self._enable_push = enable_push
         self.ssl_context = ssl_context
 
+        # Setup proxy details if applicable.
+        if proxy_host:
+            if proxy_port is None:
+                try:
+                    self.proxy_host, self.proxy_port = proxy_host.split(':')
+                except ValueError:
+                    self.proxy_host, self.proxy_port = proxy_host, 8080
+                else:
+                    self.proxy_port = int(self.proxy_port)
+            else:
+                self.proxy_host, self.proxy_port = proxy_host, proxy_port
+        else:
+            self.proxy_host = None
+            self.proxy_port = None
+
         #: The size of the in-memory buffer used to store data from the
         #: network. This is used as a performance optimisation. Increase buffer
         #: size to improve performance: decrease it to conserve memory.
@@ -222,10 +241,18 @@ def connect(self):
         :returns: Nothing.
         """
         if self._sock is None:
-            sock = socket.create_connection((self.host, self.port), 5)
+            if not self.proxy_host:
+                host = self.host
+                port = self.port
+            else:
+                host = self.proxy_host
+                port = self.proxy_port
+
+            sock = socket.create_connection((host, port), 5)
 
             if self.secure:
-                sock, proto = wrap_socket(sock, self.host, self.ssl_context)
+                assert not self.proxy_host, "Using a proxy with HTTPS not yet supported."
+                sock, proto = wrap_socket(sock, host, self.ssl_context)
             else:
                 proto = H2C_PROTOCOL
 
diff --git a/test/server.py b/test/server.py
@@ -87,10 +87,11 @@ class SocketLevelTest(object):
     A test-class that defines a few helper methods for running socket-level
     tests.
     """
-    def set_up(self, secure=True):
+    def set_up(self, secure=True, proxy=False):
         self.host = None
         self.port = None
-        self.secure = secure
+        self.secure = secure if not proxy else False
+        self.proxy = proxy
         self.server_thread = None
 
     def _start_server(self, socket_handler):
@@ -106,15 +107,27 @@ def _start_server(self, socket_handler):
         )
         self.server_thread.start()
         ready_event.wait()
+
         self.host = self.server_thread.host
         self.port = self.server_thread.port
         self.secure = self.server_thread.secure
 
     def get_connection(self):
         if self.h2:
-            return HTTP20Connection(self.host, self.port, self.secure)
+            if not self.proxy:
+                return HTTP20Connection(self.host, self.port, self.secure)
+            else:
+                return HTTP20Connection('http2bin.org', secure=self.secure, 
+                                        proxy_host=self.host, 
+                                        proxy_port=self.port)
         else:
-            return HTTP11Connection(self.host, self.port, self.secure)
+            if not self.proxy:
+                return HTTP11Connection(self.host, self.port, self.secure)
+            else:
+                return HTTP11Connection('httpbin.org', secure=self.secure, 
+                                        proxy_host=self.host, 
+                                        proxy_port=self.port)
+
 
     def get_encoder(self):
         """
diff --git a/test/test_abstraction.py b/test/test_abstraction.py
@@ -8,26 +8,30 @@ class TestHTTPConnection(object):
     def test_h1_kwargs(self):
         c = HTTPConnection(
             'test', 443, secure=False, window_manager=True, enable_push=True,
-            ssl_context=False, other_kwarg=True
+            ssl_context=False, proxy_host=False, proxy_port=False, other_kwarg=True
         )
 
         assert c._h1_kwargs == {
             'secure': False,
             'ssl_context': False,
+            'proxy_host': False,
+            'proxy_port': False,
             'other_kwarg': True,
         }
 
     def test_h2_kwargs(self):
         c = HTTPConnection(
             'test', 443, secure=False, window_manager=True, enable_push=True,
-            ssl_context=True, other_kwarg=True
+            ssl_context=True, proxy_host=False, proxy_port=False, other_kwarg=True
         )
 
         assert c._h2_kwargs == {
             'window_manager': True,
             'enable_push': True,
             'secure': False,
             'ssl_context': True,
+            'proxy_host': False,
+            'proxy_port': False,
             'other_kwarg': True,
         }
 
diff --git a/test/test_http11.py b/test/test_http11.py
@@ -43,27 +43,59 @@ def test_initialization_no_port(self):
         assert c.host == 'httpbin.org'
         assert c.port == 80
         assert not c.secure
+        assert not c.proxy_host
 
     def test_initialization_inline_port(self):
         c = HTTP11Connection('httpbin.org:443')
 
         assert c.host == 'httpbin.org'
         assert c.port == 443
         assert c.secure
+        assert not c.proxy_host
 
     def test_initialization_separate_port(self):
         c = HTTP11Connection('localhost', 8080)
 
         assert c.host == 'localhost'
         assert c.port == 8080
         assert not c.secure
+        assert not c.proxy_host
 
     def test_can_override_security(self):
         c = HTTP11Connection('localhost', 443, secure=False)
 
         assert c.host == 'localhost'
         assert c.port == 443
         assert not c.secure
+        assert not c.proxy_host
+
+    def test_initialization_proxy(self):
+        c = HTTP11Connection('httpbin.org', proxy_host='localhost')
+
+        assert c.host == 'httpbin.org'
+        assert c.port == 80
+        assert not c.secure
+        assert c.proxy_host == 'localhost'
+        assert c.proxy_port == 8080
+
+    def test_initialization_proxy_with_inline_port(self):
+        c = HTTP11Connection('httpbin.org', proxy_host='localhost:8443')
+
+        assert c.host == 'httpbin.org'
+        assert c.port == 80
+        assert not c.secure
+        assert c.proxy_host == 'localhost'
+        assert c.proxy_port == 8443
+
+    def test_initialization_proxy_with_separate_port(self):
+        c = HTTP11Connection('httpbin.org', proxy_host='localhost', proxy_port=8443)
+
+        assert c.host == 'httpbin.org'
+        assert c.port == 80
+        assert not c.secure
+        assert c.proxy_host == 'localhost'
+        assert c.proxy_port == 8443
+
 
     def test_basic_request(self):
         c = HTTP11Connection('httpbin.org')
@@ -84,6 +116,25 @@ def test_basic_request(self):
 
         assert received == expected
 
+    def test_proxy_request(self):
+        c = HTTP11Connection('httpbin.org', proxy_host='localhost')
+        c._sock = sock = DummySocket()
+
+        c.request('GET', '/get', headers={'User-Agent': 'hyper'})
+
+        expected = (
+            b"GET /get HTTP/1.1\r\n"
+            b"User-Agent: hyper\r\n"
+            b"connection: Upgrade, HTTP2-Settings\r\n"
+            b"upgrade: h2c\r\n"
+            b"HTTP2-Settings: AAQAAP//\r\n"
+            b"host: httpbin.org\r\n"
+            b"\r\n"
+        )
+        received = b''.join(sock.queue)
+
+        assert received == expected
+
     def test_request_with_bytestring_body(self):
         c = HTTP11Connection('httpbin.org')
         c._sock = sock = DummySocket()
diff --git a/test/test_hyper.py b/test/test_hyper.py
@@ -41,11 +41,27 @@ def test_connections_accept_hosts_and_ports(self):
         c = HTTP20Connection(host='www.google.com', port=8080)
         assert c.host =='www.google.com'
         assert c.port == 8080
+        assert c.proxy_host == None
 
     def test_connections_can_parse_hosts_and_ports(self):
         c = HTTP20Connection('www.google.com:8080')
         assert c.host == 'www.google.com'
         assert c.port == 8080
+        assert c.proxy_host == None
+
+    def test_connections_accept_proxy_hosts_and_ports(self):
+        c = HTTP20Connection('www.google.com', proxy_host='localhost:8443')
+        assert c.host == 'www.google.com'
+        assert c.proxy_host == 'localhost'
+        assert c.proxy_port == 8443
+
+    def test_connections_can_parse_proxy_hosts_and_ports(self):
+        c = HTTP20Connection('www.google.com', 
+                             proxy_host='localhost', 
+                             proxy_port=8443)
+        assert c.host == 'www.google.com'
+        assert c.proxy_host == 'localhost'
+        assert c.proxy_port == 8443
 
     def test_putrequest_establishes_new_stream(self):
         c = HTTP20Connection("www.google.com")
@@ -469,6 +485,19 @@ def data_cb(frame, *args):
         assert frames[0].type == WindowUpdateFrame.type
         assert frames[0].window_increment == 5535
 
+    def test_that_using_proxy_keeps_http_headers_intact(self):
+        sock = DummySocket()
+        c = HTTP20Connection('www.google.com', secure=False, proxy_host='localhost')
+        c._sock = sock
+        c.request('GET', '/')
+        s = c.recent_stream
+
+        assert s.headers == [
+            (':method', 'GET'),
+            (':scheme', 'http'),
+            (':authority', 'www.google.com'),
+            (':path', '/'),
+        ]
 
 class TestServerPush(object):
     def setup_method(self, method):
diff --git a/test/test_integration.py b/test/test_integration.py
diff --git a/test/test_integration_http11.py b/test/test_integration_http11.py
