Merge pull request #3185 from brettdh/3183-support-all-proxy-env-var

Support ALL_PROXY environment variable

8 files changed, 115 insertions, 18 deletions

diff --git a/requests/sessions.py b/requests/sessions.py
index e28419bc..d124ea0f 100644
--- a/requests/sessions.py
+++ b/requests/sessions.py
@@ -228,10 +228,10 @@ class SessionRedirectMixin(object):
         if self.trust_env and not should_bypass_proxies(url):
             environ_proxies = get_environ_proxies(url)
 
-            proxy = environ_proxies.get(scheme)
+            proxy = environ_proxies.get('all', environ_proxies.get(scheme))
 
             if proxy:
-                new_proxies.setdefault(scheme, environ_proxies[scheme])
+                new_proxies.setdefault(scheme, proxy)
 
         if 'Proxy-Authorization' in headers:
             del headers['Proxy-Authorization']
diff --git a/requests/utils.py b/requests/utils.py
index eaa89df2..8d17b6b2 100644
--- a/requests/utils.py
+++ b/requests/utils.py
@@ -584,11 +584,20 @@ def select_proxy(url, proxies):
     proxies = proxies or {}
     urlparts = urlparse(url)
     if urlparts.hostname is None:
-        proxy = None
-    else:
-        proxy = proxies.get(urlparts.scheme+'://'+urlparts.hostname)
-    if proxy is None:
-        proxy = proxies.get(urlparts.scheme)
+        return proxies.get('all', proxies.get(urlparts.scheme))
+
+    proxy_keys = [
+        'all://' + urlparts.hostname,
+        'all',
+        urlparts.scheme + '://' + urlparts.hostname,
+        urlparts.scheme,
+    ]
+    proxy = None
+    for proxy_key in proxy_keys:
+        if proxy_key in proxies:
+            proxy = proxies[proxy_key]
+            break
+
     return proxy
 
 
diff --git a/requirements.txt b/requirements.txt
index 1305d3f8..8426eecb 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -10,6 +10,7 @@ Jinja2==2.8
 MarkupSafe==0.23
 py==1.4.31
 Pygments==2.1.1
+PySocks==1.5.6
 pytest==2.8.7
 pytest-cov==2.2.1
 pytest-httpbin==0.2.0
diff --git a/tests/test_lowlevel.py b/tests/test_lowlevel.py
index eb6d6273..f3dd1b11 100644
--- a/tests/test_lowlevel.py
+++ b/tests/test_lowlevel.py
@@ -1,14 +1,18 @@
+import os
+import pytest
 import threading
 import requests
 
 from tests.testserver.server import Server
 
+from .utils import override_environ
+
 
 def test_chunked_upload():
     """can safely send generators"""
     close_server = threading.Event()
     server = Server.basic_response_server(wait_to_close_event=close_server)
-    data = (i for i in [b'a', b'b', b'c']) 
+    data = iter([b'a', b'b', b'c'])
 
     with server as (host, port):
         url = 'http://{0}:{1}/'.format(host, port)
@@ -17,3 +21,36 @@ def test_chunked_upload():
 
     assert r.status_code == 200
     assert r.request.headers['Transfer-Encoding'] == 'chunked'
+
+
+_schemes_by_var_prefix = [
+    ('http', ['http']),
+    ('https', ['https']),
+    ('all', ['http', 'https']),
+]
+
+_proxy_combos = []
+for prefix, schemes in _schemes_by_var_prefix:
+    for scheme in schemes:
+        _proxy_combos.append(("{0}_proxy".format(prefix), scheme))
+
+_proxy_combos += [(var.upper(), scheme) for var, scheme in _proxy_combos]
+
+
+@pytest.mark.parametrize("var,scheme", _proxy_combos)
+def test_use_proxy_from_environment(httpbin, var, scheme):
+    url = "{0}://httpbin.org".format(scheme)
+    fake_proxy = Server()  # do nothing with the requests; just close the socket
+    with fake_proxy as (host, port):
+        proxy_url = "socks5://{0}:{1}".format(host, port)
+        kwargs = {var: proxy_url}
+        with override_environ(**kwargs):
+            # fake proxy's lack of response will cause a ConnectionError
+            with pytest.raises(requests.exceptions.ConnectionError):
+                requests.get(url)
+
+        # the fake proxy received a request
+        assert len(fake_proxy.handler_results) == 1
+
+        # it had actual content (not checking for SOCKS protocol for now)
+        assert len(fake_proxy.handler_results[0]) > 0
diff --git a/tests/test_requests.py b/tests/test_requests.py
index d01749d2..4bcf4536 100755
--- a/tests/test_requests.py
+++ b/tests/test_requests.py
@@ -30,6 +30,7 @@ from requests.models import urlencode
 from requests.hooks import default_hooks
 
 from .compat import StringIO, u
