feat(umami): try to fix country tracking with headers and trusted pro…

…xies

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "umami-asgi"
-version = "0.2.0"
+version = "0.2.1"
 description = "A middleware for your ASGI application that enables Umami analytics."
 readme = "README.md"
 license = "MIT"

umami_asgi/middleware.py

@@ -7,15 +7,21 @@
 from starlette.responses import Response
 from starlette.background import BackgroundTask
 from dataclasses import asdict
+from typing import List
 
 
-async def send_umami_payload(api_endpoint: str, request_payload: UmamiRequest, headers: MutableHeaders,
-                             follow_redirects: bool):
+async def send_umami_payload(
+        api_endpoint: str,
+        request_payload: UmamiRequest,
+        headers: MutableHeaders,
+        follow_redirects: bool
+):
     async with httpx.AsyncClient() as client:
         try:
-            resp = await client.post(
+            payload = asdict(request_payload)
+            await client.post(
                 api_endpoint,
-                json=asdict(request_payload),
+                json=payload,
                 headers=headers,
                 follow_redirects=follow_redirects
             )
@@ -31,6 +37,8 @@ def __init__(self,
                  api_endpoint: str,
                  website_id: str,
                  follow_redirects: bool = True,
+                 proxy_enabled: bool = False,
+                 trusted_proxies: List[str] = None,
                  ) -> None:
         super().__init__(app)
         self.app = app
@@ -40,6 +48,11 @@ def __init__(self,
         self.website_id = website_id
         self.follow_redirects = follow_redirects
 
+        if proxy_enabled and trusted_proxies:
+            self.trusted_proxies = set(trusted_proxies)
+        else:
+            self.trusted_proxies = None
+
     async def dispatch(self, request: Request, call_next) -> Response:
         response = await call_next(request)
 
@@ -53,17 +66,35 @@ async def dispatch(self, request: Request, call_next) -> Response:
                 title='',
                 url=request.url.path,
                 website=self.website_id,
+                ip=request.headers.get('X-Real-IP', request.client.host)
             )
         )
 
         # set headers to track IP address correctly
         umami_headers = MutableHeaders()
         # send proper user agent or request won't be registered
-        umami_headers['User-Agent'] = request.headers['User-Agent']
+        umami_headers['User-Agent'] = request.headers.get('User-Agent', '')
         # send IP address of client, not asgi server
-        umami_headers['X-Real-IP'] = request.client.host
-        umami_headers['X-Forwarded-For'] = request.client.host
-        response.background = BackgroundTask(send_umami_payload, self.api_endpoint, umami_request, umami_headers,
-                                             self.follow_redirects)
+        umami_headers['X-Forwarded-Proto'] = 'https'
+
+        # if trusted proxies is enabled and the sender's IP is in our trusted proxy list, we can trust the headers
+        if self.trusted_proxies and (request.client.host in self.trusted_proxies or "0.0.0.0" in self.trusted_proxies):
+            umami_headers['X-Real-IP'] = request.headers.get('X-Real-IP', request.client.host)
+            umami_headers['X-Forwarded-For'] = request.headers.get('X-Forwarded-For', request.client.host)
+            umami_headers['X-Forwarded-Host'] = request.headers.get('X-Forwarded-Host', request.url.hostname)
+
+        else:
+            # check against trusted proxies is disabled, so we have to assume that an attacker sends any headers
+            # -> do not use any headers
+            umami_headers['X-Real-IP'] = request.client.host
+            umami_headers['X-Forwarded-For'] = request.client.host
+            umami_headers['X-Forwarded-Host'] = ""
+
+        response.background = BackgroundTask(
+            send_umami_payload,
+            self.api_endpoint,
+            umami_request, umami_headers,
+            self.follow_redirects
+        )
 
         return response

umami_asgi/middleware_test.py

@@ -39,8 +39,8 @@ async def test_analytics_post(httpx_mock: HTTPXMock):
         # check whether middleware sent post request to mocked umami
         umami_request = httpx_mock.get_request()
         umami_payload = loads(umami_request.content.decode("utf-8"))["payload"]
-        assert umami_payload["name"] == "GET"
         assert umami_payload["url"] == "/"
+        assert umami_payload["ip"] == "testclient"
         # check whether middleware set headers correctly to track IP address of client, not asgi server
         assert umami_request.headers["X-Real-IP"] == "testclient"
         assert umami_request.headers["X-Forwarded-For"] == "testclient"
@@ -50,5 +50,60 @@ async def test_analytics_post(httpx_mock: HTTPXMock):
 
         umami_request = httpx_mock.get_requests()[1]
         umami_payload = loads(umami_request.content.decode("utf-8"))["payload"]
-        assert umami_payload["name"] == "POST"
         assert umami_payload["url"] == "/feed"
+
+
+@pytest.mark.asyncio
+async def test_trusted_proxies(httpx_mock: HTTPXMock):
+    httpx_mock.add_response()
+    app = Starlette(
+        routes=[
+            Route("/", endpoint=homepage),
+            Route("/feed", endpoint=feed, methods=["GET", "POST"]),
+        ],
+        middleware=[
+            Middleware(UmamiMiddleware, api_endpoint="https://localhost/api", website_id="123456", proxy_enabled=True,
+                       trusted_proxies=[])
+        ])
+
+    # fake forwarding headers
+    with TestClient(app, base_url="http://testserver", headers={
+        "X-Real-IP": "123.123.123.123",
+        "X-Forwarded-For": "111.111.111.111, 222.222.222.222",
+        "X-Forwarded-Host": "example.com",
+    }) as client:
+        response = client.get("/")
+        assert response.status_code == 200
+
+        umami_request = httpx_mock.get_request()
+        h = umami_request.headers
+        assert h["X-Real-IP"] == "testclient"
+        assert h["X-Forwarded-For"] == "testclient"
+        assert h["X-Forwarded-Host"] == ""
+
+    app_trusted = Starlette(
+        routes=[
+            Route("/", endpoint=homepage),
+            Route("/feed", endpoint=feed, methods=["GET", "POST"]),
+        ],
+        middleware=[
+            Middleware(UmamiMiddleware, api_endpoint="https://localhost/api", website_id="123456", proxy_enabled=True,
+                       trusted_proxies=["testclient"])
+        ])
+
+    httpx_mock.reset(True)
+
+    with TestClient(app_trusted, base_url="http://testserver", headers={
+        "X-Real-IP": "123.123.123.123",
+        "X-Forwarded-For": "111.111.111.111, 222.222.222.222",
+        "X-Forwarded-Host": "example.com",
+    }) as client:
+        response = client.get("/")
+        assert response.status_code == 200
+
+        umami_request = httpx_mock.get_request()
+        h = umami_request.headers
+        # see also https://stackoverflow.com/questions/72557636/difference-between-x-forwarded-for-and-x-real-ip-headers
+        assert h["X-Real-IP"] == "123.123.123.123"
+        assert h["X-Forwarded-For"] == "111.111.111.111, 222.222.222.222"
+        assert h["X-Forwarded-Host"] == "example.com"

umami_asgi/umami_api.py

@@ -12,6 +12,7 @@ class UmamiPayload:
     title: str
     url: str
     website: str
+    ip: str
     # name: str
     # currently not used -> asdict() outputs "None" if not set, has to be prevented if this field is used
     # data: Optional[dict] = None