Adding proxy header to aiohttp requests

aio-libs · Jun 21, 2017 · 4f3d2df · 4f3d2df
1 parent 6f2adac
commit 4f3d2df
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 6 deletions.
diff --git a/aiohttp/client.py b/aiohttp/client.py
@@ -152,7 +152,8 @@ def _request(self, method, url, *,
                  read_until_eof=True,
                  proxy=None,
                  proxy_auth=None,
-                 timeout=sentinel):
+                 timeout=sentinel,
+                 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
@@ -183,6 +184,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
@@ -223,7 +226,7 @@ def _request(self, method, url, *,
                         expect100=expect100, loop=self._loop,
                         response_class=self._response_class,
                         proxy=proxy, proxy_auth=proxy_auth, timer=timer,
-                        session=self)
+                        session=self, proxy_headers=proxy_headers)
 
                     # connection timeout
                     try:

diff --git a/aiohttp/client_reqrep.py b/aiohttp/client_reqrep.py
@@ -66,7 +66,7 @@ def __init__(self, method, url, *,
                  chunked=None, expect100=False,
                  loop=None, response_class=None,
                  proxy=None, proxy_auth=None, proxy_from_env=False,
-                 timer=None, session=None):
+                 timer=None, session=None, proxy_headers=None):
 
         if loop is None:
             loop = asyncio.get_event_loop()
@@ -99,7 +99,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_body_from_data(data)
         self.update_transfer_encoding()
@@ -295,7 +295,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
@@ -305,6 +305,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 keep_alive(self):
         if self.version < HttpVersion10:

diff --git a/aiohttp/connector.py b/aiohttp/connector.py
@@ -738,9 +738,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)
         try: