Fix ProxyHeadersMiddleware, now it grabs the first IP from `x-forward…

…ed-for`

Also rewrite a few tests for increasing coverage for edge cases

tests/middleware/test_proxy_headers.py

@@ -1,3 +1,5 @@
+import pytest
+
 from tests.client import TestClient
 from tests.response import Response
 from uvicorn.middleware.proxy_headers import ProxyHeadersMiddleware
@@ -11,20 +13,38 @@ async def app(scope, receive, send):
     await response(scope, receive, send)
 
 
-app = ProxyHeadersMiddleware(app, trusted_hosts="*")
-
-
-def test_proxy_headers():
-    client = TestClient(app)
+@pytest.mark.parametrize(
+    ("trusted_hosts", "response_text"),
+    [
+        # always trust
+        ("*", "Remote: https://1.2.3.4:0"),
+        # trusted proxy
+        ("testclient", "Remote: https://1.2.3.4:0"),
+        (["testclient"], "Remote: https://1.2.3.4:0"),
+        # trusted proxy list
+        (["testclient", "127.0.0.1"], "Remote: https://1.2.3.4:0"),
+        # trusted proxy list
+        ("testclient, 127.0.0.1", "Remote: https://1.2.3.4:0"),
+        # request from untrusted proxy
+        ("121.0.0.1", "Remote: http://testclient:50000"),
+    ],
+)
+def test_proxy_headers_trusted_hosts(trusted_hosts, response_text):
+    app_with_middleware = ProxyHeadersMiddleware(app, trusted_hosts=trusted_hosts)
+    client = TestClient(app_with_middleware)
     headers = {"X-Forwarded-Proto": "https", "X-Forwarded-For": "1.2.3.4"}
     response = client.get("/", headers=headers)
     assert response.status_code == 200
-    assert response.text == "Remote: https://1.2.3.4:0"
+    assert response.text == response_text
 
 
-def test_proxy_headers_no_port():
-    client = TestClient(app)
-    headers = {"X-Forwarded-Proto": "https", "X-Forwarded-For": "1.2.3.4"}
+def test_proxy_headers_multiple_proxies():
+    app_with_middleware = ProxyHeadersMiddleware(app, trusted_hosts="*")
+    client = TestClient(app_with_middleware)
+    headers = {
+        "X-Forwarded-Proto": "https",
+        "X-Forwarded-For": "1.2.3.4, 10.0.2.1, 192.168.0.2",
+    }
     response = client.get("/", headers=headers)
     assert response.status_code == 200
     assert response.text == "Remote: https://1.2.3.4:0"

tests/test_main.py

@@ -1,6 +1,6 @@
+import asyncio
 import threading
 import time
-import asyncio
 
 import requests
 

uvicorn/middleware/proxy_headers.py

@@ -34,11 +34,11 @@ async def __call__(self, scope, receive, send):
                     scope["scheme"] = x_forwarded_proto.strip()
 
                 if b"x-forwarded-for" in headers:
-                    # Determine the client address from the last trusted IP in the
+                    # Determine the client address from the first trusted IP in the
                     # X-Forwarded-For header. We've lost the connecting client's port
                     # information by now, so only include the host.
                     x_forwarded_for = headers[b"x-forwarded-for"].decode("ascii")
-                    host = x_forwarded_for.split(",")[-1].strip()
+                    host = x_forwarded_for.split(",")[0].strip()
                     port = 0
                     scope["client"] = (host, port)