turns out the starlette CORS middleware sucks

Author: IAmTomahawkx

mystbin/backend/app.py

@@ -21,14 +21,13 @@
 import pathlib
 import os
 import sys
-from typing import Any, Dict, Optional
+from typing import Any, Callable, Coroutine, Dict, Optional
 
 import aiohttp
 import aioredis
 import sentry_sdk
 import ujson
-from fastapi import FastAPI, Request
-from fastapi.middleware.cors import CORSMiddleware
+from fastapi import FastAPI, Request, Response
 from sentry_sdk.integrations.asgi import SentryAsgiMiddleware
 from starlette_prometheus import metrics, PrometheusMiddleware
 
@@ -66,10 +65,20 @@ def __init__(self, *, loop: Optional[asyncio.AbstractEventLoop] = None, config:
 
 
 app = MystbinApp()
+METHODS = ("DELETE", "GET", "OPTIONS", "PATCH", "POST", "PUT")
 
 
 @app.middleware("http")
 async def request_stats(request: Request, call_next):
+    if request.method == "OPTIONS":
+        raise RuntimeError("blah")
+        return Response(headers={
+            "Access-Control-Allowed-Headers": request.headers.get("Access-Control-Request-Headers", ""),
+            "Access-Control-Allowed-Method": ", ".join(METHODS),
+            "Access-Control-Allowed-Origin": app.config["site"]["frontend_site"],
+            "Access-Control-Max-Age": "600",
+            "Vary": "Origin",
+        })
     request.app.state.request_stats["total"] += 1
 
     if request.url.path != "/admin/stats":
@@ -78,6 +87,22 @@ async def request_stats(request: Request, call_next):
     response = await call_next(request)
     return response
 
+async def cors_middleware(request: Request, call_next: Callable[[Request], Coroutine[Any, Any, Response]]):
+    headers={
+        "Access-Control-Allow-Headers": request.headers.get("Access-Control-Request-Headers", ""),
+        "Access-Control-Allow-Methods": ", ".join(METHODS),
+        "Access-Control-Allow-Origin": app.config["site"]["frontend_site"],
+        "Access-Control-Max-Age": "600",
+        "Vary": "Origin",
+    }
+    
+    if request.method == "OPTIONS":
+        return Response(headers=headers)
+    
+    resp = await call_next(request)
+    resp.headers.update(headers)
+    return resp
+
 
 @app.on_event("startup")
 async def app_startup():
@@ -98,6 +123,7 @@ async def app_startup():
 
     ratelimits.limiter.startup(app)
     app.middleware("http")(ratelimits.limiter.middleware)
+    app.middleware("http")(cors_middleware)
 
     nocli = pathlib.Path(".nocli")
     if nocli.exists():
@@ -114,17 +140,6 @@ async def app_startup():
 app.include_router(pastes.router)
 app.include_router(user.router)
 
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=[
-        app.config["site"]["frontend_site"],
-        app.config["site"]["backend_site"],
-    ],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
 
 try:
     sentry_dsn = app.config["sentry"]["dsn"]
@@ -136,5 +151,5 @@ async def app_startup():
 
     app.add_middleware(SentryAsgiMiddleware)
 
-app.add_middleware(PrometheusMiddleware)
-app.add_route("/metrics/", metrics)
+#app.add_middleware(PrometheusMiddleware)
+#app.add_route("/metrics/", metrics)

mystbin/backend/utils/ratelimits.py

@@ -152,17 +152,19 @@ async def get_key(self, zone: str, request: Request) -> str:
 
     async def _transform_and_log(self, request: Request, response: Response) -> Response:
         if isinstance(response, StreamingResponse):
-                resp_ = Response(status_code=response.status_code, background=response.background, media_type=cast(str, response.media_type))
-                resp_._headers = response._headers
-                body = b""
-                async for chunk in response.body_iterator:
-                    if not isinstance(chunk, bytes):
-                        chunk = chunk.encode(response.charset)
-                    
-                    body += chunk
+            resp_ = Response(status_code=response.status_code, background=response.background, media_type=cast(str, response.media_type))
+            del response._headers['content-type']
+            del response._headers['content-length']
+            resp_._headers = response._headers
+            body = b""
+            async for chunk in response.body_iterator:
+                if not isinstance(chunk, bytes):
+                    chunk = chunk.encode(response.charset)
 
-                resp_.body = body
-                response = resp_
+                body += chunk
+            
+            resp_.body = body
+            response = resp_
 
         await self.app.state.db.put_log(request, response)
         return response
@@ -226,7 +228,8 @@ async def middleware(self, request: Request, call_next: _CT) -> Response:
 
             resp = await call_next(request)
             resp.headers.update(headers)
-            return await self._transform_and_log(request, resp)
+            resp = await self._transform_and_log(request, resp)
+            return resp
 
 
         else: