clear untrusted headers if trusted_proxy is set but does not match

2 files changed, 128 insertions, 109 deletions

diff --git a/waitress/task.py b/waitress/task.py
index e684eb1..68aa69a 100644
--- a/waitress/task.py
+++ b/waitress/task.py
@@ -27,7 +27,6 @@ from .utilities import (
     logger,
     queue_logger,
     undquote,
-    PROXY_HEADERS,
 )
 
 rename_headers = {  # or keep them without the HTTP_ prefix added
@@ -513,8 +512,6 @@ class WSGITask(Task):
         headers,
         trusted_proxy_count=1,
         trusted_proxy_headers=None,
-        log_untrusted_proxy_headers=False,
-        clear_untrusted_proxy_headers=True,
     ):
         if trusted_proxy_headers is None:
             trusted_proxy_headers = set()
@@ -660,15 +657,6 @@ class WSGITask(Task):
                 forwarded_host = proxy.host or forwarded_host
                 forwarded_proto = proxy.proto or forwarded_proto
 
-        # Clear out the untrusted proxy headers.
-        if clear_untrusted_proxy_headers:
-            clear_untrusted_headers(
-                headers,
-                untrusted_headers,
-                log_warning=log_untrusted_proxy_headers,
-                logger=self.logger,
-            )
-
         if forwarded_proto:
             forwarded_proto = forwarded_proto.lower()
 
@@ -739,6 +727,8 @@ class WSGITask(Task):
             else:
                 environ["REMOTE_ADDR"] = client_addr.strip()
 
+        return untrusted_headers
+
     def get_environment(self):
         """Returns a WSGI environment."""
         environ = self.environ
@@ -773,28 +763,40 @@ class WSGITask(Task):
                 if path.startswith(url_prefix_with_trailing_slash):
                     path = path[len(url_prefix):]
 
-        environ = {}
-        environ['REQUEST_METHOD'] = request.command.upper()
-        environ['SERVER_PORT'] = str(server.effective_port)
-        environ['SERVER_NAME'] = server.server_name
-        environ['SERVER_SOFTWARE'] = server.adj.ident
-        environ['SERVER_PROTOCOL'] = 'HTTP/%s' % self.version
-        environ['SCRIPT_NAME'] = url_prefix
-        environ['PATH_INFO'] = path
-        environ['QUERY_STRING'] = request.query
+        environ = {
+            'REQUEST_METHOD': request.command.upper(),
+            'SERVER_PORT': str(server.effective_port),
+            'SERVER_NAME': server.server_name,
+            'SERVER_SOFTWARE': server.adj.ident,
+            'SERVER_PROTOCOL': 'HTTP/%s' % self.version,
+            'SCRIPT_NAME': url_prefix,
+            'PATH_INFO': path,
+            'QUERY_STRING': request.query,
+            'wsgi.url_scheme': request.url_scheme,
+
+            # the following environment variables are required by the WSGI spec
+            'wsgi.version': (1, 0),
+
+            # apps should use the logging module
+            'wsgi.errors': sys.stderr,
+            'wsgi.multithread': True,
+            'wsgi.multiprocess': False,
+            'wsgi.run_once': False,
+            'wsgi.input': request.get_body_stream(),
+            'wsgi.file_wrapper': ReadOnlyFileBasedBuffer,
+            'wsgi.input_terminated': True,  # wsgi.input is EOF terminated
+        }
         remote_peer = environ['REMOTE_ADDR'] = channel.addr[0]
-        environ['wsgi.url_scheme'] = request.url_scheme
 
         headers = dict(request.headers)
 
+        untrusted_headers = PROXY_HEADERS
         if remote_peer == server.adj.trusted_proxy:
-            self.parse_proxy_headers(
+            untrusted_headers = self.parse_proxy_headers(
                 environ,
                 headers,
                 trusted_proxy_count=server.adj.trusted_proxy_count,
                 trusted_proxy_headers=server.adj.trusted_proxy_headers,
-                log_untrusted_proxy_headers=server.adj.log_untrusted_proxy_headers,
-                clear_untrusted_proxy_headers=server.adj.clear_untrusted_proxy_headers,
             )
         else:
             # If we are not relying on a proxy, we still want to try and set
@@ -807,6 +809,15 @@ class WSGITask(Task):
         # so we do.
         environ["REMOTE_HOST"] = environ["REMOTE_ADDR"]
 
+        # Clear out the untrusted proxy headers
+        if server.adj.clear_untrusted_proxy_headers:
+            clear_untrusted_headers(
+                headers,
+                untrusted_headers,
+                log_warning=server.adj.log_untrusted_proxy_headers,
+                logger=self.logger,
+            )
+
         for key, value in headers.items():
             value = value.strip()
             mykey = rename_headers.get(key, None)
@@ -815,17 +826,6 @@ class WSGITask(Task):
             if mykey not in environ:
                 environ[mykey] = value
 
-        # the following environment variables are required by the WSGI spec
-        environ['wsgi.version'] = (1, 0)
-
-        # apps should use the logging module
-        environ['wsgi.errors'] = sys.stderr
-        environ['wsgi.multithread'] = True
-        environ['wsgi.multiprocess'] = False
-        environ['wsgi.run_once'] = False
-        environ['wsgi.input'] = request.get_body_stream()
-        environ['wsgi.file_wrapper'] = ReadOnlyFileBasedBuffer
-        environ['wsgi.input_terminated'] = True  # wsgi.input is EOF terminated
-
+        # cache the environ for this request
         self.environ = environ
         return environ
diff --git a/waitress/tests/test_task.py b/waitress/tests/test_task.py
index f753172..c141d79 100644
--- a/waitress/tests/test_task.py
+++ b/waitress/tests/test_task.py
@@ -854,7 +854,7 @@ class TestWSGITask(unittest.TestCase):
 
     def test_get_environment_values_w_bogus_scheme_override(self):
         inst = self._makeOne()
-        inst.channel.addr = ['192.168.1.1']
+        inst.channel.addr = ['192.168.1.1', 80]
         inst.channel.server.adj.trusted_proxy = '192.168.1.1'
         inst.channel.server.adj.trusted_proxy_headers = {'x-forwarded-proto'}
         request = DummyParser()
@@ -869,6 +869,95 @@ class TestWSGITask(unittest.TestCase):
         inst.request = request
         self.assertRaises(ValueError, inst.get_environment)
 
+    def test_get_environment_warning_other_proxy_headers(self):
+        inst = self._makeOne()
+        inst.logger = DummyLogger()
+
+        inst.request.headers = {
+            'X_FORWARDED_FOR': '[2001:db8::1]',
+            'FORWARDED': 'For=198.51.100.2;host=example.com:8080;proto=https'
+        }
+        inst.channel.addr = ['192.168.1.1', 80]
+        inst.channel.server.adj.trusted_proxy = '192.168.1.1'
+        inst.channel.server.adj.trusted_proxy_count = 1
+        inst.channel.server.adj.trusted_proxy_headers = {'forwarded'}
+        inst.channel.server.adj.log_untrusted_proxy_headers = True
+        environ = inst.get_environment()
+
+        self.assertEqual(len(inst.logger.logged), 1)
+        self.assertNotIn('HTTP_X_FORWARDED_FOR', environ)
+
+        self.assertEqual(environ['REMOTE_ADDR'], '198.51.100.2')
+        self.assertEqual(environ['SERVER_NAME'], 'example.com')
+        self.assertEqual(environ['HTTP_HOST'], 'example.com:8080')
+        self.assertEqual(environ['SERVER_PORT'], '8080')
+        self.assertEqual(environ['wsgi.url_scheme'], 'https')
+
+    def test_get_environment_contains_all_headers_including_untrusted(self):
+        inst = self._makeOne()
+        inst.logger = DummyLogger()
+
+        inst.request.headers = {
+            'X_FORWARDED_FOR': '198.51.100.2',
+            'X_FORWARDED_BY': 'Waitress',
+            'X_FORWARDED_PROTO': 'https',
+            'X_FORWARDED_HOST': 'example.org',
+        }
+        headers_orig = inst.request.headers.copy()
+        inst.channel.addr = ['192.168.1.1', 80]
+        inst.channel.server.adj.trusted_proxy = '192.168.1.1'
+        inst.channel.server.adj.trusted_proxy_count = 1
+        inst.channel.server.adj.trusted_proxy_headers = {'x-forwarded-by'}
+        inst.channel.server.adj.clear_untrusted_proxy_headers = False
+        environ = inst.get_environment()
+
+        for k, expected in headers_orig.items():
+            result = environ['HTTP_%s' % k]
+            self.assertEqual(result, expected)
+
+    def test_get_environment_contains_only_trusted_headers(self):
+        inst = self._makeOne()
+        inst.logger = DummyLogger()
+
+        inst.request.headers = {
+            'X_FORWARDED_FOR': '198.51.100.2',
+            'X_FORWARDED_BY': 'Waitress',
+            'X_FORWARDED_PROTO': 'https',
+            'X_FORWARDED_HOST': 'example.org',
+        }
+        inst.channel.addr = ['192.168.1.1', 80]
+        inst.channel.server.adj.trusted_proxy = '192.168.1.1'
+        inst.channel.server.adj.trusted_proxy_count = 1
+        inst.channel.server.adj.trusted_proxy_headers = {'x-forwarded-by'}
+        inst.channel.server.adj.clear_untrusted_proxy_headers = True
+        environ = inst.get_environment()
+
+        self.assertEqual(environ['HTTP_X_FORWARDED_BY'], 'Waitress')
+        self.assertNotIn('HTTP_X_FORWARDED_FOR', environ)
+        self.assertNotIn('HTTP_X_FORWARDED_PROTO', environ)
+        self.assertNotIn('HTTP_X_FORWARDED_HOST', environ)
+
+    def test_get_environment_clears_headers_if_untrusted_proxy(self):
+        inst = self._makeOne()
+        inst.logger = DummyLogger()
+
+        inst.request.headers = {
+            'X_FORWARDED_FOR': '198.51.100.2',
+            'X_FORWARDED_BY': 'Waitress',
+            'X_FORWARDED_PROTO': 'https',
+            'X_FORWARDED_HOST': 'example.org',
+        }
+        inst.channel.addr = ['192.168.1.255', 80]
+        inst.channel.server.adj.trusted_proxy = '192.168.1.1'
+        inst.channel.server.adj.trusted_proxy_count = 1
+        inst.channel.server.adj.trusted_proxy_headers = {'x-forwarded-by'}
+        inst.channel.server.adj.clear_untrusted_proxy_headers = True
+        environ = inst.get_environment()
+
+        self.assertNotIn('HTTP_X_FORWARDED_BY', environ)
+        self.assertNotIn('HTTP_X_FORWARDED_FOR', environ)
+        self.assertNotIn('HTTP_X_FORWARDED_PROTO', environ)
+        self.assertNotIn('HTTP_X_FORWARDED_HOST', environ)
 
     def test_parse_proxy_headers_forwarded_for(self):
         inst = self._makeOne()
@@ -1105,32 +1194,6 @@ class TestWSGITask(unittest.TestCase):
         self.assertEqual(environ['SERVER_PORT'], '8080')
         self.assertEqual(environ['wsgi.url_scheme'], 'https')
 
-    def test_parse_forwarded_warning_other_proxy_headers(self):
-        inst = self._makeOne()
-        inst.logger = DummyLogger()
-
-        headers = {
-            'X_FORWARDED_FOR': '[2001:db8::1]',
-            'FORWARDED': 'For=198.51.100.2;host=example.com:8080;proto=https'
-        }
-        environ = {}
-        inst.parse_proxy_headers(
-            environ,
-            headers,
-            trusted_proxy_count=1,
-            trusted_proxy_headers={'forwarded'},
-            log_untrusted_proxy_headers=True
-        )
-
-        self.assertEqual(len(inst.logger.logged), 1)
-        self.assertNotIn('X_FORWARDED_FOR', headers)
-
-        self.assertEqual(environ['REMOTE_ADDR'], '198.51.100.2')
-        self.assertEqual(environ['SERVER_NAME'], 'example.com')
-        self.assertEqual(environ['HTTP_HOST'], 'example.com:8080')
-        self.assertEqual(environ['SERVER_PORT'], '8080')
-        self.assertEqual(environ['wsgi.url_scheme'], 'https')
-
     def test_parse_forwarded_empty_pair(self):
         inst = self._makeOne()
 
@@ -1236,50 +1299,6 @@ class TestWSGITask(unittest.TestCase):
 
         self.assertEqual(environ, {})
 
-    def test_parse_headers_contains_only_trusted(self):
-        inst = self._makeOne()
-        inst.logger = DummyLogger()
-
-        headers = {
-            'X_FORWARDED_FOR': '198.51.100.2',
-            'X_FORWARDED_BY': 'Waitress',
-            'X_FORWARDED_PROTO': 'https',
-            'X_FORWARDED_HOST': 'example.org',
-        }
-        environ = {}
-        inst.parse_proxy_headers(
-            environ,
-            headers,
-            trusted_proxy_count=1,
-            trusted_proxy_headers={'x-forwarded-by'}
-        )
-
-        self.assertEqual(environ, {})
-        self.assertEqual(headers, {'X_FORWARDED_BY': 'Waitress'})
-
-    def test_parse_headers_contains_all_including_untrusted(self):
-        inst = self._makeOne()
-        inst.logger = DummyLogger()
-
-        headers = {
-            'X_FORWARDED_FOR': '198.51.100.2',
-            'X_FORWARDED_BY': 'Waitress',
-            'X_FORWARDED_PROTO': 'https',
-            'X_FORWARDED_HOST': 'example.org',
-        }
-        headers_orig = headers.copy()
-        environ = {}
-        inst.parse_proxy_headers(
-            environ,
-            headers,
-            trusted_proxy_count=1,
-            trusted_proxy_headers={'x-forwarded-by'},
-            clear_untrusted_proxy_headers=False
-        )
-
-        self.assertEqual(environ, {})
-        self.assertEqual(headers, headers_orig)
-
     def test_parse_multiple_x_forwarded_proto(self):
         inst = self._makeOne()
         inst.logger = DummyLogger()
@@ -1472,7 +1491,7 @@ class DummyAdj(object):
     url_prefix = ''
     trusted_proxy = None
     trusted_proxy_count = 1
-    trusted_proxy_headers = {}
+    trusted_proxy_headers = set()
     log_untrusted_proxy_headers = True
     clear_untrusted_proxy_headers = True