starlette sucks

Author: IAmTomahawkx

mystbin/backend/utils/db.py

@@ -939,16 +939,18 @@ async def put_log(
         INSERT INTO logs VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)
         """
         try:
-            body = await request.body()
+            body = request._body
         except:
             body = None
+
         try:
             resp = str(response.body)
         except AttributeError:
             resp = None
         await self._do_query(
             query,
             request.headers.get("X-Forwarded-For", request.client.host),
+            request.state.user and request.state.user.id,
             datetime.datetime.utcnow(),
             request.headers.get("CF-RAY"),
             request.headers.get("CF-IPCOUNTRY"),

mystbin/backend/utils/ratelimits.py

@@ -7,7 +7,7 @@
 
 import aioredis
 from starlette.requests import Request
-from starlette.responses import Response, JSONResponse
+from starlette.responses import Response, JSONResponse, StreamingResponse
 from starlette.routing import Match
 
 from . import tokens
@@ -150,6 +150,23 @@ async def get_key(self, zone: str, request: Request) -> str:
         key = ratelimit_id_key(request)
         return f"{zone}%{key}"
 
+    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_.body = body
+                response = resp_
+            
+        await self.app.state.db.put_log(request, response)
+        return response
+    
     async def get_bucket(self, zone: str, key: str, request: Request) -> BaseLimitBucket:
         if key in self._keys:
             return self._keys[key]
@@ -209,7 +226,8 @@ async def middleware(self, request: Request, call_next: _CT) -> Response:
 
             resp = await call_next(request)
             resp.headers.update(headers)
-            return resp
+            return await self._transform_and_log(request, resp)
+
 
         else:
             headers = {
@@ -225,8 +243,7 @@ async def middleware(self, request: Request, call_next: _CT) -> Response:
             }
             resp = await call_next(request)
             resp.headers.update(headers)
-            self.app.loop.create_task(self.app.state.db.put_log(request, resp))
-            return resp
+            return await self._transform_and_log(request, resp)
 
 
 def parse_ratelimit(limit: str) -> tuple[int, int]: