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

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)