diff --git a/main.py b/main.py index d704da5..f3b6552 100644 --- a/main.py +++ b/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)