From 8a07b349ee46ec04d0f20b4aed39bd689773bb10 Mon Sep 17 00:00:00 2001
From: Jens Langhammer <jens.langhammer@beryju.org>
Date: Thu, 10 Sep 2020 16:58:25 +0200
Subject: [PATCH] root: fix IP detection in ASGI logger, attempt to fix out of
 order issues

---
 Makefile              |  6 ++++++
 passbook/root/asgi.py | 44 ++++++++++++++++++++++++-------------------
 2 files changed, 31 insertions(+), 19 deletions(-)

diff --git a/Makefile b/Makefile
index ffb6150f2..86435e31c 100644
--- a/Makefile
+++ b/Makefile
@@ -18,3 +18,9 @@ lint:
 
 gen: coverage
 	./manage.py generate_swagger -o swagger.yaml -f yaml
+
+local-stack:
+	export PASSBOOK_TAG=testing
+	docker build -t beryju/passbook:testng .
+	docker-compose up -d
+	docker-compose run --rm server migrate
diff --git a/passbook/root/asgi.py b/passbook/root/asgi.py
index 8954e577f..a2ba3a6c0 100644
--- a/passbook/root/asgi.py
+++ b/passbook/root/asgi.py
@@ -32,6 +32,11 @@ Send = typing.Callable[[Message], typing.Awaitable[None]]
 
 ASGIApp = typing.Callable[[Scope, Receive, Send], typing.Awaitable[None]]
 
+ASGI_IP_HEADERS = (
+    b"x-forwarded-for",
+    b"x-real-ip",
+)
+
 LOGGER = get_logger("passbook.asgi")
 
 
@@ -51,7 +56,6 @@ class ASGILogger:
     """ASGI Logger, instantiated for each request"""
 
     app: ASGIApp
-    send: Send
 
     scope: Scope
     headers: Dict[ByteString, Any]
@@ -64,11 +68,26 @@ class ASGILogger:
         self.app = app
 
     async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
-        self.send = send
         self.scope = scope
         self.content_length = 0
         self.headers = dict(scope.get("headers", []))
 
+        async def send_hooked(message: Message) -> None:
+            """Hooked send method, which records status code and content-length, and for the final
+            requests logs it"""
+            headers = dict(message.get("headers", []))
+
+            if "status" in message:
+                self.status_code = message["status"]
+
+            if b"Content-Length" in headers:
+                self.content_length += int(headers.get(b"Content-Length", b"0"))
+
+            if message["type"] == "http.response.body" and not message["more_body"]:
+                runtime = int((time() - self.start) * 10 ** 6)
+                self.log(runtime)
+            await send(message)
+
         if self.headers.get(b"host", b"") == b"kubernetes-healthcheck-host":
             # Don't log kubernetes health/readiness requests
             await send({"type": "http.response.start", "status": 204, "headers": []})
@@ -80,25 +99,12 @@ class ASGILogger:
             # https://code.djangoproject.com/ticket/31508
             # https://github.com/encode/uvicorn/issues/266
             return
-        await self.app(scope, receive, self.send_hooked)
-
-    async def send_hooked(self, message: Message) -> None:
-        """Hooked send method, which records status code and content-length, and for the final
-        requests logs it"""
-        headers = dict(message.get("headers", []))
-
-        if "status" in message:
-            self.status_code = message["status"]
-
-        if b"Content-Length" in headers:
-            self.content_length += int(headers.get(b"Content-Length", b"0"))
-
-        if message["type"] == "http.response.body" and not message["more_body"]:
-            runtime = int((time() - self.start) * 10 ** 6)
-            self.log(runtime)
-        return await self.send(message)
+        await self.app(scope, receive, send_hooked)
 
     def _get_ip(self) -> str:
+        for header in ASGI_IP_HEADERS:
+            if header in self.headers:
+                return self.headers[header].decode()
         client_ip, _ = self.scope.get("client", ("", 0))
         return client_ip