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

packages/kaya-rsgi/src/kaya/rsgi/_rsgi.py

@@ -23,7 +23,7 @@ from granian._granian import (  # type: ignore[attr-defined]
 )
 from pwo import Maybe
 
-from kaya.core import AbstractKayaApp, HttpContext, HttpMethod, WebSocket, WebSocketMessage
+from kaya.core import AbstractKayaApp, HttpContext, HttpMethod, WebSocket, WebSocketMessage, resolve_client
 from kaya.core._types import StrOrStrings
 
 
@@ -51,9 +51,10 @@ class RsgiContext(HttpContext):
 
         fun = cast(Callable[[Mapping[str, Sequence[str]], tuple[str, str]], Mapping[str, Sequence[str]]], acc)
         self.headers = reduce(fun, scope.headers.items(), {})
-        self.client = (Maybe.of(scope.client.split(':'))
-                       .map(lambda it: (it[0], int(it[1])))
-                       .or_else_throw(RuntimeError))
+        self.client = resolve_client(self.headers,
+                                     Maybe.of(scope.client.split(':'))
+                                          .map(lambda it: (it[0], int(it[1])))
+                                          .or_else_throw(RuntimeError))
         self.server = (Maybe.of(scope.server.split(':'))
                        .map(lambda it: (it[0], int(it[1])))
                        .or_else_throw(RuntimeError))
@@ -134,9 +135,10 @@ class RsgiWebSocket(WebSocket):
 
         fun = cast(Callable[[Mapping[str, Sequence[str]], tuple[str, str]], Mapping[str, Sequence[str]]], acc)
         self.headers = reduce(fun, scope.headers.items(), {})
-        self.client = (Maybe.of(scope.client.split(':'))
-                       .map(lambda it: (it[0], int(it[1])))
-                       .or_else_throw(RuntimeError))
+        self.client = resolve_client(self.headers,
+                                     Maybe.of(scope.client.split(':'))
+                                          .map(lambda it: (it[0], int(it[1])))
+                                          .or_else_throw(RuntimeError))
         self.server = (Maybe.of(scope.server.split(':'))
                        .map(lambda it: (it[0], int(it[1])))
                        .or_else_throw(RuntimeError))

packages/kaya-rsgi/tests/test_rsgi.py

@@ -1,5 +1,5 @@
 import unittest
-from kaya.rsgi import RsgiWebSocket
+from kaya.rsgi import RsgiContext, RsgiWebSocket
 
 
 class RsgiWebSocketTest(unittest.TestCase):
@@ -20,3 +20,40 @@ class RsgiWebSocketTest(unittest.TestCase):
             RsgiWebSocket(FakeScope(), FakeProtocol())  # type: ignore[arg-type]
 
         self.assertIn('Granian was not configured for websockets', str(ctx.exception))
+
+
+class RsgiContextTest(unittest.TestCase):
+
+    @staticmethod
+    def _make_context(headers):
+        class FakeScope:
+            scheme = 'http'
+            method = 'GET'
+            path = '/'
+            query_string = ''
+            client = '127.0.0.1:12345'
+            server = '127.0.0.1:80'
+
+            def __init__(self, headers):
+                self.headers = headers
+
+        return RsgiContext(FakeScope(headers), object())  # type: ignore[arg-type]
+
+    def test_forwarded_header(self):
+        ctx = self._make_context({'forwarded': 'for=203.0.113.5:1234'})
+        self.assertEqual(('203.0.113.5', 1234), ctx.client)
+
+    def test_x_forwarded_headers(self):
+        ctx = self._make_context({
+            'x-forwarded-for': '203.0.113.5, 70.41.3.18',
+            'x-forwarded-port': '8443',
+        })
+        self.assertEqual(('203.0.113.5', 8443), ctx.client)
+
+    def test_x_forwarded_host_fallback(self):
+        ctx = self._make_context({'x-forwarded-host': '198.51.100.7'})
+        self.assertEqual(('198.51.100.7', 12345), ctx.client)
+
+    def test_socket_peer_fallback(self):
+        ctx = self._make_context({})
+        self.assertEqual(('127.0.0.1', 12345), ctx.client)