+from .utils import override_environ
 
 # Requests to this URL should always fail with a connection timeout (nothing
 # listening on that port)
@@ -1548,6 +1549,27 @@ def test_requests_are_updated_each_time(httpbin):
         assert session.calls[-1] == send_call
 
 
+@pytest.mark.parametrize("var,url,proxy", [
+    ('http_proxy', 'http://example.com', 'socks5://proxy.com:9876'),
+    ('https_proxy', 'https://example.com', 'socks5://proxy.com:9876'),
+    ('all_proxy', 'http://example.com', 'socks5://proxy.com:9876'),
+    ('all_proxy', 'https://example.com', 'socks5://proxy.com:9876'),
+])
+def test_proxy_env_vars_override_default(var, url, proxy):
+    session = requests.Session()
+    prep = PreparedRequest()
+    prep.prepare(method='GET', url=url)
+
+    kwargs = {
+        var: proxy
+    }
+    scheme = urlparse(url).scheme
+    with override_environ(**kwargs):
+        proxies = session.rebuild_proxies(prep, {})
+        assert scheme in proxies
+        assert proxies[scheme] == proxy
+
+
 @pytest.mark.parametrize(
     'data', (
         (('a', 'b'), ('c', 'd')),
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 68e5a173..17149d26 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -320,17 +320,28 @@ def test_dotted_netmask(mask, expected):
     assert dotted_netmask(mask) == expected
 
 
+http_proxies = {'http': 'http://http.proxy',
+                'http://some.host': 'http://some.host.proxy'}
+all_proxies = {'all': 'socks5://http.proxy',
+               'all://some.host': 'socks5://some.host.proxy'}
 @pytest.mark.parametrize(
-    'url, expected', (
-        ('hTTp://u:p@Some.Host/path', 'http://some.host.proxy'),
-        ('hTTp://u:p@Other.Host/path', 'http://http.proxy'),
-        ('hTTps://Other.Host', None),
-        ('file:///etc/motd', None),
+    'url, expected, proxies', (
+        ('hTTp://u:p@Some.Host/path', 'http://some.host.proxy', http_proxies),
+        ('hTTp://u:p@Other.Host/path', 'http://http.proxy', http_proxies),
+        ('hTTp:///path', 'http://http.proxy', http_proxies),
+        ('hTTps://Other.Host', None, http_proxies),
+        ('file:///etc/motd', None, http_proxies),
+
+        ('hTTp://u:p@Some.Host/path', 'socks5://some.host.proxy', all_proxies),
+        ('hTTp://u:p@Other.Host/path', 'socks5://http.proxy', all_proxies),
+        ('hTTp:///path', 'socks5://http.proxy', all_proxies),
+        ('hTTps://Other.Host', 'socks5://http.proxy', all_proxies),
+
+        # XXX: unsure whether this is reasonable behavior
+        ('file:///etc/motd', 'socks5://http.proxy', all_proxies),
     ))
-def test_select_proxies(url, expected):
+def test_select_proxies(url, expected, proxies):
     """Make sure we can select per-host proxies correctly."""
-    proxies = {'http': 'http://http.proxy',
-               'http://some.host': 'http://some.host.proxy'}
     assert select_proxy(url, proxies) == expected
 
 
diff --git a/tests/testserver/server.py b/tests/testserver/server.py
index 8b9643c3..5be478b3 100644
--- a/tests/testserver/server.py
+++ b/tests/testserver/server.py
@@ -25,10 +25,10 @@ class Server(threading.Thread):
     """Dummy server using for unit testing"""
     WAIT_EVENT_TIMEOUT = 5
 
-    def __init__(self, handler, host='localhost', port=0, requests_to_handle=1, wait_to_close_event=None):
+    def __init__(self, handler=None, host='localhost', port=0, requests_to_handle=1, wait_to_close_event=None):
         super(Server, self).__init__()
 
-        self.handler = handler
+        self.handler = handler or consume_socket_content
         self.handler_results = []
 
         self.host = host
diff --git a/tests/utils.py b/tests/utils.py
new file mode 100644
index 00000000..6cb75bfb
--- /dev/null
+++ b/tests/utils.py
@@ -0,0 +1,17 @@
+import contextlib
+import os
+
+
+@contextlib.contextmanager
+def override_environ(**kwargs):
+    save_env = dict(os.environ)
+    for key, value in kwargs.items():
+        if value is None:
+            del os.environ[key]
+        else:
+            os.environ[key] = value
+    try:
+        yield
+    finally:
+        os.environ.clear()
+        os.environ.update(save_env)