From 2e99aadbd9e5a543519b1c7d6c619a94613b3311 Mon Sep 17 00:00:00 2001 From: Ben Wilson Date: Sun, 30 Aug 2026 19:18:02 -0700 Subject: [PATCH] feat(server): add --api-key bearer authentication Require `Authorization: Bearer ` on every route except /health when `ft serve --api-key` (or FREETOKEN_API_KEY) is set; 401 with a WWW-Authenticate challenge otherwise. Unset, nothing changes. - The check is a middleware registered after the request-ring middleware, so it runs first: a rejected request never reaches a handler or the ring. OPTIONS passes through for CORS preflights; the CORS middleware installed at startup stays outermost. Constant-time compare. - /health stays open: liveness probes and the desktop app's load-progress polling cannot carry a header and reveal only status. - Shell mode keeps working: the attached client is handed the server's key; `ft shell` attaching to a running server reads FREETOKEN_API_KEY. - The control-plane requests of the shell client now send the same bearer the OpenAI client already sends. Closes #152 --- docs/cli.md | 1 + python/freetoken/server/api_server.py | 56 ++++++- python/freetoken/server/args.py | 19 +++ python/freetoken/shell/client.py | 6 +- python/freetoken/shell/tui.py | 14 +- tests/server/test_api_key.py | 212 ++++++++++++++++++++++++++ 6 files changed, 301 insertions(+), 7 deletions(-) create mode 100644 tests/server/test_api_key.py diff --git a/docs/cli.md b/docs/cli.md index ff4af382d..6e05a39fa 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -40,6 +40,7 @@ parsers all resolve automatically from the checkpoint and the GPU. |---|---|---| | `--host` | 127.0.0.1 | Bind address | | `--port` | 1919 | Bind port | +| `--api-key` | disabled | Require `Authorization: Bearer ` on every route except `/health` (401 otherwise); `FREETOKEN_API_KEY` is read when the flag is absent | | `--gpu` | GPU 0 | GPU to run on: a UUID from `nvidia-smi -L` or an `nvidia-smi` index; see [below](#choosing-a-gpu) | | `--max-running-requests` | 4 | Max concurrently running requests | | `--max-output-tokens` | 32768 | Default output budget for requests that omit one | diff --git a/python/freetoken/server/api_server.py b/python/freetoken/server/api_server.py index 3e2acc854..8d07991b2 100644 --- a/python/freetoken/server/api_server.py +++ b/python/freetoken/server/api_server.py @@ -2,6 +2,7 @@ import asyncio import contextlib +import hmac import json import os import signal @@ -437,6 +438,23 @@ def install_cors(app: FastAPI, origins_csv: str) -> None: _UNTRACKED_REQUEST_PREFIXES = ("/v1/messages/count_tokens",) +# --api-key (ServerArgs.api_key), installed by run_api_server before serving. None means no +# authentication -- every route answers as before. Module-level rather than closed over so the +# middleware below can be registered at import time like the others and tests can set it. +_API_KEY: str | None = None +# Routes that stay open with a key set: liveness probes (load balancers, Docker healthchecks, +# the desktop app's load-progress polling) cannot carry a header and reveal nothing but status. +_API_KEY_OPEN_PATHS = ("/health",) + + +def install_api_key(key: str | None) -> None: + """Arm the bearer check for every route but ``_API_KEY_OPEN_PATHS``; ``None``/"" disarms.""" + global _API_KEY + _API_KEY = key or None + if _API_KEY is not None: + logger.info("API key required (Authorization: Bearer) on every route except /health") + + def _served_model_name() -> str | None: st = _GLOBAL_STATE cfg = getattr(st, "config", None) if st is not None else None @@ -566,6 +584,36 @@ def _resolve_num_swa_pages(state: FrontendManager, req: CacheRebuildRequest) -> return max(1, -(-window_tokens // swa_page_size)) # ceil-div to the pool's page unit +@app.middleware("http") +async def _api_key_middleware(request: Request, call_next): + """Reject any request without ``Authorization: Bearer `` when a key is set. + + Registered after ``_record_request_middleware`` so it runs *before* it (Starlette wraps the + last-added middleware outermost): a 401 never lands in the request ring or a handler. CORS + preflights (OPTIONS) carry no credentials by design and pass through; the CORS middleware + installed at startup is outer still, so it answers them. Constant-time compare, and the + same body shape the OpenAI-compatible routes use for errors.""" + key = _API_KEY + if key is None or request.method == "OPTIONS" or request.url.path in _API_KEY_OPEN_PATHS: + return await call_next(request) + scheme, _, token = request.headers.get("authorization", "").partition(" ") + if scheme.lower() != "bearer" or not hmac.compare_digest( + token.strip().encode("utf-8"), key.encode("utf-8") + ): + return JSONResponse( + status_code=401, + headers={"WWW-Authenticate": "Bearer"}, + content={ + "error": { + "message": "Invalid or missing API key (Authorization: Bearer ).", + "type": "authentication_error", + "code": 401, + } + }, + ) + return await call_next(request) + + @app.post("/v1/cache/rebuild") async def cache_rebuild(req: CacheRebuildRequest): """Trigger a runtime KV/MoE cache resize. Blocks until the scheduler reports a result @@ -873,7 +921,7 @@ def _flag_shutdown(signum, frame) -> None: signal.signal(sig, _flag_shutdown) -def _serve_and_run_shell(host: str, port: int) -> None: +def _serve_and_run_shell(host: str, port: int, api_key: str | None = None) -> None: """Shell mode: serve the API here, and attach the terminal client to it over the loopback. The shell is an ordinary API client (see ``freetoken.shell``), so shell mode is just @@ -906,7 +954,8 @@ def _serve_and_run_shell(host: str, port: int) -> None: # /health and echoes the same load progress the desktop app polls for. A ^C during that # wait is a stop, not a crash -- exit through the teardown below, not a traceback. with contextlib.suppress(KeyboardInterrupt): - asyncio.run(run_shell(origin, connect_grace=30.0)) + # The attached shell is an ordinary client, so it carries the server's own key. + asyncio.run(run_shell(origin, connect_grace=30.0, api_key=api_key)) finally: server.should_exit = True thread.join(timeout=15) @@ -951,6 +1000,7 @@ def run_api_server(config: ServerArgs, start_backend: Callable[[], "Any"], run_s # Create/validate FREETOKEN_API_LOG_DIR and start the writer thread up front, so a # bad path is reported at boot rather than silently on the first request. install_cors(app, config.cors_origins) + install_api_key(config.api_key) init_request_logging() # Hide the frequent health/stats/requests/cache-status polling of the desktop app (and of # the shell's status bar) from uvicorn's access log; non-polling access lines are @@ -1034,7 +1084,7 @@ def _on_meta(meta: dict) -> None: ).start() if run_shell: - _serve_and_run_shell(host, port) + _serve_and_run_shell(host, port, config.api_key) return # uvicorn stays on the main thread (signal handling unchanged); ^C reaches the worker group. uvicorn.run(app, host=host, port=port) diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index 6696f65dd..e19106346 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -15,6 +15,9 @@ class ServerArgs(SchedulerConfig): server_host: str = "127.0.0.1" server_port: int = 1919 + # Bearer token every request must carry (except /health). None = no authentication, + # today's behaviour. Read from FREETOKEN_API_KEY when --api-key is not given. + api_key: str | None = None num_tokenizer: int = 0 silent_output: bool = False # The terminal shell is attached to this server (ft shell --model / ft serve --shell-mode). @@ -308,6 +311,17 @@ def _infer_reasoning_parser(model_path: str) -> str | None: help="The port number for the server to listen on.", ) + parser.add_argument( + "--api-key", + type=str, + default=ServerArgs.api_key, + help=( + "Require `Authorization: Bearer ` on every route except /health " + "(401 otherwise). Unset: no authentication. When the flag is absent, " + "FREETOKEN_API_KEY is read instead so the key need not appear in `ps`." + ), + ) + parser.add_argument( "--cuda-graph-max-bs", "--graph", @@ -654,6 +668,11 @@ def _infer_reasoning_parser(model_path: str) -> str | None: if kwargs["model_path"].startswith("~"): kwargs["model_path"] = os.path.expanduser(kwargs["model_path"]) + if kwargs["api_key"] is None: + kwargs["api_key"] = os.environ.get("FREETOKEN_API_KEY") or None + elif not kwargs["api_key"].strip(): + parser.error("--api-key must not be empty (omit it to serve without authentication)") + if kwargs["served_model_name"] is None: kwargs["served_model_name"] = ( os.path.basename(os.path.normpath(kwargs["model_path"])) or kwargs["model_path"] diff --git a/python/freetoken/shell/client.py b/python/freetoken/shell/client.py index 183fa4dff..3454d4cc8 100644 --- a/python/freetoken/shell/client.py +++ b/python/freetoken/shell/client.py @@ -102,6 +102,10 @@ def __init__( ) -> None: self.origin = origin.rstrip("/") self.timeout = timeout + # Sent on the control plane too (the OpenAI client already sends it on chat), so a + # server started with --api-key accepts both halves of the shell. The default is the + # placeholder an unauthenticated server ignores. + self.api_key = api_key self._openai = AsyncOpenAI( base_url=f"{self.origin}/v1", api_key=api_key, @@ -118,7 +122,7 @@ def _request_json_blocking( self, method: str, path: str, body: dict[str, Any] | None, timeout: float ) -> dict[str, Any]: data = None - headers = {"Accept": "application/json"} + headers = {"Accept": "application/json", "Authorization": f"Bearer {self.api_key}"} if body is not None: data = json.dumps(body).encode("utf-8") headers["Content-Type"] = "application/json" diff --git a/python/freetoken/shell/tui.py b/python/freetoken/shell/tui.py index 93bd18cc3..36d54db1e 100644 --- a/python/freetoken/shell/tui.py +++ b/python/freetoken/shell/tui.py @@ -8,6 +8,7 @@ from __future__ import annotations import asyncio +import os import contextlib import re import shutil @@ -33,6 +34,7 @@ from prompt_toolkit.styles import Style from .client import ( + LOCAL_API_KEY, ContentDelta, ReasoningDelta, Sampling, @@ -515,13 +517,19 @@ def _format_load_progress(doc: dict) -> str: return f"loading ({phase})..." -async def run_shell(origin: str, *, connect_grace: float = 0.0) -> int: +async def run_shell( + origin: str, *, connect_grace: float = 0.0, api_key: str | None = None +) -> int: """Attach to the FreeToken server at ``origin`` and run the terminal chat. ``connect_grace`` is how long to keep retrying a refused connection before giving up -- left at 0 when attaching to a server the user says is already running, raised when the - caller just started one in this process (see ``server/api_server.py``).""" - client = ShellClient(origin) + caller just started one in this process (see ``server/api_server.py``). + + ``api_key`` is the server's ``--api-key`` when it has one: passed in by shell mode, read + from ``FREETOKEN_API_KEY`` when attaching to a running server, else the local placeholder.""" + key = api_key or os.environ.get("FREETOKEN_API_KEY") or LOCAL_API_KEY + client = ShellClient(origin, api_key=key) try: return await _run_shell(client, origin, connect_grace=connect_grace) finally: diff --git a/tests/server/test_api_key.py b/tests/server/test_api_key.py new file mode 100644 index 000000000..bbbe64446 --- /dev/null +++ b/tests/server/test_api_key.py @@ -0,0 +1,212 @@ +"""--api-key: the bearer gate on the API server, its argument parsing, and the shell client. + +The gate is a middleware registered at import time and armed by ``install_api_key`` (what +``run_api_server`` calls with ``ServerArgs.api_key``). With no key it is inert, so every +existing test keeps its behaviour. These tests drive the real ``api.app`` through a TestClient +and flip ``_API_KEY`` directly, the way ``test_rebuild_maintenance`` flips ``_GLOBAL_STATE``. +""" + +from __future__ import annotations + +import urllib.request +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +from fastapi.testclient import TestClient + +from freetoken.server.args import parse_args + + +class _Config: + def to_dict(self) -> dict: + return {"architectures": ["DeepseekV4ForCausalLM"], "torch_dtype": "bfloat16"} + + +def _parse(extra: list[str]): + with patch("freetoken.utils.cached_load_hf_config", lambda _path: _Config()): + return parse_args(["--model", "/models/anon", *extra]) + + +# ------------------------------------------------------------------ argument parsing + + +def test_api_key_defaults_to_none(monkeypatch): + monkeypatch.delenv("FREETOKEN_API_KEY", raising=False) + args, _ = _parse([]) + assert args.api_key is None + + +def test_api_key_flag_is_parsed(monkeypatch): + monkeypatch.delenv("FREETOKEN_API_KEY", raising=False) + args, _ = _parse(["--api-key", "s3cret"]) + assert args.api_key == "s3cret" + + +def test_api_key_falls_back_to_the_environment(monkeypatch): + monkeypatch.setenv("FREETOKEN_API_KEY", "from-env") + args, _ = _parse([]) + assert args.api_key == "from-env" + + +def test_api_key_flag_wins_over_the_environment(monkeypatch): + monkeypatch.setenv("FREETOKEN_API_KEY", "from-env") + args, _ = _parse(["--api-key", "from-flag"]) + assert args.api_key == "from-flag" + + +def test_empty_environment_key_means_unset(monkeypatch): + monkeypatch.setenv("FREETOKEN_API_KEY", "") + args, _ = _parse([]) + assert args.api_key is None + + +def test_empty_api_key_flag_is_rejected(monkeypatch): + monkeypatch.delenv("FREETOKEN_API_KEY", raising=False) + with pytest.raises(SystemExit, match="2"): + _parse(["--api-key", " "]) + + +def test_shell_mode_accepts_a_key(monkeypatch): + # Unlike TLS, a key is fine in shell mode: the attached client carries it (see below). + monkeypatch.delenv("FREETOKEN_API_KEY", raising=False) + args, run_shell = _parse(["--shell-mode", "--api-key", "s3cret"]) + assert run_shell is True + assert args.api_key == "s3cret" + + +# ------------------------------------------------------------------ the middleware + + +def _serving_state(): + return SimpleNamespace( + maintenance_state="serving", + config=SimpleNamespace(served_model_name="anon"), + fatal_error=None, + ready_at=None, + instance_id=None, + rebuild_futures={}, + last_rebuild=None, + ) + + +@pytest.fixture +def keyed_client(monkeypatch): + import freetoken.server.api_server as api + + monkeypatch.setattr(api, "_API_KEY", "s3cret") + monkeypatch.setattr(api, "_GLOBAL_STATE", _serving_state()) + return TestClient(api.app) + + +def test_no_key_configured_leaves_every_route_open(monkeypatch): + import freetoken.server.api_server as api + + monkeypatch.setattr(api, "_API_KEY", None) + monkeypatch.setattr(api, "_GLOBAL_STATE", _serving_state()) + client = TestClient(api.app) + assert client.get("/v1").status_code == 200 + assert client.get("/health").status_code == 200 + + +def test_missing_header_is_401_with_a_bearer_challenge(keyed_client): + r = keyed_client.get("/v1") + assert r.status_code == 401 + assert r.headers["WWW-Authenticate"] == "Bearer" + assert r.json()["error"]["type"] == "authentication_error" + + +@pytest.mark.parametrize( + "authorization", + ["Bearer wrong", "Bearer s3cre", "Bearer s3cret-and-more", "Basic s3cret", "s3cret"], +) +def test_wrong_or_malformed_credentials_are_401(keyed_client, authorization): + r = keyed_client.get("/v1", headers={"Authorization": authorization}) + assert r.status_code == 401 + + +def test_matching_bearer_passes(keyed_client): + r = keyed_client.get("/v1", headers={"Authorization": "Bearer s3cret"}) + assert r.status_code == 200 + assert r.json() == {"status": "ok"} + # scheme is case-insensitive, surrounding whitespace on the token is not the token + r = keyed_client.get("/v1", headers={"Authorization": "bearer s3cret "}) + assert r.status_code == 200 + + +def test_health_stays_open_for_liveness_probes(keyed_client): + r = keyed_client.get("/health") + assert r.status_code == 200 + assert r.json()["status"] == "ok" + + +@pytest.mark.parametrize( + "method, path", + [ + ("POST", "/v1/chat/completions"), + ("POST", "/v1/messages"), + ("POST", "/v1/responses"), + ("GET", "/v1/models"), + ("GET", "/v1/stats"), + ("GET", "/v1/requests"), + ("POST", "/v1/cache/rebuild"), + ("POST", "/v1/admin/prepare-stop"), + ("POST", "/generate"), + ], +) +def test_every_other_route_is_gated_before_its_handler(keyed_client, method, path): + # No body and no engine behind these: a 401 proves the gate answered first, since the + # handler would have produced a 422/503 or needed the state. + r = keyed_client.request(method, path) + assert r.status_code == 401 + + +def test_cors_preflight_is_not_challenged(keyed_client): + # A preflight carries no credentials by design; the CORS middleware (installed at startup, + # outermost) answers it. Without CORS configured FastAPI's own answer is what we see -- + # anything but a 401 is the assertion. + r = keyed_client.options("/v1/chat/completions") + assert r.status_code != 401 + + +def test_install_api_key_arms_and_disarms(): + import freetoken.server.api_server as api + + prev = api._API_KEY + try: + api.install_api_key("k") + assert api._API_KEY == "k" + api.install_api_key("") + assert api._API_KEY is None + api.install_api_key(None) + assert api._API_KEY is None + finally: + api._API_KEY = prev + + +# ------------------------------------------------------------------ the shell client + + +def test_shell_client_sends_the_key_on_the_control_plane(monkeypatch): + from freetoken.shell.client import ShellClient + + seen: dict = {} + + class _Resp: + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + def read(self): + return b'{"status": "ok"}' + + def _urlopen(request, timeout): + seen["authorization"] = request.get_header("Authorization") + return _Resp() + + monkeypatch.setattr(urllib.request, "urlopen", _urlopen) + client = ShellClient("http://127.0.0.1:1", api_key="s3cret") + assert client._request_json_blocking("GET", "/health", None, 1.0) == {"status": "ok"} + assert seen["authorization"] == "Bearer s3cret"