From c814f6fc6bc67a0fe588482f52820956c71f641c Mon Sep 17 00:00:00 2001
From: Thomas Carmet <thomas.carmet@scality.com>
Date: Wed, 9 Aug 2023 02:46:33 +0000
Subject: [PATCH] Add middleware for healthcheck

---
 runner_manager/auth.py            | 16 ++++++++++++++++
 runner_manager/main.py            | 15 ++++++++++-----
 runner_manager/models/settings.py |  6 +++++-
 runner_manager/routers/_health.py | 10 ++++++++++
 tests/api/tests_auth.py           |  1 -
 5 files changed, 41 insertions(+), 7 deletions(-)
 create mode 100644 runner_manager/auth.py
 create mode 100644 runner_manager/routers/_health.py

diff --git a/runner_manager/auth.py b/runner_manager/auth.py
new file mode 100644
index 00000000..eceb6622
--- /dev/null
+++ b/runner_manager/auth.py
@@ -0,0 +1,16 @@
+
+from fastapi.middleware.trustedhost import TrustedHostMiddleware
+from starlette.types import Receive, Scope, Send
+
+
+class TrustedHostHealthRoutes(TrustedHostMiddleware):
+    """A healthcheck endpoint that answers to GET requests on /_health"""
+
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        """If the request is made on the path _health then execute the check on hosts"""
+        if scope["path"] == "/_health/":
+            print(scope)
+            await super().__call__(scope, receive, send)
+        else:
+            await self.app(scope, receive, send)
diff --git a/runner_manager/main.py b/runner_manager/main.py
index bae9b695..8f9c512b 100644
--- a/runner_manager/main.py
+++ b/runner_manager/main.py
@@ -1,18 +1,20 @@
 import logging
 
-from fastapi import Depends, FastAPI, HTTPException, Response, Security, status
+from fastapi import Depends, FastAPI, HTTPException, Security, status
 from fastapi.security import APIKeyHeader, APIKeyQuery
+from runner_manager.auth import TrustedHostHealthRoutes
 
 from runner_manager.dependencies import get_queue, get_settings
 from runner_manager.jobs.startup import startup
 from runner_manager.models.settings import Settings
-from runner_manager.routers import webhook
+from runner_manager.routers import webhook, _health
 
 log = logging.getLogger(__name__)
 
 app = FastAPI()
 api_key_query = APIKeyQuery(name="api-key", auto_error=False)
 api_key_header = APIKeyHeader(name="x-api-key", auto_error=False)
+settings = get_settings()
 
 
 def get_api_key(
@@ -20,6 +22,7 @@ def get_api_key(
     api_key_header: str = Security(api_key_header),
     settings: Settings = Depends(get_settings),
 ) -> str:
+    """Get the API key from either the query parameter or the header"""
     if not settings.api_key:
         return ""
     if api_key_query in [settings.api_key.get_secret_value()]:
@@ -33,6 +36,11 @@ def get_api_key(
 
 
 app.include_router(webhook.router)
+app.include_router(_health.router)
+app.add_middleware(
+    TrustedHostHealthRoutes,
+    allowed_hosts=settings.allowed_hosts
+)
 
 
 @app.on_event("startup")
@@ -43,9 +51,6 @@ def startup_event():
     log.info(f"Startup job is {status}")
 
 
-@app.get("/_health")
-def health():
-    return Response(status_code=200)
 
 
 @app.get("/public")
diff --git a/runner_manager/models/settings.py b/runner_manager/models/settings.py
index 7ad29351..d720d6d8 100644
--- a/runner_manager/models/settings.py
+++ b/runner_manager/models/settings.py
@@ -1,5 +1,5 @@
 from pathlib import Path
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Optional, Sequence
 
 import yaml
 from pydantic import AnyHttpUrl, BaseSettings, RedisDsn, SecretStr
@@ -26,6 +26,10 @@ class Settings(BaseSettings):
     redis_om_url: Optional[RedisDsn] = None
     github_base_url: Optional[AnyHttpUrl] = None
     api_key: Optional[SecretStr] = None
+    allowed_hosts: Optional[Sequence[str]] = [
+        "localhost",
+        "testserver",
+    ]
 
     class Config:
         smart_union = True
diff --git a/runner_manager/routers/_health.py b/runner_manager/routers/_health.py
new file mode 100644
index 00000000..52a50bfb
--- /dev/null
+++ b/runner_manager/routers/_health.py
@@ -0,0 +1,10 @@
+
+from fastapi import APIRouter, Response
+
+router = APIRouter(prefix="/_health")
+
+
+@router.get("/", status_code=200)
+def healthcheck():
+    """Healthcheck endpoint that answers to GET requests on /_healthz"""
+    return Response(status_code=200)
diff --git a/tests/api/tests_auth.py b/tests/api/tests_auth.py
index 6d11d36c..b68d19ee 100644
--- a/tests/api/tests_auth.py
+++ b/tests/api/tests_auth.py
@@ -27,5 +27,4 @@ def test_private_endpoint_with_valid_api_key(fastapp, client):
     fastapp.dependency_overrides[get_settings] = settings_api_key
     headers = {"x-api-key": "secret"}
     response = client.get("/private", headers=headers)
-    print(response.text)
     assert response.status_code == 200