Resolve client address from Forwarded and X-Forwarded-* headers

packages/kaya-core/tests/test_asgi.py

@@ -62,6 +62,11 @@ class AsgiTest(unittest.TestCase):
         async def handle_request(ctx: HttpContext, _: List[str]) -> None:
             await ctx.stream_body(200, (chunk async for chunk in ctx.request_body))
 
+        @self.app.GET('/client-ip')
+        async def handle_request(ctx: HttpContext) -> None:
+            host, port = ctx.client if ctx.client is not None else (None, None)
+            await ctx.send_str(200, json.dumps({'host': host, 'port': port}))
+
     @async_test
     async def test_hello(self):
         transport = httpx.ASGITransport(app=self.app)
@@ -190,6 +195,60 @@ class AsgiTest(unittest.TestCase):
                 'employee_id': 101325
             }, response)
 
+    @async_test
+    async def test_client_ip_forwarded(self):
+        transport = httpx.ASGITransport(app=self.app)
+
+        async with httpx.AsyncClient(transport=transport, base_url="http://127.0.0.1:80") as client:
+            # socket peer, no forwarded headers
+            r = await client.get("/client-ip")
+            socket_client = json.loads(r.text)
+            self.assertEqual('127.0.0.1', socket_client['host'])
+
+            # RFC 7239 Forwarded header, with port
+            r = await client.get("/client-ip", headers={'Forwarded': 'for=203.0.113.5:1234'})
+            self.assertEqual({'host': '203.0.113.5', 'port': 1234}, json.loads(r.text))
+
+            # RFC 7239 Forwarded header, bracketed IPv6 with port
+            r = await client.get("/client-ip", headers={'Forwarded': 'for="[2001:db8::1]:4711"'})
+            self.assertEqual({'host': '2001:db8::1', 'port': 4711}, json.loads(r.text))
+
+            # RFC 7239 Forwarded header without port keeps the socket port
+            r = await client.get("/client-ip", headers={'Forwarded': 'for=203.0.113.5'})
+            self.assertEqual({'host': '203.0.113.5', 'port': socket_client['port']}, json.loads(r.text))
+
+            # Forwarded with for=unknown falls through to X-Forwarded-For
+            r = await client.get("/client-ip", headers={
+                'Forwarded': 'for=unknown',
+                'X-Forwarded-For': '198.51.100.7',
+            })
+            self.assertEqual('198.51.100.7', json.loads(r.text)['host'])
+
+            # Forwarded takes precedence over X-Forwarded-For
+            r = await client.get("/client-ip", headers={
+                'Forwarded': 'for=203.0.113.5',
+                'X-Forwarded-For': '198.51.100.7',
+            })
+            self.assertEqual('203.0.113.5', json.loads(r.text)['host'])
+
+            # X-Forwarded-For: first entry of the chain, port from X-Forwarded-Port
+            r = await client.get("/client-ip", headers={
+                'X-Forwarded-For': '203.0.113.5, 70.41.3.18',
+                'X-Forwarded-Port': '8443',
+            })
+            self.assertEqual({'host': '203.0.113.5', 'port': 8443}, json.loads(r.text))
+
+            # X-Forwarded-Host fallback
+            r = await client.get("/client-ip", headers={'X-Forwarded-Host': '198.51.100.7'})
+            self.assertEqual('198.51.100.7', json.loads(r.text)['host'])
+
+            # invalid X-Forwarded-Port is ignored, socket port is kept
+            r = await client.get("/client-ip", headers={
+                'X-Forwarded-For': '203.0.113.5',
+                'X-Forwarded-Port': 'not-a-port',
+            })
+            self.assertEqual({'host': '203.0.113.5', 'port': socket_client['port']}, json.loads(r.text))
+
     @async_test
     async def test_nested_param_routes(self):
         app = KayaApp()

packages/kaya-core/tests/test_websocket.py

@@ -132,6 +132,37 @@ class WebSocketTest(unittest.TestCase):
         self.assertEqual(1, len(sent_messages))
         self.assertEqual({'type': 'websocket.accept'}, sent_messages[0])
 
+    @async_test
+    async def test_client_forwarded_header(self):
+        async def send(message):
+            pass
+
+        async def receive():
+            return {'type': 'websocket.connect'}
+
+        scope = {
+            'type': 'websocket',
+            'path': '/echo',
+            'query_string': b'',
+            'scheme': 'ws',
+            'client': ('127.0.0.1', 12345),
+            'server': ('127.0.0.1', 80),
+            'headers': [(b'forwarded', b'for=203.0.113.5:1234')],
+        }
+        ws = AsgiWebSocket(scope, receive, send)
+        self.assertEqual(('203.0.113.5', 1234), ws.client)
+
+        scope['headers'] = [
+            (b'x-forwarded-for', b'203.0.113.5, 70.41.3.18'),
+            (b'x-forwarded-port', b'8443'),
+        ]
+        ws = AsgiWebSocket(scope, receive, send)
+        self.assertEqual(('203.0.113.5', 8443), ws.client)
+
+        scope['headers'] = []
+        ws = AsgiWebSocket(scope, receive, send)
+        self.assertEqual(('127.0.0.1', 12345), ws.client)
+
     @async_test
     async def test_websocket_scope_without_scheme(self):
         # Daphne omits the optional `scheme` key from websocket scopes.