From cc1a4fac232c9258a4c540184bcd91d590f74975 Mon Sep 17 00:00:00 2001
From: Andy Freeland <andy@andyfreeland.net>
Date: Mon, 15 Nov 2021 22:11:44 -0800
Subject: [PATCH] Check application health over HTTP instead of TCP

This may also fix a memory leak:
https://github.com/encode/uvicorn/issues/1226

Closes #35.
---
 fly.toml                       |  6 ++++--
 src/wanikani_apprentice/app.py | 18 ++++++++++++++++++
 tests/test_app.py              | 14 ++++++++++++++
 3 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/fly.toml b/fly.toml
index 6044dba6..d7d4a5cb 100644
--- a/fly.toml
+++ b/fly.toml
@@ -24,7 +24,6 @@ processes = []
   url_prefix = "/static"
 
 [[services]]
-  http_checks = []
   internal_port = 8080
   processes = ["app"]
   protocol = "tcp"
@@ -43,7 +42,10 @@ processes = []
     handlers = ["tls", "http"]
     port = 443
 
-  [[services.tcp_checks]]
+  [[services.http_checks]]
     grace_period = "10s"
     interval = "10s"
     timeout = "2s"
+    method = "get"
+    path = "/__lbheartbeat__"
+    protocol = "http"
diff --git a/src/wanikani_apprentice/app.py b/src/wanikani_apprentice/app.py
index e64827f6..8ac708ef 100644
--- a/src/wanikani_apprentice/app.py
+++ b/src/wanikani_apprentice/app.py
@@ -3,6 +3,7 @@
 import operator
 import os.path
 
+import attr
 import httpx
 import sentry_sdk
 from sentry_sdk.integrations.asgi import SentryAsgiMiddleware
@@ -18,6 +19,10 @@
 from starlette.routing import Route
 from starlette.staticfiles import StaticFiles
 from starlette.templating import _TemplateResponse
+from starlette.types import ASGIApp
+from starlette.types import Receive
+from starlette.types import Scope
+from starlette.types import Send
 import structlog
 
 from . import config
@@ -36,6 +41,18 @@
 log = structlog.get_logger()
 
 
+@attr.frozen
+class LBHeartbeatMiddleware:
+    app: ASGIApp
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        if scope["type"] == "http" and scope["path"] == "/__lbheartbeat__":
+            response = Response("OK")
+            await response(scope, receive, send)
+        else:
+            await self.app(scope, receive, send)
+
+
 async def index(request: Request) -> RedirectResponse:
     if is_logged_in(request):
         return RedirectResponse(request.url_for("assignments"))
@@ -143,6 +160,7 @@ def create_app() -> Starlette:
     )
     middleware = [
         Middleware(SentryAsgiMiddleware),
+        Middleware(LBHeartbeatMiddleware),
     ]
 
     api = WaniKaniAPIClient(str(config.WANIKANI_API_KEY), client=httpx_client)
diff --git a/tests/test_app.py b/tests/test_app.py
index 245a194e..a5c4a19f 100644
--- a/tests/test_app.py
+++ b/tests/test_app.py
@@ -227,3 +227,17 @@ def test_https_only(mocker):
     resp = test_client.get("/", allow_redirects=False)
     assert resp.status_code == 307
     assert resp.headers["Location"].startswith("https://")
+
+
+def test_lbheartbeat_bypass_https_only(mocker):
+    from wanikani_apprentice import config
+    from wanikani_apprentice.app import create_app
+
+    mocker.patch.object(config, "HTTPS_ONLY", True)
+
+    app = create_app()
+    test_client = TestClient(app)
+
+    resp = test_client.get("/__lbheartbeat__", allow_redirects=False)
+    assert resp.status_code == 200
+    assert resp.text == "OK"