iterative / mlem

🐶 A tool to package, serve, and deploy any ML model on any platform. Archived to be resurrected one day🤞
https://mlem.ai
Apache License 2.0
718 stars 44 forks source link

Right way to add authentification middleware #692

Open aguschin opened 1 year ago

aguschin commented 1 year ago

Adding middleware ended up dirtier than if I wrote the code for FastAPI serving myself, because the FastAPIMiddleware seems to be applied after the serializer, and I wanted a middleware applied before serializers (for authentication purposes).

Other question: If I'm understanding correctly, the middleware on_request methods are called after the request_serializer. Is there a way to add layers before the request_serializer.

I was trying to add authentification middleware (for FastApi). I first tried to inherit from FastAPIMiddleware , and do a test in the on_request method (to test if the field api_key of the request header has the correct value). The problem is, since I am using the pil_numpy serializer, the request argument of on_request is a Numpy array.

So the only solution I have found is:

class KeyAuthMiddleware(FastAPIMiddleware):
    api_key: str

    def on_app_init(self, app: FastAPI):
        @app.middleware("http")
        async def api_key_auth_middleware(request: Request, call_next):
            api_key = request.headers.get("api_key")
            if api_key != self.api_key:
                raise HTTPException(status_code=401, detail="Invalid API Key")
            response = await call_next(request)
            return response

Is there a better way to do it?

reported in discord