up

main.py

@@ -1,7 +1,6 @@
 import fastapi
 import uvicorn
-
-from app.middlewares import check_token
+import os
 
 from app.routers import take
 from app.routers import put
@@ -10,9 +9,19 @@ from app.routers import finish
 from app.storage import mongo
 
 
+QUEUES_TOKEN = os.getenv('QUEUES_TOKEN')
+
 app = fastapi.FastAPI()
 
-app.add_middleware(check_token.CheckToken())
+
+@app.middleware('http')
+async def check_token(request: fastapi.Request, call_next):
+    if QUEUES_TOKEN:
+        token = request.headers.get('X-Queues-Token')
+        if not token or token != QUEUES_TOKEN:
+            raise fastapi.HTTPException(403)
+    return await call_next(request)
+
 
 app.include_router(take.router)
 app.include_router(put.router)