Add forwarded headers wrappers (#620)

README.md

@@ -408,6 +408,25 @@ Given you specify N threads with the relevant option, in **st** mode Granian wil
 
 Benchmarks suggests **st** mode to be more efficient with a small amount of processes, while **mt** mode seems to scale more efficiently where you have a large number of CPUs. Real performance will though depend on specific application code, and thus *your mileage might vary*.
 
+### Proxies and forwarded headers
+
+Since none of the supported applications protocols define a strategy for proxies' *forwarded headers*, Granian doesn't provide any option to configure its behaviour around them.
+
+What Granian provides instead, for contexts in which is being run behind a reverse proxy, are *wrappers* you can use on top of your application, in order to alter the request scope based on the headers forwarded by the proxy itself. You can find such wrappers in the `granian.utils.proxies` module:
+
+```python
+from granian.utils.proxies import wrap_asgi_with_proxy_headers, wrap_wsgi_with_proxy_headers
+
+async def my_asgi_app(scope, receive, send):
+    ...
+
+def my_wsgi_app(environ, start_response):
+    ...
+
+my_asgi_app = wrap_asgi_with_proxy_headers(my_asgi_app, trusted_hosts="1.2.3.4")
+my_wsgi_app = wrap_wsgi_with_proxy_headers(my_wsgi_app, trusted_hosts="1.2.3.4")
+```
+
 ## Free-threaded Python
 
 > **Warning:** free-threaded Python support is still experimental and highly discouraged in *production environments*.

granian/utils/proxies.py

@@ -0,0 +1,107 @@
+import ipaddress
+from functools import wraps as _wraps
+
+
+class _Forwarders:
+    def __init__(self, trusted_hosts: list[str] | str) -> None:
+        self.always_trust: bool = trusted_hosts in ('*', ['*'])
+        self.literals: set[str] = set()
+        self.hosts: set[ipaddress.IPv4Address | ipaddress.IPv6Address] = set()
+        self.networks: set[ipaddress.IPv4Network | ipaddress.IPv6Network] = set()
+
+        if self.always_trust:
+            return
+
+        if isinstance(trusted_hosts, str):
+            trusted_hosts = _parse_raw_hosts(trusted_hosts)
+
+        for host in trusted_hosts:
+            try:
+                if '/' in host:
+                    self.networks.add(ipaddress.ip_network(host))
+                    continue
+                self.hosts.add(ipaddress.ip_address(host))
+            except ValueError:
+                self.literals.add(host)
+
+    def __contains__(self, host: str | None) -> bool:
+        if self.always_trust:
+            return True
+
+        if not host:
+            return False
+
+        try:
+            ip = ipaddress.ip_address(host)
+            if ip in self.hosts:
+                return True
+            return any(ip in net for net in self.networks)
+        except ValueError:
+            return host in self.literals
+
+    def get_client_host(self, x_forwarded_for: str) -> str:
+        x_forwarded_for_hosts = _parse_raw_hosts(x_forwarded_for)
+
+        if self.always_trust:
+            return x_forwarded_for_hosts[0]
+
+        for host in reversed(x_forwarded_for_hosts):
+            if host not in self:
+                return host
+
+        return x_forwarded_for_hosts[0]
+
+
+def _parse_raw_hosts(value: str) -> list[str]:
+    return [item.strip() for item in value.split(',')]
+
+
+def wrap_asgi_with_proxy_headers(app, trusted_hosts: list[str] | str = '127.0.0.1'):
+    forwarders = _Forwarders(trusted_hosts)
+
+    @_wraps(app)
+    def wrapped(scope, receive, send):
+        if scope['type'] == 'lifespan':
+            return app(scope, receive, send)
+
+        client_addr = scope.get('client')
+        client_host = client_addr[0] if client_addr else None
+
+        if client_host in forwarders:
+            headers = dict(scope['headers'])
+
+            if x_forwarded_proto := headers.get(b'x-forwarded-proto', b'').decode('latin1').strip():
+                if x_forwarded_proto in {'http', 'https', 'ws', 'wss'}:
+                    if scope['type'] == 'websocket':
+                        scope['scheme'] = x_forwarded_proto.replace('http', 'ws')
+                    else:
+                        scope['scheme'] = x_forwarded_proto
+
+            if x_forwarded_for := headers.get(b'x-forwarded-for', b'').decode('latin1'):
+                if host := forwarders.get_client_host(x_forwarded_for):
+                    scope['client'] = (host, 0)
+
+        return app(scope, receive, send)
+
+    return wrapped
+
+
+def wrap_wsgi_with_proxy_headers(app, trusted_hosts: list[str] | str = '127.0.0.1'):
+    forwarders = _Forwarders(trusted_hosts)
+
+    @_wraps(app)
+    def wrapped(scope, resp):
+        client_host = scope.get('REMOTE_ADDR')
+
+        if client_host in forwarders:
+            if x_forwarded_proto := scope.get('HTTP_X_FORWARDED_PROTO'):
+                if x_forwarded_proto in {'http', 'https'}:
+                    scope['wsgi.url_scheme'] = x_forwarded_proto
+
+            if x_forwarded_for := scope.get('HTTP_X_FORWARDED_FOR'):
+                if host := forwarders.get_client_host(x_forwarded_for):
+                    scope['REMOTE_ADDR'] = host
+
+        return app(scope, resp)
+
+    return wrapped