feat(host-agent): add BearerAuthMiddleware for MCP server
This commit is contained in:
@@ -0,0 +1,36 @@
|
|||||||
|
"""Bearer-token auth middleware for the MCP sub-app.
|
||||||
|
|
||||||
|
Mounted on the FastMCP ``streamable_http_app()`` (NOT the console FastAPI),
|
||||||
|
so cookie-session auth on console routes is unaffected.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from starlette.middleware.base import BaseHTTPMiddleware
|
||||||
|
from starlette.requests import Request
|
||||||
|
from starlette.responses import JSONResponse, Response
|
||||||
|
|
||||||
|
from host_agent.mcp_token import McpTokenStore
|
||||||
|
|
||||||
|
|
||||||
|
class BearerAuthMiddleware(BaseHTTPMiddleware):
|
||||||
|
def __init__(self, app, token_store: McpTokenStore) -> None:
|
||||||
|
super().__init__(app)
|
||||||
|
self._store = token_store
|
||||||
|
|
||||||
|
async def dispatch(self, request: Request, call_next) -> Response: # type: ignore[no-untyped-def]
|
||||||
|
header = request.headers.get("Authorization")
|
||||||
|
if not header or not header.lower().startswith("bearer "):
|
||||||
|
return _unauthorized()
|
||||||
|
presented = header.split(" ", 1)[1].strip()
|
||||||
|
if not self._store.verify(presented):
|
||||||
|
return _unauthorized()
|
||||||
|
return await call_next(request)
|
||||||
|
|
||||||
|
|
||||||
|
def _unauthorized() -> JSONResponse:
|
||||||
|
return JSONResponse(
|
||||||
|
status_code=401,
|
||||||
|
content={"error": "invalid token"},
|
||||||
|
headers={"WWW-Authenticate": "Bearer"},
|
||||||
|
)
|
||||||
@@ -0,0 +1,58 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
from starlette.applications import Starlette
|
||||||
|
from starlette.responses import JSONResponse
|
||||||
|
from starlette.testclient import TestClient
|
||||||
|
|
||||||
|
from host_agent.mcp_token import McpTokenStore
|
||||||
|
from host_agent.web.mcp_auth import BearerAuthMiddleware
|
||||||
|
|
||||||
|
|
||||||
|
def _make_client(tmp_path: Path) -> tuple[TestClient, str]:
|
||||||
|
store = McpTokenStore(tmp_path / "host_mcp_token.json")
|
||||||
|
token = store.load_or_create().token
|
||||||
|
|
||||||
|
async def hello(request): # type: ignore[no-untyped-def]
|
||||||
|
return JSONResponse({"ok": True})
|
||||||
|
|
||||||
|
inner = Starlette(routes=[])
|
||||||
|
inner.router.add_route("/", hello, methods=["GET"])
|
||||||
|
wrapped = Starlette()
|
||||||
|
wrapped.add_middleware(BearerAuthMiddleware, token_store=store)
|
||||||
|
wrapped.mount("/", inner)
|
||||||
|
return TestClient(wrapped), token
|
||||||
|
|
||||||
|
|
||||||
|
def test_no_header_returns_401(tmp_path: Path) -> None:
|
||||||
|
client, _ = _make_client(tmp_path)
|
||||||
|
resp = client.get("/")
|
||||||
|
assert resp.status_code == 401
|
||||||
|
assert resp.headers["WWW-Authenticate"] == "Bearer"
|
||||||
|
assert resp.json() == {"error": "invalid token"}
|
||||||
|
|
||||||
|
|
||||||
|
def test_wrong_token_returns_401(tmp_path: Path) -> None:
|
||||||
|
client, _ = _make_client(tmp_path)
|
||||||
|
resp = client.get("/", headers={"Authorization": "Bearer wrong"})
|
||||||
|
assert resp.status_code == 401
|
||||||
|
|
||||||
|
|
||||||
|
def test_correct_token_passes_through(tmp_path: Path) -> None:
|
||||||
|
client, token = _make_client(tmp_path)
|
||||||
|
resp = client.get("/", headers={"Authorization": f"Bearer {token}"})
|
||||||
|
assert resp.status_code == 200
|
||||||
|
assert resp.json() == {"ok": True}
|
||||||
|
|
||||||
|
|
||||||
|
def test_non_bearer_scheme_returns_401(tmp_path: Path) -> None:
|
||||||
|
client, token = _make_client(tmp_path)
|
||||||
|
resp = client.get("/", headers={"Authorization": f"Basic {token}"})
|
||||||
|
assert resp.status_code == 401
|
||||||
|
|
||||||
|
|
||||||
|
def test_header_case_insensitive(tmp_path: Path) -> None:
|
||||||
|
client, token = _make_client(tmp_path)
|
||||||
|
resp = client.get("/", headers={"authorization": f"Bearer {token}"})
|
||||||
|
assert resp.status_code == 200
|
||||||
Reference in New Issue
Block a user