Merge pull request #7 from yola/fix_scheme

Explicitly set url scheme when passing X-Forwarded-Proto

proxyprefix/wsgi.py

@@ -6,8 +6,14 @@ def __init__(self, app):
 
     def __call__(self, environ, start_response):
         prefix = environ.get('HTTP_X_FORWARDED_PREFIX')
+        scheme = environ.get('HTTP_X_FORWARDED_PROTO')
+
         if prefix:
             prefix_paths(environ, prefix)
+
+        if scheme:
+            set_scheme(environ, scheme)
+
         return self.app(environ, start_response)
 
 
@@ -25,3 +31,10 @@ def prefix_paths(environ, prefix):
     # this django quirk:
     if environ.get('SCRIPT_URL'):
         environ['SCRIPT_URL'] = ''
+
+
+def set_scheme(environ, scheme):
+    """Force environ to be http or https."""
+    scheme = 'https' if scheme == 'https' else 'http'
+    environ['wsgi.url_scheme'] = scheme
+    environ['HTTPS'] = 'on' if scheme == 'https' else 'off'

setup.py

@@ -19,7 +19,7 @@
         'Programming Language :: Python :: 2',
         'Topic :: Internet :: WWW/HTTP :: WSGI :: Middleware',
     ],
-    extras_require = {
+    extras_require={
         'djproxy': ['djproxy>=2.0.0'],
     },
 )

tests/test_reverse_proxied_app.py

@@ -1,14 +1,18 @@
 from mock import Mock, patch
 from unittest2 import TestCase
 
-from proxyprefix.wsgi import prefix_paths, ReverseProxiedApp
+from proxyprefix.wsgi import prefix_paths, set_scheme, ReverseProxiedApp
 
 
 class TestReverseProxiedApp(TestCase):
     def setUp(self):
-        self.prefix_paths_patcher = patch('proxyprefix.wsgi.prefix_paths')
-        self.prefix_paths = self.prefix_paths_patcher.start()
-        self.addCleanup(self.prefix_paths_patcher.stop)
+        prefix_paths_patcher = patch('proxyprefix.wsgi.prefix_paths')
+        self.prefix_paths = prefix_paths_patcher.start()
+        self.addCleanup(prefix_paths_patcher.stop)
+
+        set_scheme_patcher = patch('proxyprefix.wsgi.set_scheme')
+        self.set_scheme = set_scheme_patcher.start()
+        self.addCleanup(set_scheme_patcher.stop)
 
         self.environ = {}
         self.start_response = Mock()
@@ -30,6 +34,15 @@ def test_it_prefixes_paths_with_HTTP_X_FORWARDED_PREFIX(self):
         self.proxied_app(self.environ, self.start_response)
         self.prefix_paths.assert_called_with(self.environ, 'prefix')
 
+    def test_it_does_not_set_scheme_if_no_HTTP_X_FORWARDED_PROTO(self):
+        self.proxied_app(self.environ, self.start_response)
+        self.assertFalse(self.set_scheme.called)
+
+    def test_it_sets_scheme_to_HTTP_X_FORWARDED_PROTO(self):
+        self.environ['HTTP_X_FORWARDED_PROTO'] = 'http'
+        self.proxied_app(self.environ, self.start_response)
+        self.set_scheme.assert_called_with(self.environ, 'http')
+
 
 class TestPrefixPaths(TestCase):
     """prefix_paths(environ, prefix)"""
@@ -51,3 +64,26 @@ def test_it_removes_SCRIPT_URL_if_present(self):
         self.environ['SCRIPT_URL'] = '/script/url'
         prefix_paths(self.environ, 'prefix')
         self.assertFalse(self.environ['SCRIPT_URL'])
+
+
+class TestSetScheme(TestCase):
+    """set_scheme(environ, scheme)"""
+
+    def setUp(self):
+        self.environ = {}
+
+    def test_sets_wsgi_url_scheme_to_https_if_scheme_is_https(self):
+        set_scheme(self.environ, 'https')
+        self.assertEqual(self.environ['wsgi.url_scheme'], 'https')
+
+    def test_sets_wsgi_url_scheme_to_http_if_scheme_is_not_https(self):
+        set_scheme(self.environ, 'foo')
+        self.assertEqual(self.environ['wsgi.url_scheme'], 'http')
+
+    def test_sets_HTTPS_to_on_if_scheme_is_https(self):
+        set_scheme(self.environ, 'https')
+        self.assertEqual(self.environ['HTTPS'], 'on')
+
+    def test_sets_HTTPS_to_off_if_scheme_is_not_https(self):
+        set_scheme(self.environ, 'foo')
+        self.assertEqual(self.environ['HTTPS'], 'off')