Add redirect feature based on proto headers

Starlette/FastAPI don't support https redirect
behind the proxy.

app/core/settings.py

@@ -149,6 +149,12 @@ class FeatureFlags(BaseModel):
         required=False,
         default=False,
     )
+    enable_https_redirect: bool = get_env_var(
+        "ENABLE_HTTPS_REDIRECT",
+        cast=bool,
+        required=False,
+        default=False,
+    )
 
 
 feature_flags = FeatureFlags()

app/main.py

@@ -5,7 +5,7 @@
 from fastapi.staticfiles import StaticFiles
 from fastapi_pagination import add_pagination
 
-from app.core import settings
+from app.core import feature_flags, settings
 from app.graph_ql.routers import router as graphql_router
 from app.routers.callbacks import router as callbacks_router
 from app.routers.characters import router as characters_router
@@ -15,6 +15,7 @@
 from app.routers.seasons import router as seasons_router
 from app.routers.tokens import router as tokens_router
 from app.routers.users import router as users_router
+from app.middlewares.secure import HTTPSRedirectMiddleware
 
 mimetypes.add_type("image/webp", ".webp")
 
@@ -23,6 +24,9 @@
     redoc_url=None,
 )
 
+if feature_flags.enable_https_redirect:
+    app.add_middleware(HTTPSRedirectMiddleware)
+
 app.add_middleware(
     CORSMiddleware,
     allow_origins=settings.allow_origins,

app/middlewares/secure.py

@@ -0,0 +1,59 @@
+from typing import Dict
+
+from starlette import status
+from starlette.datastructures import URL
+from starlette.responses import RedirectResponse
+from starlette.types import ASGIApp, Receive, Scope, Send
+
+from app.core import settings
+
+
+class HTTPSRedirectMiddleware:
+    https_port = 443
+    http_port = 80
+    proto_header = "x-forwarded-proto"
+    port_header = "x-forwarded-port"
+
+    def __init__(self, app: ASGIApp) -> None:
+        self.app = app
+
+    def is_secure(self, headers: Dict):
+        try:
+            host: str = headers["host"]
+        except KeyError:
+            return False
+        try:
+            proto: str = headers[self.proto_header]
+        except KeyError:
+            return False
+        try:
+            port: str = headers[self.port_header]
+        except KeyError:
+            return False
+
+        if (
+            host == settings.trusted_host
+            and proto in ("https", "wss")
+            and int(port) == self.https_port
+        ):
+            return True
+        return False
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        headers: Dict = {h[0].decode().lower(): h[1].decode() for h in scope["headers"]}
+        if not self.is_secure(headers):
+            url = URL(scope=scope)
+            redirect_scheme = {"http": "https", "ws": "wss"}[url.scheme]
+            netloc = (
+                url.hostname
+                if url.port in (self.http_port, self.https_port)
+                else url.netloc
+            )
+            url = url.replace(scheme=redirect_scheme, netloc=netloc)
+            response = RedirectResponse(
+                url,
+                status_code=status.HTTP_307_TEMPORARY_REDIRECT,
+            )
+            await response(scope, receive, send)
+        else:
+            await self.app(scope, receive, send)