| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import hashlib |
| 4 | import secrets |
| 5 | from collections.abc import Awaitable |
| 6 | |
| 7 | from starlette.datastructures import Headers |
| 8 | from starlette.responses import JSONResponse |
| 9 | from starlette.types import ASGIApp, Receive, Scope, Send |
| 10 | |
| 11 | GUARDED_PREFIX = ("/v1", "/v2", "/inference", "/cohere") |
| 12 | |
| 13 | |
| 14 | class AuthenticationMiddleware: |
| 15 | """Pure ASGI middleware that authenticates each request by checking |
| 16 | if the Authorization Bearer token exists and equals anyof "{api_key}". |
| 17 | |
| 18 | Notes |
| 19 | ----- |
| 20 | There are two cases in which authentication is skipped: |
| 21 | 1. The HTTP method is OPTIONS. |
| 22 | 2. The request path doesn't start with GUARDED_PREFIX (e.g. /health). |
| 23 | |
| 24 | """ |
| 25 | |
| 26 | def __init__(self, app: ASGIApp, tokens: list[str]) -> None: |
| 27 | self.app = app |
| 28 | self.api_tokens = [hashlib.sha256(t.encode("utf-8")).digest() for t in tokens] |
| 29 | |
| 30 | def verify_token(self, headers: Headers) -> bool: |
| 31 | authorization_header_value = headers.get("Authorization") |
| 32 | if not authorization_header_value: |
| 33 | return False |
| 34 | |
| 35 | scheme, _, param = authorization_header_value.partition(" ") |
| 36 | if scheme.lower() != "bearer": |
| 37 | return False |
| 38 | |
| 39 | param_hash = hashlib.sha256(param.encode("utf-8")).digest() |
| 40 | |
| 41 | token_match = False |
| 42 | for token_hash in self.api_tokens: |
| 43 | token_match |= secrets.compare_digest(param_hash, token_hash) |
| 44 | |
| 45 | return token_match |
| 46 | |
| 47 | def __call__(self, scope: Scope, receive: Receive, send: Send) -> Awaitable[None]: |
| 48 | if ( |
| 49 | scope["type"] not in ("http", "websocket") |
| 50 | or scope.get("method") == "OPTIONS" |
| 51 | ): |
| 52 | # scope["type"] can be "lifespan" or "startup" for example, |
| 53 | # in which case we don't need to do anything |
| 54 | return self.app(scope, receive, send) |
| 55 | root_path = scope.get("root_path", "") |
| 56 | url_path = scope["path"].removeprefix(root_path) |
| 57 | headers = Headers(scope=scope) |
| 58 | # Type narrow to satisfy mypy. |
| 59 | if url_path.startswith(GUARDED_PREFIX) and not self.verify_token(headers): |
| 60 | response = JSONResponse(content={"error": "Unauthorized"}, status_code=401) |
| 61 | return response(scope, receive, send) |
| 62 | return self.app(scope, receive, send) |