Add custom proxy headers (#2272)

* Adding proxy header to aiohttp requests

* Added contributor

CONTRIBUTORS.txt

@@ -48,6 +48,7 @@ Chris Moore
 Christopher Schmitt
 Claudiu Popa
 Damien Nadé
+Dan Xu
 Daniel García
 Daniel Nelson
 Danny Song

aiohttp/client.py

@@ -157,7 +157,8 @@ def _request(self, method, url, *,
                  timeout=sentinel,
                  verify_ssl=None,
                  fingerprint=None,
-                 ssl_context=None):
+                 ssl_context=None,
+                 proxy_headers=None):
 
         # NOTE: timeout clamps existing connect and read timeouts.  We cannot
         # set the default to None because we need to detect if the user wants
@@ -188,6 +189,8 @@ def _request(self, method, url, *,
 
         # Merge with default headers and transform to CIMultiDict
         headers = self._prepare_headers(headers)
+        proxy_headers = self._prepare_headers(proxy_headers)
+
         if auth is None:
             auth = self._default_auth
         # It would be confusing if we support explicit Authorization header
@@ -238,7 +241,7 @@ def _request(self, method, url, *,
                         proxy=proxy, proxy_auth=proxy_auth, timer=timer,
                         session=self, auto_decompress=self._auto_decompress,
                         verify_ssl=verify_ssl, fingerprint=fingerprint,
-                        ssl_context=ssl_context)
+                        ssl_context=ssl_context, proxy_headers=proxy_headers)
 
                     # connection timeout
                     try:

aiohttp/client_reqrep.py

@@ -80,7 +80,8 @@ def __init__(self, method, url, *,
                  loop=None, response_class=None,
                  proxy=None, proxy_auth=None, proxy_from_env=False,
                  timer=None, session=None, auto_decompress=True,
-                 verify_ssl=None, fingerprint=None, ssl_context=None):
+                 verify_ssl=None, fingerprint=None, ssl_context=None,
+                 proxy_headers=None):
 
         if verify_ssl is False and ssl_context is not None:
             raise ValueError(
@@ -121,7 +122,7 @@ def __init__(self, method, url, *,
         self.update_cookies(cookies)
         self.update_content_encoding(data)
         self.update_auth(auth)
-        self.update_proxy(proxy, proxy_auth, proxy_from_env)
+        self.update_proxy(proxy, proxy_auth, proxy_from_env, proxy_headers)
         self.update_fingerprint(fingerprint)
 
         self.update_body_from_data(data)
@@ -317,7 +318,7 @@ def update_expect_continue(self, expect=False):
         if expect:
             self._continue = helpers.create_future(self.loop)
 
-    def update_proxy(self, proxy, proxy_auth, proxy_from_env):
+    def update_proxy(self, proxy, proxy_auth, proxy_from_env, proxy_headers):
         if proxy_from_env and not proxy:
             proxy_url = getproxies().get(self.original_url.scheme)
             proxy = URL(proxy_url) if proxy_url else None
@@ -327,6 +328,7 @@ def update_proxy(self, proxy, proxy_auth, proxy_from_env):
             raise ValueError("proxy_auth must be None or BasicAuth() tuple")
         self.proxy = proxy
         self.proxy_auth = proxy_auth
+        self.proxy_headers = proxy_headers
 
     def update_fingerprint(self, fingerprint):
         if fingerprint:

aiohttp/connector.py

@@ -809,9 +809,14 @@ def _create_direct_connection(self, req):
 
     @asyncio.coroutine
     def _create_proxy_connection(self, req):
+        headers = {}
+        if req.proxy_headers is not None:
+            headers = req.proxy_headers
+        headers[hdrs.HOST] = req.headers[hdrs.HOST]
+
         proxy_req = ClientRequest(
             hdrs.METH_GET, req.proxy,
-            headers={hdrs.HOST: req.headers[hdrs.HOST]},
+            headers=headers,
             auth=req.proxy_auth,
             loop=self._loop,
             verify_ssl=req.verify_ssl,

changes/2001.feature

@@ -0,0 +1 @@
+Added variable to customize proxy headers

docs/client_reference.rst

@@ -291,6 +291,11 @@ The client session supports the context manager protocol for self closing.
 
          .. versionadded:: 2.3
 
+      :param dict proxy_headers: HTTP headers to send to the proxy if the
+         parameter proxy has been provided.
+
+         .. versionadded:: 2.3
+
       :return ClientResponse: a :class:`client response <ClientResponse>`
          object.
 

tests/test_proxy.py

@@ -70,6 +70,38 @@ def test_connect(self, ClientRequestMock):
             ssl_context=None,
             verify_ssl=None)
 
+    @mock.patch('aiohttp.connector.ClientRequest')
+    def test_proxy_headers(self, ClientRequestMock):
+        req = ClientRequest(
+            'GET', URL('http://www.python.org'),
+            proxy=URL('http://proxy.example.com'),
+            proxy_headers={'Foo': 'Bar'},
+            loop=self.loop)
+        self.assertEqual(str(req.proxy), 'http://proxy.example.com')
+
+        # mock all the things!
+        connector = aiohttp.TCPConnector(loop=self.loop)
+        connector._resolve_host = make_mocked_coro([mock.MagicMock()])
+
+        proto = mock.Mock(**{
+            'transport.get_extra_info.return_value': False,
+        })
+        self.loop.create_connection = make_mocked_coro(
+            (proto.transport, proto))
+        conn = self.loop.run_until_complete(connector.connect(req))
+        self.assertEqual(req.url, URL('http://www.python.org'))
+        self.assertIs(conn._protocol, proto)
+        self.assertIs(conn.transport, proto.transport)
+
+        ClientRequestMock.assert_called_with(
+            'GET', URL('http://proxy.example.com'),
+            auth=None,
+            fingerprint=None,
+            headers={'Host': 'www.python.org', 'Foo': 'Bar'},
+            loop=self.loop,
+            ssl_context=None,
+            verify_ssl=None)
+
     @mock.patch('aiohttp.connector.ClientRequest', **clientrequest_mock_attrs)
     def test_connect_req_verify_ssl_true(self, ClientRequestMock):
         req = ClientRequest(