diff --git a/.github/copy-pr-bot.yaml b/.github/copy-pr-bot.yaml index 7fe36225..d4ca9418 100644 --- a/.github/copy-pr-bot.yaml +++ b/.github/copy-pr-bot.yaml @@ -1 +1 @@ -enabled: true \ No newline at end of file +enabled: true diff --git a/Cargo.lock b/Cargo.lock index 2864fa4e..a072f4ba 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -156,6 +156,17 @@ dependencies = [ "pin-project-lite", ] +[[package]] +name = "async-lock" +version = "3.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "290f7f2596bd5b78a9fec8088ccd89180d7f9f55b94b0576823bbbdc72ee8311" +dependencies = [ + "event-listener", + "event-listener-strategy", + "pin-project-lite", +] + [[package]] name = "async-stream" version = "0.3.6" @@ -638,6 +649,15 @@ dependencies = [ "itertools 0.10.5", ] +[[package]] +name = "crossbeam-channel" +version = "0.5.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "82b8f8f868b36967f9606790d1903570de9ceaf870a7bf9fbbd3016d636a2cb2" +dependencies = [ + "crossbeam-utils", +] + [[package]] name = "crossbeam-deque" version = "0.8.6" @@ -1523,12 +1543,11 @@ dependencies = [ [[package]] name = "http" -version = "1.3.1" +version = "1.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f4a85d31aea989eead29a3aaf9e1115a180df8282431156e533de47660892565" +checksum = "8be7462df143984c4598a256ef469b251d7d7f9e271135073e78fc535414f3d0" dependencies = [ "bytes", - "fnv", "itoa", ] @@ -2455,21 +2474,25 @@ dependencies = [ "config", "futures", "hf-hub", + "http", "jiff", "k8s-openapi", "kube", "mockall 0.14.0", "modelexpress-client", "modelexpress-common", + "moka", "once_cell", "prost 0.13.5", "redis", "schemars 0.8.22", + "secrecy", "serde", "serde_json", "serde_yaml", "sha2", "tempfile", + "thiserror 2.0.16", "tokio", "tokio-stream", "tokio-test", @@ -2477,11 +2500,32 @@ dependencies = [ "tonic 0.13.1", "tonic-build", "tonic-health", + "tower", "tracing", "tracing-subscriber", "uuid", ] +[[package]] +name = "moka" +version = "0.12.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "957228ad12042ee839f93c8f257b62b4c0ab5eaae1d4fa60de53b27c9d7c5046" +dependencies = [ + "async-lock", + "crossbeam-channel", + "crossbeam-epoch", + "crossbeam-utils", + "equivalent", + "event-listener", + "futures-util", + "parking_lot", + "portable-atomic", + "smallvec", + "tagptr", + "uuid", +] + [[package]] name = "multimap" version = "0.10.1" @@ -3908,6 +3952,12 @@ dependencies = [ "libc", ] +[[package]] +name = "tagptr" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b2093cf4c8eb1e67749a6762251bc9cd836b6fc171623bd0a9d324d37af2417" + [[package]] name = "tempfile" version = "3.22.0" @@ -4277,9 +4327,9 @@ dependencies = [ [[package]] name = "tower" -version = "0.5.2" +version = "0.5.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d039ad9159c98b70ecfd540b2573b97f7f52c3e8d9f8ad57a24b916a536975f9" +checksum = "ebe5ef63511595f1344e2d5cfa636d973292adc0eec1f0ad45fae9f0851ab1d4" dependencies = [ "futures-core", "futures-util", diff --git a/Cargo.toml b/Cargo.toml index 05093a47..d0895d6f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -58,6 +58,10 @@ tracing-subscriber = { version = "0.3", features = ["env-filter"] } futures = "0.3" uuid = { version = "1.17", features = ["v4", "serde"] } thiserror = "2.0" +tower = { version = "0.5.3", features = ["util"] } +http = "1.4.1" +moka = { version = "0.12.15", features = ["future"] } +secrecy = "0.10.3" redis = { version = "0.27", features = ["tokio-comp", "connection-manager"] } reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls", "stream"] } diff --git a/docs/DEPLOYMENT.md b/docs/DEPLOYMENT.md index 8178f5dd..1f3a4ae6 100644 --- a/docs/DEPLOYMENT.md +++ b/docs/DEPLOYMENT.md @@ -175,6 +175,75 @@ GCS uses the configured/default ModelExpress cache root; `MODEL_EXPRESS_CACHE_DI See [`CLI.md`](CLI.md) for full CLI usage documentation. +## ServiceAccount Authentication + +Optional, off by default. When enabled, the server authenticates every gRPC caller +(except health checks) against a Kubernetes ServiceAccount token and authorizes them +against an exact-match allowlist. No sidecar or service mesh is required: the server +calls the Kubernetes `TokenReview` API in-process. + +- **AuthN**: the caller presents a projected ServiceAccount token; the server verifies it + via `TokenReview` and extracts `system:serviceaccount::`. +- **AuthZ**: that `:` must exactly match a configured allowlist + entry. + +### Modes + +| Mode | Behavior | +|------|----------| +| `off` (default) | No auth. Tokens are ignored. | +| `enforce` | Verify every call and reject unauthenticated or non-allowlisted callers. | + +### Server configuration + +| Env Var | Default | Description | +|---------|---------|-------------| +| `MODEL_EXPRESS_SECURITY_MODE` | `off` | `off` \| `enforce` | +| `MODEL_EXPRESS_SECURITY_TOKEN_AUDIENCES` | (none) | Comma-separated audiences the token must carry. Required for `enforce`. | +| `MODEL_EXPRESS_SECURITY_ALLOWED_SERVICE_ACCOUNTS` | (none) | Comma-separated `:` allowlist. Required for `enforce`. | +| `MODEL_EXPRESS_SECURITY_CACHE_TTL_SECS` | `60` | TTL for the verified-token and rejection caches. | + +`enforce` fails config validation if either the audience list or the allowlist is empty, +so a typo can't silently deny-all or accept-any-audience. + +The server's ServiceAccount needs permission to create `TokenReview`s (a cluster-scoped +subresource), via a `ClusterRoleBinding` to the built-in `system:auth-delegator` role. The +Helm chart creates this automatically when `security.enabled=true`: + +```yaml +security: + enabled: true + mode: enforce + tokenAudiences: ["modelexpress"] + allowedServiceAccounts: + - "vllm:worker" + - "vllm:router" +``` + +### Client configuration + +Clients (Rust and Python) attach the token automatically when a projected token file is +present, and send nothing when it is absent (so the same client works against an `off` +server, including off-cluster). Mount a projected token into each worker pod with an +audience that matches the server's, then point the client at it: + +```yaml +volumes: + - name: mx-token + projected: + sources: + - serviceAccountToken: + path: modelexpress + audience: modelexpress + expirationSeconds: 3600 +# mounted at /var/run/secrets/tokens/modelexpress +``` + +| Env Var | Default | Description | +|---------|---------|-------------| +| `MX_AUTH_TOKEN_PATH` | `/var/run/secrets/tokens/modelexpress` | Projected token file path | +| `MX_AUTH_TOKEN_TTL_SECONDS` | `60` | How often to re-read the token (rotation is also picked up on mtime change) | + ## Docker ### Production Image diff --git a/helm/templates/deployment.yaml b/helm/templates/deployment.yaml index 2f5d32de..d114ba1a 100644 --- a/helm/templates/deployment.yaml +++ b/helm/templates/deployment.yaml @@ -65,6 +65,28 @@ spec: {{- if .Values.extraEnv }} {{- toYaml .Values.extraEnv | nindent 12 }} {{- end }} + {{- if .Values.security.enabled }} + {{- if and (eq .Values.security.mode "enforce") (not .Values.security.tokenAudiences) }} + {{- fail "security.mode=enforce requires a non-empty security.tokenAudiences" }} + {{- end }} + {{- if and (eq .Values.security.mode "enforce") (not .Values.security.allowedServiceAccounts) }} + {{- fail "security.mode=enforce requires a non-empty security.allowedServiceAccounts" }} + {{- end }} + - name: MODEL_EXPRESS_SECURITY_MODE + value: {{ .Values.security.mode | quote }} + {{- with .Values.security.tokenAudiences }} + - name: MODEL_EXPRESS_SECURITY_TOKEN_AUDIENCES + value: {{ join "," . | quote }} + {{- end }} + {{- with .Values.security.allowedServiceAccounts }} + - name: MODEL_EXPRESS_SECURITY_ALLOWED_SERVICE_ACCOUNTS + value: {{ join "," . | quote }} + {{- end }} + {{- if hasKey .Values.security "cacheTtlSecs" }} + - name: MODEL_EXPRESS_SECURITY_CACHE_TTL_SECS + value: {{ .Values.security.cacheTtlSecs | quote }} + {{- end }} + {{- end }} {{- if .Values.command }} command: {{- toYaml .Values.command | nindent 12 }} diff --git a/helm/templates/rbac.yaml b/helm/templates/rbac.yaml index 35ea7ac7..a2c66f50 100644 --- a/helm/templates/rbac.yaml +++ b/helm/templates/rbac.yaml @@ -42,3 +42,24 @@ subjects: name: {{ include "modelexpress.serviceAccountName" . }} namespace: {{ .Release.Namespace }} {{- end }} + +# TokenReview access for ServiceAccount auth (security.enabled). The server verifies +# caller tokens by POSTing TokenReviews, a cluster-scoped subresource, so this needs a +# ClusterRoleBinding to the built-in system:auth-delegator role. +{{- if and .Values.serviceAccount.create .Values.security.enabled }} +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRoleBinding +metadata: + name: {{ include "modelexpress.fullname" . }}-auth-delegator + labels: + {{- include "modelexpress.labels" . | nindent 4 }} +roleRef: + apiGroup: rbac.authorization.k8s.io + kind: ClusterRole + name: system:auth-delegator +subjects: + - kind: ServiceAccount + name: {{ include "modelexpress.serviceAccountName" . }} + namespace: {{ .Release.Namespace }} +{{- end }} diff --git a/helm/values.yaml b/helm/values.yaml index f2533e15..6660a169 100644 --- a/helm/values.yaml +++ b/helm/values.yaml @@ -27,6 +27,21 @@ serviceAccount: # Required when MX_METADATA_BACKEND=kubernetes. enabled: false +# ServiceAccount auth (AuthN via Kubernetes TokenReview + namespace/SA allowlist AuthZ). +# Disabled by default. When enabled together with serviceAccount.create, the chart also +# creates a ClusterRoleBinding to system:auth-delegator so the server can POST +# TokenReviews. If you bring your own ServiceAccount (serviceAccount.create=false), bind +# it to system:auth-delegator yourself or every request fails closed. +security: + enabled: false + # off | enforce. enforce rejects unauthenticated or non-allowlisted callers. + mode: enforce + # Audiences the caller's projected token must carry. Required for enforce. + tokenAudiences: [] + # Exact-match allowlist of ":". Required for enforce. + allowedServiceAccounts: [] + # cacheTtlSecs: 60 + podAnnotations: {} podLabels: {} diff --git a/modelexpress_client/python/modelexpress/auth.py b/modelexpress_client/python/modelexpress/auth.py new file mode 100644 index 00000000..a83c231a --- /dev/null +++ b/modelexpress_client/python/modelexpress/auth.py @@ -0,0 +1,152 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Bearer-token auth for ModelExpress gRPC clients; a no-op when no token file is present.""" + +import logging +import os +import threading +import time + +import grpc + +from modelexpress import envs + +logger = logging.getLogger("modelexpress.auth") + +DEFAULT_TOKEN_PATH = "/var/run/secrets/tokens/modelexpress" +DEFAULT_TTL_SECONDS = 60.0 + + +class TokenProvider: + """Caches a projected ServiceAccount token, re-reading on TTL expiry or mtime change.""" + + def __init__(self, path: str | None = None, ttl_seconds: float | None = None): + self._path = path or envs.MX_AUTH_TOKEN_PATH or DEFAULT_TOKEN_PATH + if ttl_seconds is None: + raw_ttl = envs.MX_AUTH_TOKEN_TTL_SECONDS + try: + ttl_seconds = float(raw_ttl) if raw_ttl else DEFAULT_TTL_SECONDS + except ValueError: + logger.warning( + "Invalid %s value %r; using default %ss.", + "MX_AUTH_TOKEN_TTL_SECONDS", + raw_ttl, + DEFAULT_TTL_SECONDS, + ) + ttl_seconds = DEFAULT_TTL_SECONDS + self._ttl = ttl_seconds + self._lock = threading.Lock() + self._cached: str | None = None + self._cached_at: float = 0.0 + self._cached_mtime: float = 0.0 + self._warned_missing = False + + def get(self) -> str | None: + """Return the current token, or None when the token file is absent.""" + now = time.monotonic() + with self._lock: + if self._cached is not None and (now - self._cached_at) < self._ttl: + return self._cached + try: + mtime = os.stat(self._path).st_mtime + if self._cached is not None and mtime == self._cached_mtime: + self._cached_at = now + return self._cached + with open(self._path, encoding="utf-8") as handle: + token = handle.read().strip() + self._cached = token or None + self._cached_at = now + self._cached_mtime = mtime + self._warned_missing = False + return self._cached + except FileNotFoundError: + if not self._warned_missing: + logger.warning( + "Auth token file %s not found; RPCs sent without a bearer token " + "(server must have auth off).", + self._path, + ) + self._warned_missing = True + self._cached = None + return None + except OSError as exc: + if not self._warned_missing: + logger.warning( + "Auth token file %s could not be read (%s); RPCs sent without a " + "bearer token.", + self._path, + exc, + ) + self._warned_missing = True + self._cached = None + return None + + +class _BearerInterceptor( + grpc.UnaryUnaryClientInterceptor, + grpc.UnaryStreamClientInterceptor, + grpc.StreamUnaryClientInterceptor, + grpc.StreamStreamClientInterceptor, +): + def __init__(self, provider: TokenProvider): + self._provider = provider + + def _with_auth(self, client_call_details): + token = self._provider.get() + if not token: + return client_call_details + metadata = list(client_call_details.metadata or []) + metadata.append(("authorization", f"Bearer {token}")) + return _ClientCallDetails( + client_call_details.method, + client_call_details.timeout, + metadata, + client_call_details.credentials, + client_call_details.wait_for_ready, + client_call_details.compression, + ) + + def intercept_unary_unary(self, continuation, client_call_details, request): + return continuation(self._with_auth(client_call_details), request) + + def intercept_unary_stream(self, continuation, client_call_details, request): + return continuation(self._with_auth(client_call_details), request) + + def intercept_stream_unary(self, continuation, client_call_details, request_iterator): + return continuation(self._with_auth(client_call_details), request_iterator) + + def intercept_stream_stream(self, continuation, client_call_details, request_iterator): + return continuation(self._with_auth(client_call_details), request_iterator) + + +class _ClientCallDetails(grpc.ClientCallDetails): + """Concrete carrier; the gRPC namedtuple is not public API.""" + + def __init__(self, method, timeout, metadata, credentials, wait_for_ready, compression): + self.method = method + self.timeout = timeout + self.metadata = metadata + self.credentials = credentials + self.wait_for_ready = wait_for_ready + self.compression = compression + + +_SHARED_PROVIDER: TokenProvider | None = None +_SHARED_LOCK = threading.Lock() + + +def shared_provider() -> TokenProvider: + """Process-wide token provider, created on first use.""" + global _SHARED_PROVIDER + with _SHARED_LOCK: + if _SHARED_PROVIDER is None: + _SHARED_PROVIDER = TokenProvider() + return _SHARED_PROVIDER + + +def with_auth(channel: grpc.Channel, provider: TokenProvider | None = None) -> grpc.Channel: + """Wrap a channel so every RPC carries the bearer token (no-op when absent).""" + return grpc.intercept_channel( + channel, _BearerInterceptor(provider or shared_provider()) + ) diff --git a/modelexpress_client/python/modelexpress/client.py b/modelexpress_client/python/modelexpress/client.py index c48ab65b..abd040df 100644 --- a/modelexpress_client/python/modelexpress/client.py +++ b/modelexpress_client/python/modelexpress/client.py @@ -17,6 +17,7 @@ import grpc +from . import auth from . import envs from . import p2p_pb2 from . import p2p_pb2_grpc @@ -148,7 +149,9 @@ def stub(self) -> p2p_pb2_grpc.P2pServiceStub: ("grpc.max_send_message_length", self._max_message_size), ("grpc.max_receive_message_length", self._max_message_size), ] - self._channel = grpc.insecure_channel(self.server_url, options=options) + self._channel = auth.with_auth( + grpc.insecure_channel(self.server_url, options=options) + ) self._stub = p2p_pb2_grpc.P2pServiceStub(self._channel) logger.debug("MxClient connected to %s", self.server_url) return self._stub diff --git a/modelexpress_client/python/modelexpress/envs.py b/modelexpress_client/python/modelexpress/envs.py index 64e1b27d..9a6889b1 100644 --- a/modelexpress_client/python/modelexpress/envs.py +++ b/modelexpress_client/python/modelexpress/envs.py @@ -44,6 +44,9 @@ MX_SERVER_ADDRESS: Optional[str] MODEL_EXPRESS_LOG_LEVEL: str MODEL_NAME: Optional[str] + # Auth (client) + MX_AUTH_TOKEN_PATH: Optional[str] + MX_AUTH_TOKEN_TTL_SECONDS: Optional[str] # Metadata / worker MX_METADATA_BACKEND: str MX_METADATA_PORT: int @@ -143,6 +146,9 @@ def _env_float(name: str, default: float) -> float: "MX_SERVER_ADDRESS": lambda: os.environ.get("MX_SERVER_ADDRESS"), "MODEL_EXPRESS_LOG_LEVEL": lambda: os.environ.get("MODEL_EXPRESS_LOG_LEVEL", "").upper(), "MODEL_NAME": lambda: os.environ.get("MODEL_NAME"), + # ── Auth (client) ────────────────────────────────────────────────────── + "MX_AUTH_TOKEN_PATH": lambda: os.environ.get("MX_AUTH_TOKEN_PATH"), + "MX_AUTH_TOKEN_TTL_SECONDS": lambda: os.environ.get("MX_AUTH_TOKEN_TTL_SECONDS"), # ── Metadata / worker ────────────────────────────────────────────────── "MX_METADATA_BACKEND": lambda: os.environ.get("MX_METADATA_BACKEND", "").lower().strip(), "MX_METADATA_PORT": lambda: _env_int("MX_METADATA_PORT", 5555), diff --git a/modelexpress_client/python/tests/test_auth.py b/modelexpress_client/python/tests/test_auth.py new file mode 100644 index 00000000..86b11c83 --- /dev/null +++ b/modelexpress_client/python/tests/test_auth.py @@ -0,0 +1,117 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Bearer-token interceptor tests: real temp-file token reads against an in-process server.""" + +import logging +import os +from concurrent.futures import ThreadPoolExecutor + +import grpc +import pytest + +from modelexpress.auth import TokenProvider, with_auth + +_METHOD = "/test.Echo/Call" + + +class _CaptureHandler(grpc.GenericRpcHandler): + """Generic handler that records invocation metadata for any method.""" + + def __init__(self): + self.last_metadata: list[tuple[str, str]] | None = None + + def service(self, handler_call_details): + def handler(request, context): + self.last_metadata = list(context.invocation_metadata()) + return b"ok" + + return grpc.unary_unary_rpc_method_handler(handler) + + +def _serve(handler: _CaptureHandler) -> tuple[grpc.Server, int]: + server = grpc.server(ThreadPoolExecutor(max_workers=1)) + server.add_generic_rpc_handlers((handler,)) + port = server.add_insecure_port("127.0.0.1:0") + server.start() + return server, port + + +def _call(channel: grpc.Channel) -> None: + rpc = channel.unary_unary(_METHOD) + rpc(b"hi", timeout=5) + + +def test_token_provider_reads_file(tmp_path): + token_file = tmp_path / "token" + token_file.write_text("abc.def") + provider = TokenProvider(path=str(token_file)) + assert provider.get() == "abc.def" + + +def test_token_provider_refreshes_on_change(tmp_path): + token_file = tmp_path / "token" + token_file.write_text("tok1") + # ttl=0 forces a re-read every call so we can observe rotation. + provider = TokenProvider(path=str(token_file), ttl_seconds=0.0) + assert provider.get() == "tok1" + + token_file.write_text("tok2") + later = os.stat(str(token_file)).st_mtime + 10 + os.utime(str(token_file), (later, later)) + assert provider.get() == "tok2" + + +def test_token_provider_missing_file_returns_none_and_warns_once(tmp_path, caplog): + provider = TokenProvider(path=str(tmp_path / "absent")) + with caplog.at_level(logging.WARNING, logger="modelexpress.auth"): + assert provider.get() is None + assert provider.get() is None + warnings = [r for r in caplog.records if r.name == "modelexpress.auth"] + assert len(warnings) == 1 + + +def test_token_provider_invalid_ttl_env_falls_back_to_default(monkeypatch, caplog): + monkeypatch.setenv("MX_AUTH_TOKEN_TTL_SECONDS", "not-a-number") + with caplog.at_level(logging.WARNING, logger="modelexpress.auth"): + provider = TokenProvider(path="/unused") + assert provider._ttl == 60.0 + assert any("MX_AUTH_TOKEN_TTL_SECONDS" in r.getMessage() for r in caplog.records) + + +def test_interceptor_attaches_bearer_when_token_present(tmp_path): + token_file = tmp_path / "token" + token_file.write_text("secret-token") + provider = TokenProvider(path=str(token_file)) + + handler = _CaptureHandler() + server, port = _serve(handler) + try: + channel = with_auth(grpc.insecure_channel(f"127.0.0.1:{port}"), provider) + _call(channel) + finally: + server.stop(0) + + assert handler.last_metadata is not None + metadata = dict(handler.last_metadata) + assert metadata.get("authorization") == "Bearer secret-token" + + +def test_interceptor_omits_header_when_token_absent(tmp_path): + provider = TokenProvider(path=str(tmp_path / "absent")) + + handler = _CaptureHandler() + server, port = _serve(handler) + try: + channel = with_auth(grpc.insecure_channel(f"127.0.0.1:{port}"), provider) + _call(channel) + finally: + server.stop(0) + + assert handler.last_metadata is not None + metadata = dict(handler.last_metadata) + assert "authorization" not in metadata + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-v"])) diff --git a/modelexpress_client/src/auth.rs b/modelexpress_client/src/auth.rs new file mode 100644 index 00000000..7196763e --- /dev/null +++ b/modelexpress_client/src/auth.rs @@ -0,0 +1,288 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Attaches a Kubernetes projected ServiceAccount token as `authorization: Bearer` on +//! every RPC, re-reading on TTL/mtime change. A no-op when no token file is present. + +use std::path::PathBuf; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant, SystemTime}; + +use tonic::metadata::MetadataValue; +use tonic::service::Interceptor; +use tonic::{Request, Status}; +use tracing::warn; + +const DEFAULT_TOKEN_PATH: &str = "/var/run/secrets/tokens/modelexpress"; +const DEFAULT_TTL: Duration = Duration::from_secs(60); + +#[derive(Default)] +struct CachedToken { + token: Option, + read_at: Option, + mtime: Option, + warned_missing: bool, +} + +pub struct TokenProvider { + path: PathBuf, + ttl: Duration, + cache: Mutex, +} + +impl TokenProvider { + #[must_use] + pub fn from_env() -> Self { + let path = std::env::var(modelexpress_common::envs::MX_AUTH_TOKEN_PATH) + .map(PathBuf::from) + .unwrap_or_else(|_| PathBuf::from(DEFAULT_TOKEN_PATH)); + let ttl = std::env::var(modelexpress_common::envs::MX_AUTH_TOKEN_TTL_SECONDS) + .ok() + .and_then(|raw| raw.parse::().ok()) + .map_or(DEFAULT_TTL, Duration::from_secs); + Self { + path, + ttl, + cache: Mutex::new(CachedToken::default()), + } + } + + fn token(&self) -> Option { + let now = Instant::now(); + let mut cache = self + .cache + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + + if let Some(read_at) = cache.read_at + && now.duration_since(read_at) < self.ttl + { + return cache.token.clone(); + } + + match std::fs::metadata(&self.path).and_then(|meta| meta.modified()) { + Ok(mtime) => { + if cache.token.is_some() && cache.mtime == Some(mtime) { + cache.read_at = Some(now); + return cache.token.clone(); + } + match std::fs::read_to_string(&self.path) { + Ok(contents) => { + let token = contents.trim(); + cache.token = (!token.is_empty()).then(|| token.to_string()); + cache.read_at = Some(now); + cache.mtime = Some(mtime); + cache.warned_missing = false; + cache.token.clone() + } + Err(e) => { + Self::warn_once(&mut cache, &format!("failed to read token: {e}")); + None + } + } + } + Err(_) => { + Self::warn_once( + &mut cache, + "token file not found; RPCs sent without a bearer token \ + (server must have auth off)", + ); + cache.token = None; + cache.mtime = None; + None + } + } + } + + fn warn_once(cache: &mut CachedToken, message: &str) { + if !cache.warned_missing { + warn!("modelexpress auth: {message}"); + cache.warned_missing = true; + } + } +} + +#[derive(Clone)] +pub struct AuthInterceptor { + provider: Arc, +} + +impl AuthInterceptor { + #[must_use] + pub fn new(provider: Arc) -> Self { + Self { provider } + } +} + +impl Interceptor for AuthInterceptor { + fn call(&mut self, mut request: Request<()>) -> Result, Status> { + if let Some(token) = self.provider.token() { + let value: MetadataValue<_> = format!("Bearer {token}") + .parse() + .map_err(|_| Status::internal("auth token is not a valid header value"))?; + request.metadata_mut().insert("authorization", value); + } + Ok(request) + } +} + +#[cfg(test)] +#[allow(clippy::expect_used)] +mod tests { + use super::*; + use std::fs; + use std::io::Write; + + #[test] + fn reads_token_from_file() { + let mut file = tempfile::NamedTempFile::new().expect("temp file"); + writeln!(file, " my-token").expect("write"); + let provider = TokenProvider { + path: file.path().to_path_buf(), + ttl: DEFAULT_TTL, + cache: Mutex::new(CachedToken::default()), + }; + assert_eq!(provider.token().as_deref(), Some("my-token")); + } + + #[test] + fn missing_file_is_none() { + let provider = TokenProvider { + path: PathBuf::from("/no/such/token/file"), + ttl: DEFAULT_TTL, + cache: Mutex::new(CachedToken::default()), + }; + assert!(provider.token().is_none()); + } + + #[test] + fn empty_file_is_none() { + let file = tempfile::NamedTempFile::new().expect("temp file"); + let provider = TokenProvider { + path: file.path().to_path_buf(), + ttl: DEFAULT_TTL, + cache: Mutex::new(CachedToken::default()), + }; + assert!(provider.token().is_none()); + } + + #[test] + fn cached_token_is_returned_until_ttl_expires() { + let mut file = tempfile::NamedTempFile::new().expect("temp file"); + writeln!(file, "old-token").expect("write old token"); + let provider = TokenProvider { + path: file.path().to_path_buf(), + ttl: DEFAULT_TTL, + cache: Mutex::new(CachedToken::default()), + }; + + assert_eq!(provider.token().as_deref(), Some("old-token")); + fs::write(file.path(), "new-token\n").expect("write new token"); + + assert_eq!(provider.token().as_deref(), Some("old-token")); + } + + #[test] + fn cached_token_reuses_mtime_when_ttl_expires_without_file_change() { + let mut file = tempfile::NamedTempFile::new().expect("temp file"); + writeln!(file, "same-token").expect("write token"); + let provider = TokenProvider { + path: file.path().to_path_buf(), + ttl: Duration::ZERO, + cache: Mutex::new(CachedToken::default()), + }; + + assert_eq!(provider.token().as_deref(), Some("same-token")); + let first_read_at = provider + .cache + .lock() + .expect("cache lock") + .read_at + .expect("read timestamp"); + + assert_eq!(provider.token().as_deref(), Some("same-token")); + let second_read_at = provider + .cache + .lock() + .expect("cache lock") + .read_at + .expect("read timestamp"); + assert!(second_read_at >= first_read_at); + } + + #[test] + fn token_is_reread_after_ttl_when_mtime_changes() { + let mut file = tempfile::NamedTempFile::new().expect("temp file"); + writeln!(file, "old-token").expect("write old token"); + let provider = TokenProvider { + path: file.path().to_path_buf(), + ttl: Duration::ZERO, + cache: Mutex::new(CachedToken::default()), + }; + + assert_eq!(provider.token().as_deref(), Some("old-token")); + fs::write(file.path(), "new-token\n").expect("write new token"); + provider.cache.lock().expect("cache lock").mtime = Some(SystemTime::UNIX_EPOCH); + + assert_eq!(provider.token().as_deref(), Some("new-token")); + } + + #[test] + fn interceptor_attaches_authorization_metadata() { + let mut file = tempfile::NamedTempFile::new().expect("temp file"); + writeln!(file, "rpc-token").expect("write token"); + let provider = Arc::new(TokenProvider { + path: file.path().to_path_buf(), + ttl: DEFAULT_TTL, + cache: Mutex::new(CachedToken::default()), + }); + let mut interceptor = AuthInterceptor::new(provider); + + let request = interceptor + .call(Request::new(())) + .expect("interceptor request"); + + assert_eq!( + request + .metadata() + .get("authorization") + .and_then(|value| value.to_str().ok()), + Some("Bearer rpc-token") + ); + } + + #[test] + fn interceptor_rejects_invalid_header_value() { + let mut file = tempfile::NamedTempFile::new().expect("temp file"); + writeln!(file, "bad\nvalue").expect("write token"); + let provider = Arc::new(TokenProvider { + path: file.path().to_path_buf(), + ttl: DEFAULT_TTL, + cache: Mutex::new(CachedToken::default()), + }); + let mut interceptor = AuthInterceptor::new(provider); + + let status = interceptor + .call(Request::new(())) + .expect_err("invalid metadata value"); + + assert_eq!(status.code(), tonic::Code::Internal); + } + + #[test] + fn missing_file_is_rechecked_immediately() { + let dir = tempfile::tempdir().expect("temp dir"); + let path = dir.path().join("late-mounted-token"); + let provider = TokenProvider { + path: path.clone(), + ttl: DEFAULT_TTL, + cache: Mutex::new(CachedToken::default()), + }; + + assert!(provider.token().is_none()); + + fs::write(&path, "appeared-token\n").expect("write token"); + + assert_eq!(provider.token().as_deref(), Some("appeared-token")); + } +} diff --git a/modelexpress_client/src/lib.rs b/modelexpress_client/src/lib.rs index 52a9f4e0..e83f3d22 100644 --- a/modelexpress_client/src/lib.rs +++ b/modelexpress_client/src/lib.rs @@ -27,15 +27,23 @@ use tracing::info; use tracing::warn; use uuid::Uuid; +pub mod auth; + // Re-export for public use pub use modelexpress_common::client_config::{ClientArgs, ClientConfig}; pub use modelexpress_common::models::ModelProvider; +use auth::{AuthInterceptor, TokenProvider}; +use std::sync::Arc; +use tonic::service::interceptor::InterceptedService; + +type AuthChannel = InterceptedService; + /// The main client for interacting with the `modelexpress_server` via gRPC pub struct Client { - health_client: HealthServiceClient, - api_client: ApiServiceClient, - model_client: ModelServiceClient, + health_client: HealthServiceClient, + api_client: ApiServiceClient, + model_client: ModelServiceClient, cache_config: Option, } @@ -176,9 +184,11 @@ impl Client { .connect() .await?; - let health_client = HealthServiceClient::new(channel.clone()); - let api_client = ApiServiceClient::new(channel.clone()); - let model_client = ModelServiceClient::new(channel); + let interceptor = AuthInterceptor::new(Arc::new(TokenProvider::from_env())); + let health_client = + HealthServiceClient::with_interceptor(channel.clone(), interceptor.clone()); + let api_client = ApiServiceClient::with_interceptor(channel.clone(), interceptor.clone()); + let model_client = ModelServiceClient::with_interceptor(channel, interceptor); // Use the cache config from the client configuration let cache_config = Some(config.cache.clone()); diff --git a/modelexpress_common/src/envs.rs b/modelexpress_common/src/envs.rs index 33520b91..7eddd11f 100644 --- a/modelexpress_common/src/envs.rs +++ b/modelexpress_common/src/envs.rs @@ -116,6 +116,23 @@ pub const MX_HEARTBEAT_TIMEOUT_SECS: &str = "MX_HEARTBEAT_TIMEOUT_SECS"; /// Age (seconds) after which a STALE worker is garbage-collected. pub const MX_GC_TIMEOUT_SECS: &str = "MX_GC_TIMEOUT_SECS"; +// ── Security / auth (server) ──────────────────────────────────────────────── +/// ServiceAccount auth mode (`off`, `enforce`). Off by default. +pub const MODEL_EXPRESS_SECURITY_MODE: &str = "MODEL_EXPRESS_SECURITY_MODE"; +/// Comma-separated SA token audiences the caller's token must carry. +pub const MODEL_EXPRESS_SECURITY_TOKEN_AUDIENCES: &str = "MODEL_EXPRESS_SECURITY_TOKEN_AUDIENCES"; +/// Comma-separated allowed callers as `:`. +pub const MODEL_EXPRESS_SECURITY_ALLOWED_SERVICE_ACCOUNTS: &str = + "MODEL_EXPRESS_SECURITY_ALLOWED_SERVICE_ACCOUNTS"; +/// TTL for the verified-token and rejection caches, in seconds. +pub const MODEL_EXPRESS_SECURITY_CACHE_TTL_SECS: &str = "MODEL_EXPRESS_SECURITY_CACHE_TTL_SECS"; + +// ── Auth (client) ─────────────────────────────────────────────────────────── +/// Path to the Kubernetes projected ServiceAccount token file. +pub const MX_AUTH_TOKEN_PATH: &str = "MX_AUTH_TOKEN_PATH"; +/// TTL in seconds for the cached token. +pub const MX_AUTH_TOKEN_TTL_SECONDS: &str = "MX_AUTH_TOKEN_TTL_SECONDS"; + // ── System ────────────────────────────────────────────────────────────────── /// Primary source for the user's home directory. pub const HOME: &str = "HOME"; @@ -321,6 +338,21 @@ mod tests { assert_eq!(MX_REAPER_SCAN_INTERVAL_SECS, "MX_REAPER_SCAN_INTERVAL_SECS"); assert_eq!(MX_HEARTBEAT_TIMEOUT_SECS, "MX_HEARTBEAT_TIMEOUT_SECS"); assert_eq!(MX_GC_TIMEOUT_SECS, "MX_GC_TIMEOUT_SECS"); + assert_eq!(MODEL_EXPRESS_SECURITY_MODE, "MODEL_EXPRESS_SECURITY_MODE"); + assert_eq!( + MODEL_EXPRESS_SECURITY_TOKEN_AUDIENCES, + "MODEL_EXPRESS_SECURITY_TOKEN_AUDIENCES" + ); + assert_eq!( + MODEL_EXPRESS_SECURITY_ALLOWED_SERVICE_ACCOUNTS, + "MODEL_EXPRESS_SECURITY_ALLOWED_SERVICE_ACCOUNTS" + ); + assert_eq!( + MODEL_EXPRESS_SECURITY_CACHE_TTL_SECS, + "MODEL_EXPRESS_SECURITY_CACHE_TTL_SECS" + ); + assert_eq!(MX_AUTH_TOKEN_PATH, "MX_AUTH_TOKEN_PATH"); + assert_eq!(MX_AUTH_TOKEN_TTL_SECONDS, "MX_AUTH_TOKEN_TTL_SECONDS"); assert_eq!(HOME, "HOME"); assert_eq!(USERPROFILE, "USERPROFILE"); assert_eq!(KUBECONFIG, "KUBECONFIG"); diff --git a/modelexpress_server/Cargo.toml b/modelexpress_server/Cargo.toml index ef165dbd..a69a4c20 100644 --- a/modelexpress_server/Cargo.toml +++ b/modelexpress_server/Cargo.toml @@ -52,6 +52,11 @@ schemars = { workspace = true } base64 = { workspace = true } sha2 = { workspace = true } tonic-health = { workspace = true } +thiserror = { workspace = true } +tower = { workspace = true } +http = { workspace = true } +moka = { workspace = true } +secrecy = { workspace = true } [dev-dependencies] modelexpress-common = { workspace = true, features = ["test-support"] } @@ -66,5 +71,10 @@ name = "in_process_server" path = "tests/in_process_server.rs" required-features = ["integration-tests"] +[[test]] +name = "auth_e2e" +path = "tests/auth_e2e.rs" +required-features = ["integration-tests"] + [lints] workspace = true diff --git a/modelexpress_server/src/auth.rs b/modelexpress_server/src/auth.rs new file mode 100644 index 00000000..6d700cc4 --- /dev/null +++ b/modelexpress_server/src/auth.rs @@ -0,0 +1,396 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! ServiceAccount `TokenReview` authentication plus an exact-match allowlist. + +mod layer; +mod token; + +pub use layer::AuthLayer; +pub use token::CallerIdentity; + +use std::collections::HashSet; +use std::fmt; +use std::hash::{Hash, Hasher}; +use std::time::Duration; + +use http::HeaderMap; +use kube::client::Client; +use moka::future::Cache; +use secrecy::ExposeSecret; +use secrecy::zeroize::Zeroizing; +use tracing::warn; + +use crate::config::{SecurityConfig, ServiceAccountRef}; +use token::{extract_bearer, review_token}; + +#[derive(Debug, thiserror::Error)] +pub enum Denial { + #[error("missing or invalid service account token")] + Unauthenticated, + #[error("{0}")] + PermissionDenied(String), + #[error("authentication backend unavailable")] + Unavailable, +} + +impl Denial { + pub(crate) fn into_status(self) -> tonic::Status { + match self { + Self::Unauthenticated => { + tonic::Status::unauthenticated("missing or invalid service account token") + } + Self::PermissionDenied(message) => tonic::Status::permission_denied(message), + Self::Unavailable => tonic::Status::unavailable("authentication backend unavailable"), + } + } +} + +#[derive(Clone)] +struct TokenCacheKey(Zeroizing); + +impl TokenCacheKey { + fn new(token: &str) -> Self { + Self(Zeroizing::new(token.to_string())) + } + + fn as_str(&self) -> &str { + self.0.as_str() + } +} + +impl fmt::Debug for TokenCacheKey { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str("TokenCacheKey([REDACTED])") + } +} + +impl PartialEq for TokenCacheKey { + fn eq(&self, other: &Self) -> bool { + self.as_str() == other.as_str() + } +} + +impl Eq for TokenCacheKey {} + +impl Hash for TokenCacheKey { + fn hash(&self, state: &mut H) { + self.as_str().hash(state); + } +} + +pub struct AuthState { + client: Client, + audiences: Vec, + allowlist: HashSet, + token_cache: Cache, + negative_cache: Cache, +} + +impl AuthState { + const TOKEN_CACHE_MAX_ENTRIES: u64 = 10_000; + + #[must_use] + pub fn new(client: Client, config: &SecurityConfig) -> Self { + let ttl = Duration::from_secs(config.cache_ttl_secs); + Self { + client, + audiences: config.token_audiences.clone(), + allowlist: config.allowed_service_accounts.iter().cloned().collect(), + token_cache: Cache::builder() + .time_to_live(ttl) + .max_capacity(Self::TOKEN_CACHE_MAX_ENTRIES) + .build(), + negative_cache: Cache::builder() + .time_to_live(ttl) + .max_capacity(Self::TOKEN_CACHE_MAX_ENTRIES) + .build(), + } + } + + pub(crate) async fn verify(&self, headers: &HeaderMap) -> Result { + let token = extract_bearer(headers).ok_or(Denial::Unauthenticated)?; + let key = TokenCacheKey::new(token.expose_secret()); + + if self.negative_cache.get(&key).await.is_some() { + return Err(Denial::Unauthenticated); + } + + let client = self.client.clone(); + let audiences = self.audiences.clone(); + let key_for_review = key.clone(); + let identity = match self + .token_cache + .try_get_with(key.clone(), async move { + review_token(&client, key_for_review.as_str(), &audiences).await + }) + .await + { + Ok(identity) => identity, + Err(error) => { + if error.is_token_rejection() { + self.negative_cache.insert(key, ()).await; + return Err(Denial::Unauthenticated); + } + warn!(error = %error, "TokenReview backend error"); + return Err(Denial::Unavailable); + } + }; + + let caller = ServiceAccountRef { + namespace: identity.namespace.clone(), + service_account: identity.service_account.clone(), + }; + if !self.allowlist.contains(&caller) { + return Err(Denial::PermissionDenied(format!( + "service account {}:{} is not in the allowlist", + identity.namespace, identity.service_account + ))); + } + + Ok(identity) + } +} + +#[cfg(test)] +#[allow(clippy::expect_used)] +pub(crate) mod test_util { + use std::collections::VecDeque; + use std::sync::{ + Arc, Mutex, + atomic::{AtomicUsize, Ordering}, + }; + + use http::HeaderMap; + use http::header::{AUTHORIZATION, CONTENT_TYPE}; + use k8s_openapi::api::authentication::v1::{TokenReview, TokenReviewStatus, UserInfo}; + use kube::client::{Body, Client}; + use tower::BoxError; + + use crate::config::{AuthMode, SecurityConfig, ServiceAccountRef}; + + pub(crate) fn allowed_ref(namespace: &str, service_account: &str) -> ServiceAccountRef { + ServiceAccountRef { + namespace: namespace.to_string(), + service_account: service_account.to_string(), + } + } + + pub(crate) fn security_config(allowed: Vec) -> SecurityConfig { + SecurityConfig { + mode: Some(AuthMode::Enforce), + token_audiences: vec!["modelexpress".to_string()], + allowed_service_accounts: allowed, + cache_ttl_secs: 60, + } + } + + pub(crate) fn bearer_headers(token: &str) -> HeaderMap { + let mut headers = HeaderMap::new(); + headers.insert( + AUTHORIZATION, + format!("Bearer {token}") + .parse() + .expect("valid bearer header"), + ); + headers + } + + pub(crate) fn token_review_for(username: &str) -> TokenReview { + TokenReview { + status: Some(TokenReviewStatus { + authenticated: Some(true), + audiences: Some(vec!["modelexpress".to_string()]), + user: Some(UserInfo { + username: Some(username.to_string()), + ..Default::default() + }), + ..Default::default() + }), + ..Default::default() + } + } + + pub(crate) fn unauthenticated_review(error: &str) -> TokenReview { + TokenReview { + status: Some(TokenReviewStatus { + authenticated: Some(false), + error: Some(error.to_string()), + ..Default::default() + }), + ..Default::default() + } + } + + pub(crate) fn review_without_user() -> TokenReview { + TokenReview { + status: Some(TokenReviewStatus { + authenticated: Some(true), + ..Default::default() + }), + ..Default::default() + } + } + + pub(crate) enum FakeResponse { + Review(Box), + Error(http::StatusCode), + } + + pub(crate) fn fake_kube_client(reviews: Vec) -> (Client, Arc) { + fake_kube_client_with_responses( + reviews + .into_iter() + .map(|r| FakeResponse::Review(Box::new(r))) + .collect(), + ) + } + + pub(crate) fn fake_kube_client_with_responses( + responses: Vec, + ) -> (Client, Arc) { + let responses = Arc::new(Mutex::new(VecDeque::from(responses))); + let calls = Arc::new(AtomicUsize::new(0)); + + let service = tower::service_fn({ + let responses = responses.clone(); + let calls = calls.clone(); + move |_request: http::Request| { + let responses = responses.clone(); + let calls = calls.clone(); + async move { + calls.fetch_add(1, Ordering::SeqCst); + let response_item = responses + .lock() + .expect("fake response queue lock") + .pop_front() + .expect("fake response"); + let response = match response_item { + FakeResponse::Review(review) => { + let body = serde_json::to_vec(&*review).expect("serialize TokenReview"); + http::Response::builder() + .status(http::StatusCode::OK) + .header(CONTENT_TYPE, "application/json") + .body(Body::from(body)) + .expect("fake TokenReview response body") + } + FakeResponse::Error(status) => http::Response::builder() + .status(status) + .body(Body::empty()) + .expect("fake error response body"), + }; + Ok::<_, BoxError>(response) + } + } + }); + + (Client::new(service, "default"), calls) + } +} + +#[cfg(test)] +#[allow(clippy::expect_used)] +mod tests { + use super::*; + use std::sync::atomic::Ordering; + + use test_util::{ + allowed_ref, bearer_headers, fake_kube_client, security_config, token_review_for, + unauthenticated_review, + }; + + #[test] + fn unauthenticated_maps_to_grpc_code() { + let status = Denial::Unauthenticated.into_status(); + assert_eq!(status.code(), tonic::Code::Unauthenticated); + } + + #[test] + fn permission_denied_maps_to_grpc_code() { + let status = Denial::PermissionDenied("nope".to_string()).into_status(); + assert_eq!(status.code(), tonic::Code::PermissionDenied); + assert_eq!(status.message(), "nope"); + } + + #[test] + fn unavailable_maps_to_grpc_code() { + let status = Denial::Unavailable.into_status(); + assert_eq!(status.code(), tonic::Code::Unavailable); + } + + #[tokio::test] + async fn verify_accepts_allowlisted_identity_and_uses_positive_cache() { + let (client, calls) = + fake_kube_client(vec![token_review_for("system:serviceaccount:vllm:worker")]); + let config = security_config(vec![allowed_ref("vllm", "worker")]); + let state = AuthState::new(client, &config); + let headers = bearer_headers("good-token"); + + let first = state.verify(&headers).await.expect("first token review"); + let second = state.verify(&headers).await.expect("cached token review"); + + assert_eq!(first.namespace, "vllm"); + assert_eq!(first.service_account, "worker"); + assert_eq!(second.namespace, "vllm"); + assert_eq!(calls.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn verify_rejects_identity_outside_allowlist() { + let (client, _calls) = + fake_kube_client(vec![token_review_for("system:serviceaccount:other:worker")]); + let config = security_config(vec![allowed_ref("vllm", "worker")]); + let state = AuthState::new(client, &config); + let headers = bearer_headers("other-token"); + + let error = state + .verify(&headers) + .await + .expect_err("non-allowlisted service account should be denied"); + + assert!( + matches!(error, Denial::PermissionDenied(message) if message.contains("other:worker")) + ); + } + + #[tokio::test] + async fn verify_negative_caches_definitive_token_rejections() { + let (client, calls) = fake_kube_client(vec![unauthenticated_review("expired")]); + let config = security_config(vec![allowed_ref("vllm", "worker")]); + let state = AuthState::new(client, &config); + let headers = bearer_headers("expired-token"); + + assert!(matches!( + state.verify(&headers).await, + Err(Denial::Unauthenticated) + )); + assert!(matches!( + state.verify(&headers).await, + Err(Denial::Unauthenticated) + )); + assert_eq!(calls.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn verify_returns_unavailable_for_api_errors_without_caching() { + use test_util::FakeResponse; + + let (client, calls) = test_util::fake_kube_client_with_responses(vec![ + FakeResponse::Error(http::StatusCode::INTERNAL_SERVER_ERROR), + FakeResponse::Error(http::StatusCode::INTERNAL_SERVER_ERROR), + ]); + let config = security_config(vec![allowed_ref("vllm", "worker")]); + let state = AuthState::new(client, &config); + let headers = bearer_headers("token"); + + assert!(matches!( + state.verify(&headers).await, + Err(Denial::Unavailable) + )); + assert!(matches!( + state.verify(&headers).await, + Err(Denial::Unavailable) + )); + assert_eq!(calls.load(Ordering::SeqCst), 2); + } +} diff --git a/modelexpress_server/src/auth/layer.rs b/modelexpress_server/src/auth/layer.rs new file mode 100644 index 00000000..b3a89bb6 --- /dev/null +++ b/modelexpress_server/src/auth/layer.rs @@ -0,0 +1,201 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Tower middleware that authenticates a wrapped gRPC service via [`AuthState::verify`]. + +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use http::{Request, Response}; +use tonic::body::Body; +use tonic::server::NamedService; +use tower::{Layer, Service}; +use tracing::{debug, warn}; + +use crate::auth::AuthState; + +#[derive(Clone)] +pub struct AuthLayer { + state: Arc, +} + +impl AuthLayer { + #[must_use] + pub fn new(state: Arc) -> Self { + Self { state } + } +} + +impl Layer for AuthLayer { + type Service = AuthService; + + fn layer(&self, inner: S) -> Self::Service { + AuthService { + inner, + state: self.state.clone(), + } + } +} + +#[derive(Clone)] +pub struct AuthService { + inner: S, + state: Arc, +} + +impl NamedService for AuthService { + const NAME: &'static str = S::NAME; +} + +impl Service> for AuthService +where + S: Service, Response = Response> + Clone + Send + 'static, + S::Future: Send + 'static, +{ + type Response = Response; + type Error = S::Error; + type Future = Pin, S::Error>> + Send>>; + + fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { + self.inner.poll_ready(cx) + } + + fn call(&mut self, req: Request) -> Self::Future { + let clone = self.inner.clone(); + let mut inner = std::mem::replace(&mut self.inner, clone); + let state = self.state.clone(); + + Box::pin(async move { + match state.verify(req.headers()).await { + Ok(caller) => { + debug!( + namespace = %caller.namespace, + service_account = %caller.service_account, + pod = caller.pod_name.as_deref().unwrap_or("?"), + path = %req.uri().path(), + "auth ok" + ); + let mut req = req; + req.extensions_mut().insert(caller); + inner.call(req).await + } + Err(denial) => { + warn!(reason = %denial, path = %req.uri().path(), "auth denied"); + Ok(denial.into_status().into_http::()) + } + } + }) + } +} + +#[cfg(test)] +#[allow(clippy::expect_used)] +mod tests { + use super::*; + use std::convert::Infallible; + use std::future::{Ready, ready}; + use std::sync::atomic::Ordering; + use std::sync::{Arc, Mutex}; + + use http::header::AUTHORIZATION; + use http::{HeaderValue, StatusCode}; + + use crate::auth::CallerIdentity; + use crate::auth::test_util::{ + allowed_ref, fake_kube_client, security_config, token_review_for, + }; + + struct Dummy; + impl NamedService for Dummy { + const NAME: &'static str = "model_express.p2p.P2pService"; + } + + #[test] + fn forwards_wrapped_service_name() { + assert_eq!( + as NamedService>::NAME, + ::NAME + ); + } + + #[derive(Clone, Default)] + struct RecordingService { + caller: Arc>>, + calls: Arc, + } + + impl Service> for RecordingService { + type Response = Response; + type Error = Infallible; + type Future = Ready, Self::Error>>; + + fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + fn call(&mut self, req: Request) -> Self::Future { + self.calls.fetch_add(1, Ordering::SeqCst); + *self.caller.lock().expect("recording service caller lock") = + req.extensions().get::().cloned(); + ready(Ok(Response::new(Body::empty()))) + } + } + + #[tokio::test] + async fn authenticated_request_reaches_inner_service_with_identity() { + let (client, kube_calls) = + fake_kube_client(vec![token_review_for("system:serviceaccount:vllm:worker")]); + let config = security_config(vec![allowed_ref("vllm", "worker")]); + let state = Arc::new(AuthState::new(client, &config)); + let inner = RecordingService::default(); + let caller = inner.caller.clone(); + let calls = inner.calls.clone(); + let mut service = AuthLayer::new(state).layer(inner); + let request = Request::builder() + .uri("/model_express.ModelService/ListModelFiles") + .header(AUTHORIZATION, HeaderValue::from_static("Bearer good-token")) + .body(Body::empty()) + .expect("request"); + + let response = service.call(request).await.expect("inner service response"); + + assert_eq!(response.status(), StatusCode::OK); + assert_eq!(calls.load(Ordering::SeqCst), 1); + let caller = caller + .lock() + .expect("recording service caller lock") + .clone() + .expect("caller identity extension"); + assert_eq!(caller.namespace, "vllm"); + assert_eq!(caller.service_account, "worker"); + assert_eq!(kube_calls.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn denied_request_returns_grpc_error_without_calling_inner() { + let (client, kube_calls) = fake_kube_client(Vec::new()); + let config = security_config(vec![allowed_ref("vllm", "worker")]); + let state = Arc::new(AuthState::new(client, &config)); + let inner = RecordingService::default(); + let calls = inner.calls.clone(); + let mut service = AuthLayer::new(state).layer(inner); + let request = Request::builder() + .uri("/model_express.ModelService/ListModelFiles") + .body(Body::empty()) + .expect("request"); + + let response = service.call(request).await.expect("denial response"); + + assert_eq!(calls.load(Ordering::SeqCst), 0); + assert_eq!(kube_calls.load(Ordering::SeqCst), 0); + assert_eq!( + response + .headers() + .get("grpc-status") + .and_then(|value| value.to_str().ok()), + Some("16") + ); + } +} diff --git a/modelexpress_server/src/auth/token.rs b/modelexpress_server/src/auth/token.rs new file mode 100644 index 00000000..61024139 --- /dev/null +++ b/modelexpress_server/src/auth/token.rs @@ -0,0 +1,416 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Bearer-token extraction and Kubernetes `TokenReview` verification. + +use http::HeaderMap; +use http::header::AUTHORIZATION; +use k8s_openapi::api::authentication::v1::{TokenReview, TokenReviewSpec, UserInfo}; +use kube::api::{Api, PostParams}; +use kube::client::Client; +use secrecy::SecretString; + +const EXTRA_POD_NAME: &str = "authentication.kubernetes.io/pod-name"; +const EXTRA_POD_UID: &str = "authentication.kubernetes.io/pod-uid"; + +#[derive(Debug, Clone)] +pub struct CallerIdentity { + pub namespace: String, + pub service_account: String, + pub pod_name: Option, + pub pod_uid: Option, +} + +#[derive(Debug, thiserror::Error)] +pub(crate) enum TokenError { + #[error("kube api error during TokenReview: {0}")] + Api(#[source] kube::Error), + #[error("TokenReview returned no status")] + NoStatus, + #[error("token not authenticated: {0:?}")] + NotAuthenticated(Option), + #[error("TokenReview returned no user info")] + NoUser, + #[error("caller is not a kubernetes service account")] + NotServiceAccount, + #[error("token audience mismatch: requested {requested:?}, got {actual:?}")] + AudienceMismatch { + requested: Vec, + actual: Option>, + }, +} + +impl TokenError { + pub(crate) fn is_token_rejection(&self) -> bool { + matches!( + self, + Self::NotAuthenticated(_) | Self::NotServiceAccount | Self::AudienceMismatch { .. } + ) + } +} + +pub(crate) fn extract_bearer(headers: &HeaderMap) -> Option { + let value = headers.get(AUTHORIZATION)?.to_str().ok()?; + let prefix = value.get(..7)?; + if !prefix.eq_ignore_ascii_case("bearer ") { + return None; + } + let raw = value.get(7..)?; + if raw.is_empty() { + return None; + } + Some(SecretString::from(raw)) +} + +pub(crate) fn parse_sa_username(username: &str) -> Option<(String, String)> { + let rest = username.strip_prefix("system:serviceaccount:")?; + let (namespace, service_account) = rest.split_once(':')?; + if namespace.is_empty() || service_account.is_empty() { + return None; + } + Some((namespace.to_string(), service_account.to_string())) +} + +pub(crate) fn caller_from_userinfo(user: &UserInfo) -> Option { + let username = user.username.as_deref()?; + let (namespace, service_account) = parse_sa_username(username)?; + let first = |key: &str| -> Option { + user.extra + .as_ref() + .and_then(|map| map.get(key)) + .and_then(|values| values.first().cloned()) + }; + Some(CallerIdentity { + namespace, + service_account, + pod_name: first(EXTRA_POD_NAME), + pod_uid: first(EXTRA_POD_UID), + }) +} + +pub(crate) async fn review_token( + client: &Client, + token: &str, + audiences: &[String], +) -> Result { + let api: Api = Api::all(client.clone()); + let review = TokenReview { + spec: TokenReviewSpec { + token: Some(token.to_string()), + audiences: if audiences.is_empty() { + None + } else { + Some(audiences.to_vec()) + }, + }, + ..Default::default() + }; + let reviewed = api + .create(&PostParams::default(), &review) + .await + .map_err(TokenError::Api)?; + let status = reviewed.status.ok_or(TokenError::NoStatus)?; + if !status.authenticated.unwrap_or(false) { + return Err(TokenError::NotAuthenticated(status.error)); + } + + if !audiences.is_empty() { + let status_audiences = status.audiences.as_deref().unwrap_or(&[]); + let has_overlap = audiences + .iter() + .any(|req| status_audiences.iter().any(|stat| req == stat)); + if !has_overlap { + return Err(TokenError::AudienceMismatch { + requested: audiences.to_vec(), + actual: status.audiences.clone(), + }); + } + } + + let user = status.user.ok_or(TokenError::NoUser)?; + caller_from_userinfo(&user).ok_or(TokenError::NotServiceAccount) +} + +#[cfg(test)] +#[allow(clippy::expect_used, clippy::unwrap_used)] +mod tests { + use super::*; + use crate::auth::test_util::{ + fake_kube_client, review_without_user, token_review_for, unauthenticated_review, + }; + use secrecy::ExposeSecret; + use std::sync::atomic::Ordering; + + #[test] + fn extracts_bearer_token_case_insensitively() { + let mut headers = HeaderMap::new(); + headers.insert(AUTHORIZATION, "Bearer abc.def".parse().unwrap()); + assert_eq!(extract_bearer(&headers).unwrap().expose_secret(), "abc.def"); + + let mut lower = HeaderMap::new(); + lower.insert(AUTHORIZATION, "bearer xyz".parse().unwrap()); + assert_eq!(extract_bearer(&lower).unwrap().expose_secret(), "xyz"); + + let mut upper = HeaderMap::new(); + upper.insert(AUTHORIZATION, "BEARER token123".parse().unwrap()); + assert_eq!(extract_bearer(&upper).unwrap().expose_secret(), "token123"); + + let mut mixed = HeaderMap::new(); + mixed.insert(AUTHORIZATION, "BeArEr abc".parse().unwrap()); + assert_eq!(extract_bearer(&mixed).unwrap().expose_secret(), "abc"); + } + + #[test] + fn rejects_missing_or_malformed_bearer() { + assert!(extract_bearer(&HeaderMap::new()).is_none()); + + let mut basic = HeaderMap::new(); + basic.insert(AUTHORIZATION, "Basic abc".parse().unwrap()); + assert!(extract_bearer(&basic).is_none()); + + let mut empty = HeaderMap::new(); + empty.insert(AUTHORIZATION, "Bearer ".parse().unwrap()); + assert!(extract_bearer(&empty).is_none()); + + let mut short = HeaderMap::new(); + short.insert(AUTHORIZATION, "Bear".parse().unwrap()); + assert!(extract_bearer(&short).is_none()); + } + + #[test] + fn parses_valid_sa_username() { + assert_eq!( + parse_sa_username("system:serviceaccount:foo:bar"), + Some(("foo".to_string(), "bar".to_string())) + ); + } + + #[test] + fn rejects_non_sa_usernames() { + assert_eq!(parse_sa_username("system:node:node-1"), None); + assert_eq!(parse_sa_username("alice"), None); + assert_eq!(parse_sa_username("system:serviceaccount:foo:"), None); + assert_eq!(parse_sa_username("system:serviceaccount::bar"), None); + assert_eq!(parse_sa_username("system:serviceaccount:foo"), None); + } + + #[test] + fn caller_from_userinfo_extracts_pod_identity() { + let mut extra = std::collections::BTreeMap::new(); + extra.insert(EXTRA_POD_NAME.to_string(), vec!["worker-0".to_string()]); + extra.insert(EXTRA_POD_UID.to_string(), vec!["uid-123".to_string()]); + let user = UserInfo { + username: Some("system:serviceaccount:ns:sa".to_string()), + extra: Some(extra), + ..Default::default() + }; + let caller = caller_from_userinfo(&user).unwrap(); + assert_eq!(caller.namespace, "ns"); + assert_eq!(caller.service_account, "sa"); + assert_eq!(caller.pod_name.as_deref(), Some("worker-0")); + assert_eq!(caller.pod_uid.as_deref(), Some("uid-123")); + } + + #[test] + fn caller_from_userinfo_without_extra_has_no_pod_identity() { + let user = UserInfo { + username: Some("system:serviceaccount:ns:sa".to_string()), + ..Default::default() + }; + let caller = caller_from_userinfo(&user).unwrap(); + assert!(caller.pod_name.is_none()); + assert!(caller.pod_uid.is_none()); + } + + #[test] + fn only_definitive_rejections_are_cacheable() { + assert!(TokenError::NotAuthenticated(None).is_token_rejection()); + assert!(TokenError::NotServiceAccount.is_token_rejection()); + assert!( + TokenError::AudienceMismatch { + requested: vec!["modelexpress".to_string()], + actual: None, + } + .is_token_rejection() + ); + assert!(!TokenError::NoStatus.is_token_rejection()); + assert!(!TokenError::NoUser.is_token_rejection()); + } + + #[test] + fn caller_from_userinfo_rejects_non_sa() { + let user = UserInfo { + username: Some("system:node:node-1".to_string()), + ..Default::default() + }; + assert!(caller_from_userinfo(&user).is_none()); + } + + #[tokio::test] + async fn review_token_accepts_service_account_identity() { + let (client, calls) = + fake_kube_client(vec![token_review_for("system:serviceaccount:vllm:worker")]); + let audiences = vec!["modelexpress".to_string()]; + + let caller = review_token(&client, "token", &audiences) + .await + .expect("authenticated TokenReview"); + + assert_eq!(caller.namespace, "vllm"); + assert_eq!(caller.service_account, "worker"); + assert_eq!(calls.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn review_token_rejects_unauthenticated_status() { + let (client, _calls) = fake_kube_client(vec![unauthenticated_review("expired")]); + + let error = review_token(&client, "token", &[]) + .await + .expect_err("unauthenticated TokenReview"); + + assert!(matches!( + error, + TokenError::NotAuthenticated(Some(message)) if message == "expired" + )); + } + + #[tokio::test] + async fn review_token_rejects_missing_status() { + let (client, _calls) = fake_kube_client(vec![TokenReview::default()]); + + let error = review_token(&client, "token", &[]) + .await + .expect_err("missing TokenReview status"); + + assert!(matches!(error, TokenError::NoStatus)); + } + + #[tokio::test] + async fn review_token_rejects_missing_user() { + let (client, _calls) = fake_kube_client(vec![review_without_user()]); + + let error = review_token(&client, "token", &[]) + .await + .expect_err("missing TokenReview user"); + + assert!(matches!(error, TokenError::NoUser)); + } + + #[tokio::test] + async fn review_token_rejects_authenticated_with_no_audiences() { + use k8s_openapi::api::authentication::v1::TokenReviewStatus; + + let review = TokenReview { + status: Some(TokenReviewStatus { + authenticated: Some(true), + audiences: None, + user: Some(UserInfo { + username: Some("system:serviceaccount:vllm:worker".to_string()), + ..Default::default() + }), + ..Default::default() + }), + ..Default::default() + }; + let (client, _calls) = fake_kube_client(vec![review]); + let audiences = vec!["modelexpress".to_string()]; + + let error = review_token(&client, "token", &audiences) + .await + .expect_err("TokenReview with None audiences"); + + assert!(matches!( + error, + TokenError::AudienceMismatch { + requested, + actual: None, + } if requested == audiences + )); + } + + #[tokio::test] + async fn review_token_rejects_mismatched_audiences() { + use k8s_openapi::api::authentication::v1::TokenReviewStatus; + + let review = TokenReview { + status: Some(TokenReviewStatus { + authenticated: Some(true), + audiences: Some(vec!["other".to_string()]), + user: Some(UserInfo { + username: Some("system:serviceaccount:vllm:worker".to_string()), + ..Default::default() + }), + ..Default::default() + }), + ..Default::default() + }; + let (client, _calls) = fake_kube_client(vec![review]); + let audiences = vec!["modelexpress".to_string()]; + + let error = review_token(&client, "token", &audiences) + .await + .expect_err("TokenReview with mismatched audiences"); + + assert!(matches!( + error, + TokenError::AudienceMismatch { + requested, + actual: Some(ref actual_vec), + } if requested == audiences && actual_vec == &["other".to_string()] + )); + } + + #[tokio::test] + async fn review_token_accepts_overlapping_audiences() { + use k8s_openapi::api::authentication::v1::TokenReviewStatus; + + let review = TokenReview { + status: Some(TokenReviewStatus { + authenticated: Some(true), + audiences: Some(vec!["other".to_string(), "modelexpress".to_string()]), + user: Some(UserInfo { + username: Some("system:serviceaccount:vllm:worker".to_string()), + ..Default::default() + }), + ..Default::default() + }), + ..Default::default() + }; + let (client, _calls) = fake_kube_client(vec![review]); + let audiences = vec!["modelexpress".to_string()]; + + let caller = review_token(&client, "token", &audiences) + .await + .expect("overlapping audiences"); + + assert_eq!(caller.namespace, "vllm"); + assert_eq!(caller.service_account, "worker"); + } + + #[tokio::test] + async fn review_token_accepts_empty_requested_audiences_with_none_status() { + use k8s_openapi::api::authentication::v1::TokenReviewStatus; + + let review = TokenReview { + status: Some(TokenReviewStatus { + authenticated: Some(true), + audiences: None, + user: Some(UserInfo { + username: Some("system:serviceaccount:vllm:worker".to_string()), + ..Default::default() + }), + ..Default::default() + }), + ..Default::default() + }; + let (client, _calls) = fake_kube_client(vec![review]); + + let caller = review_token(&client, "token", &[]) + .await + .expect("empty requested audiences"); + + assert_eq!(caller.namespace, "vllm"); + assert_eq!(caller.service_account, "worker"); + } +} diff --git a/modelexpress_server/src/config.rs b/modelexpress_server/src/config.rs index d2f30eb9..4fe98fad 100644 --- a/modelexpress_server/src/config.rs +++ b/modelexpress_server/src/config.rs @@ -12,6 +12,152 @@ use tracing::Level; use crate::cache::CacheEvictionConfig; +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, clap::ValueEnum, Default)] +#[serde(rename_all = "lowercase")] +pub enum AuthMode { + #[default] + Off, + Enforce, +} + +/// A `:` pair, parsed from `ns:sa`. +#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)] +pub struct ServiceAccountRef { + pub namespace: String, + pub service_account: String, +} + +impl std::fmt::Display for ServiceAccountRef { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}:{}", self.namespace, self.service_account) + } +} + +impl std::str::FromStr for ServiceAccountRef { + type Err = String; + + fn from_str(s: &str) -> Result { + let (namespace, service_account) = s + .split_once(':') + .ok_or_else(|| format!("expected ':', got '{s}'"))?; + let namespace = namespace.trim(); + let service_account = service_account.trim(); + if namespace.is_empty() || service_account.is_empty() { + return Err(format!( + "namespace and service account must both be non-empty in '{s}'" + )); + } + Ok(Self { + namespace: namespace.to_string(), + service_account: service_account.to_string(), + }) + } +} + +/// Comma-separated CLI/env list of `T`, ignoring blank entries. +#[derive(Debug, Clone)] +pub struct CommaList(Vec); + +impl CommaList { + #[must_use] + pub fn into_inner(self) -> Vec { + self.0 + } +} + +impl std::str::FromStr for CommaList { + type Err = T::Err; + + fn from_str(s: &str) -> Result { + let items = s + .split(',') + .map(str::trim) + .filter(|part| !part.is_empty()) + .map(T::from_str) + .collect::, _>>()?; + Ok(Self(items)) + } +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(default)] +pub struct SecurityConfig { + pub mode: Option, + pub token_audiences: Vec, + pub allowed_service_accounts: Vec, + pub cache_ttl_secs: u64, +} + +impl Default for SecurityConfig { + fn default() -> Self { + Self { + mode: None, + token_audiences: Vec::new(), + allowed_service_accounts: Vec::new(), + cache_ttl_secs: 60, + } + } +} + +impl SecurityConfig { + #[must_use] + pub fn resolve_mode(&self) -> AuthMode { + self.mode.unwrap_or(AuthMode::Off) + } + + pub fn validate_resolved(&self, mode: AuthMode) -> Result<(), String> { + if mode == AuthMode::Enforce { + if self.token_audiences.is_empty() { + return Err( + "security mode 'enforce' requires at least one token audience \ + (security.token_audiences)" + .to_string(), + ); + } + if self.allowed_service_accounts.is_empty() { + return Err( + "security mode 'enforce' requires a non-empty service account \ + allowlist (security.allowed_service_accounts)" + .to_string(), + ); + } + } + Ok(()) + } +} + +#[derive(clap::Args, Debug, Default)] +pub struct SecurityArgs { + /// ServiceAccount auth mode (off, enforce). Off by default. + #[arg( + long = "security-mode", + env = modelexpress_common::envs::MODEL_EXPRESS_SECURITY_MODE, + value_enum + )] + pub mode: Option, + + /// Comma-separated SA token audiences the caller's token must carry. + #[arg( + long = "security-token-audiences", + env = modelexpress_common::envs::MODEL_EXPRESS_SECURITY_TOKEN_AUDIENCES + )] + pub token_audiences: Option>, + + /// Comma-separated allowed callers as `:`. + #[arg( + long = "security-allowed-service-accounts", + env = modelexpress_common::envs::MODEL_EXPRESS_SECURITY_ALLOWED_SERVICE_ACCOUNTS + )] + pub allowed_service_accounts: Option>, + + /// TTL for the verified-token and rejection caches, in seconds. + #[arg( + long = "security-cache-ttl-secs", + env = modelexpress_common::envs::MODEL_EXPRESS_SECURITY_CACHE_TTL_SECS + )] + pub cache_ttl_secs: Option, +} + /// Command line arguments for the server #[derive(Parser, Debug)] #[command(author, version, about, long_about = None)] @@ -44,6 +190,9 @@ pub struct ServerArgs { #[arg(long, env = modelexpress_common::envs::MODEL_EXPRESS_CACHE_EVICTION_ENABLED)] pub cache_eviction_enabled: Option, + #[command(flatten)] + pub security: SecurityArgs, + /// Validate configuration and exit #[arg(long)] pub validate_config: bool, @@ -58,6 +207,8 @@ pub struct ServerConfig { pub cache: CacheConfig, /// Logging configuration pub logging: LoggingConfig, + #[serde(default)] + pub security: SecurityConfig, } /// Server-specific settings @@ -177,6 +328,20 @@ impl ServerConfig { config.cache.eviction.enabled = cache_eviction_enabled; } + // Apply security overrides + if let Some(mode) = args.security.mode { + config.security.mode = Some(mode); + } + if let Some(token_audiences) = args.security.token_audiences { + config.security.token_audiences = token_audiences.into_inner(); + } + if let Some(allowed) = args.security.allowed_service_accounts { + config.security.allowed_service_accounts = allowed.into_inner(); + } + if let Some(cache_ttl_secs) = args.security.cache_ttl_secs { + config.security.cache_ttl_secs = cache_ttl_secs; + } + // Validate the final configuration config.validate()?; @@ -195,6 +360,11 @@ impl ServerConfig { ))); } + let mode = self.security.resolve_mode(); + self.security + .validate_resolved(mode) + .map_err(ConfigError::Message)?; + Ok(()) } @@ -232,6 +402,20 @@ impl ServerConfig { info!(" Format: {}", self.logging.format); info!(" File: {:?}", self.logging.file); info!(" Structured: {}", self.logging.structured); + + info!("Security Configuration:"); + info!(" Mode: {:?}", self.security.resolve_mode()); + info!(" Token audiences: {:?}", self.security.token_audiences); + info!( + " Allowed service accounts: {}", + self.security + .allowed_service_accounts + .iter() + .map(ToString::to_string) + .collect::>() + .join(", ") + ); + info!(" Cache TTL: {}s", self.security.cache_ttl_secs); } } @@ -496,6 +680,7 @@ mod tests { log_format: None, cache_directory: None, cache_eviction_enabled: None, + security: SecurityArgs::default(), validate_config: false, }; @@ -535,6 +720,7 @@ mod tests { log_format: None, cache_directory: None, cache_eviction_enabled: None, + security: SecurityArgs::default(), validate_config: false, }; @@ -580,6 +766,7 @@ mod tests { log_format: Some(LogFormat::Json), cache_directory: Some(PathBuf::from("/tmp/override_cache")), cache_eviction_enabled: Some(false), + security: SecurityArgs::default(), validate_config: false, }; @@ -607,6 +794,7 @@ mod tests { log_format: None, cache_directory: None, cache_eviction_enabled: None, + security: SecurityArgs::default(), validate_config: false, }; @@ -619,4 +807,62 @@ mod tests { assert_eq!(config.server.port.get(), 9001); assert_eq!(config.logging.level, LogLevel::Warn); } + + #[test] + fn security_defaults_to_off() { + let security = SecurityConfig::default(); + assert_eq!(security.resolve_mode(), AuthMode::Off); + assert!(security.token_audiences.is_empty()); + assert!(security.allowed_service_accounts.is_empty()); + } + + #[test] + #[allow(clippy::expect_used)] + fn service_account_ref_parses_namespace_and_sa() { + let parsed: ServiceAccountRef = "vllm:worker".parse().expect("valid ns:sa"); + assert_eq!(parsed.namespace, "vllm"); + assert_eq!(parsed.service_account, "worker"); + assert_eq!(parsed.to_string(), "vllm:worker"); + let trimmed: ServiceAccountRef = " ns : sa ".parse().expect("trims"); + assert_eq!(trimmed, "ns:sa".parse().expect("valid")); + } + + #[test] + fn service_account_ref_rejects_malformed() { + assert!("noseparator".parse::().is_err()); + assert!(":sa".parse::().is_err()); + assert!("ns:".parse::().is_err()); + } + + #[test] + #[allow(clippy::expect_used)] + fn comma_list_parses_and_skips_blanks() { + let list: CommaList = + "a:one, b:two ,, c:three".parse().expect("valid list"); + let items = list.into_inner(); + assert_eq!(items.len(), 3); + assert_eq!(items[1].namespace, "b"); + assert_eq!(items[2].service_account, "three"); + + let empty: CommaList = " ".parse().expect("empty"); + assert!(empty.into_inner().is_empty()); + } + + #[test] + fn enforce_requires_audience_and_allowlist() { + let mut security = SecurityConfig { + mode: Some(AuthMode::Enforce), + ..SecurityConfig::default() + }; + assert!(security.validate_resolved(AuthMode::Enforce).is_err()); + + security.token_audiences = vec!["modelexpress".to_string()]; + assert!(security.validate_resolved(AuthMode::Enforce).is_err()); + + security.allowed_service_accounts = vec![ServiceAccountRef { + namespace: "vllm".to_string(), + service_account: "worker".to_string(), + }]; + assert!(security.validate_resolved(AuthMode::Enforce).is_ok()); + } } diff --git a/modelexpress_server/src/lib.rs b/modelexpress_server/src/lib.rs index 00c63f51..83ca1905 100644 --- a/modelexpress_server/src/lib.rs +++ b/modelexpress_server/src/lib.rs @@ -1,6 +1,7 @@ // SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +pub mod auth; pub mod backend_config; pub mod cache; pub mod config; diff --git a/modelexpress_server/src/server.rs b/modelexpress_server/src/server.rs index 2cef51a1..a16716f8 100644 --- a/modelexpress_server/src/server.rs +++ b/modelexpress_server/src/server.rs @@ -13,11 +13,13 @@ use modelexpress_common::grpc::{ model::model_service_server::ModelServiceServer, p2p::p2p_service_server::P2pServiceServer, }; use tonic::transport::Server; +use tower::Layer; use tracing::{error, info}; +use crate::auth::{AuthLayer, AuthState}; use crate::backend_config::BackendConfig; use crate::cache::CacheEvictionService; -use crate::config::ServerConfig; +use crate::config::{AuthMode, ServerConfig}; use crate::p2p::{service::P2pServiceImpl, state::P2pStateManager}; use crate::registry::state::RegistryManager; use crate::services::{ApiServiceImpl, HealthServiceImpl, ModelDownloadTracker, ModelServiceImpl}; @@ -151,20 +153,48 @@ pub async fn run_server( } }; - // Start the gRPC server + let mode = config.security.resolve_mode(); + let auth_layer = if mode == AuthMode::Off { + info!("ServiceAccount auth: disabled"); + None + } else { + config + .security + .validate_resolved(mode) + .map_err(|e| -> Box { e.into() })?; + let client = kube::Client::try_default().await.map_err(|e| { + error!("ServiceAccount auth enabled but kube client init failed: {e}"); + e + })?; + info!( + "ServiceAccount auth: enforce ({} allowed service account(s), {} audience(s))", + config.security.allowed_service_accounts.len(), + config.security.token_audiences.len() + ); + Some(AuthLayer::new(Arc::new(AuthState::new( + client, + &config.security, + )))) + }; + + let api = ApiServiceServer::new(api_service); + let model = ModelServiceServer::new(model_service); + let p2p = P2pServiceServer::new(p2p_service) + .max_decoding_message_size(MAX_MESSAGE_SIZE) + .max_encoding_message_size(MAX_MESSAGE_SIZE); + info!("Starting gRPC server on: {addr}"); - let server_result = Server::builder() + let router = Server::builder() .add_service(health_service_v1) - .add_service(HealthServiceServer::new(health_service)) - .add_service(ApiServiceServer::new(api_service)) - .add_service(ModelServiceServer::new(model_service)) - .add_service( - P2pServiceServer::new(p2p_service) - .max_decoding_message_size(MAX_MESSAGE_SIZE) - .max_encoding_message_size(MAX_MESSAGE_SIZE), - ) - .serve_with_shutdown(addr, shutdown_signal) - .await; + .add_service(HealthServiceServer::new(health_service)); + let router = match &auth_layer { + Some(layer) => router + .add_service(layer.layer(api)) + .add_service(layer.layer(model)) + .add_service(layer.layer(p2p)), + None => router.add_service(api).add_service(model).add_service(p2p), + }; + let server_result = router.serve_with_shutdown(addr, shutdown_signal).await; // Wait for background services to complete if let Some(handle) = cache_handle diff --git a/modelexpress_server/tests/auth_e2e.rs b/modelexpress_server/tests/auth_e2e.rs new file mode 100644 index 00000000..5c9f0e44 --- /dev/null +++ b/modelexpress_server/tests/auth_e2e.rs @@ -0,0 +1,200 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Full-stack ServiceAccount token auth: the real client (interceptor reading a projected +//! token file) against a real tonic server wrapped in [`AuthLayer`], with a fake +//! TokenReview backend. Scenarios run sequentially in one test because the token path +//! env var is process-global. + +#![allow(clippy::expect_used)] + +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::time::Duration; + +use http::header::CONTENT_TYPE; +use k8s_openapi::api::authentication::v1::{TokenReview, TokenReviewStatus, UserInfo}; +use kube::client::{Body, Client as KubeClient}; +use modelexpress_client::Client; +use modelexpress_common::Error; +use modelexpress_common::client_config::ClientConfig; +use modelexpress_common::config::ConnectionConfig; +use modelexpress_common::grpc::health::health_service_server::HealthServiceServer; +use modelexpress_server::auth::{AuthLayer, AuthState}; +use modelexpress_server::config::{AuthMode, SecurityConfig, ServiceAccountRef}; +use modelexpress_server::services::HealthServiceImpl; +use tokio::sync::oneshot; +use tokio::task::JoinHandle; +use tonic::transport::Server; +use tower::{BoxError, Layer}; + +type ServerResult = Result<(), Box>; + +fn free_port() -> u16 { + let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("bind ephemeral port"); + listener.local_addr().expect("local addr").port() +} + +fn fake_kube_client_for_review(review: TokenReview) -> (KubeClient, Arc) { + let calls = Arc::new(AtomicUsize::new(0)); + let service = tower::service_fn({ + let calls = calls.clone(); + move |_request: http::Request| { + calls.fetch_add(1, Ordering::SeqCst); + let body = serde_json::to_vec(&review).expect("serialize TokenReview"); + let response = http::Response::builder() + .status(http::StatusCode::OK) + .header(CONTENT_TYPE, "application/json") + .body(Body::from(body)) + .expect("fake TokenReview response body"); + async move { Ok::<_, BoxError>(response) } + } + }); + (KubeClient::new(service, "default"), calls) +} + +fn start_auth_server( + port: u16, + kube_client: KubeClient, + security_config: SecurityConfig, +) -> (oneshot::Sender<()>, JoinHandle) { + let (tx, rx) = oneshot::channel::<()>(); + let shutdown = async move { + let _ = rx.await; + }; + let handle = tokio::spawn(async move { + let addr = format!("127.0.0.1:{port}").parse()?; + let state = Arc::new(AuthState::new(kube_client, &security_config)); + let auth_layer = AuthLayer::new(state); + Server::builder() + .add_service(auth_layer.layer(HealthServiceServer::new(HealthServiceImpl))) + .serve_with_shutdown(addr, shutdown) + .await + .map_err(Into::into) + }); + (tx, handle) +} + +async fn connect_client(port: u16) -> Client { + let config = ClientConfig { + connection: ConnectionConfig::new(format!("http://127.0.0.1:{port}")), + ..Default::default() + }; + for _ in 0..100 { + if let Ok(client) = Client::new(config.clone()).await { + return client; + } + tokio::time::sleep(Duration::from_millis(50)).await; + } + panic!("server on port {port} never became reachable"); +} + +async fn stop_and_join(shutdown: oneshot::Sender<()>, handle: JoinHandle) { + let _ = shutdown.send(()); + tokio::time::timeout(Duration::from_secs(10), handle) + .await + .expect("server task did not exit in time") + .expect("server task panicked") + .expect("server returned an error"); +} + +fn token_review_for(username: &str) -> TokenReview { + TokenReview { + status: Some(TokenReviewStatus { + authenticated: Some(true), + audiences: Some(vec!["modelexpress".to_string()]), + user: Some(UserInfo { + username: Some(username.to_string()), + ..Default::default() + }), + ..Default::default() + }), + ..Default::default() + } +} + +fn unauthenticated_review() -> TokenReview { + TokenReview { + status: Some(TokenReviewStatus { + authenticated: Some(false), + ..Default::default() + }), + ..Default::default() + } +} + +fn security_config() -> SecurityConfig { + SecurityConfig { + mode: Some(AuthMode::Enforce), + token_audiences: vec!["modelexpress".to_string()], + allowed_service_accounts: vec![ServiceAccountRef { + namespace: "vllm".to_string(), + service_account: "worker".to_string(), + }], + cache_ttl_secs: 60, + } +} + +fn grpc_code(error: &Error) -> Option { + match error { + Error::Grpc(status) => Some(status.code()), + _ => None, + } +} + +#[tokio::test] +async fn token_auth_end_to_end() { + let temp_dir = tempfile::tempdir().expect("temp dir"); + let token_path = temp_dir.path().join("token"); + std::fs::write(&token_path, "worker-token\n").expect("write token"); + unsafe { + std::env::set_var( + modelexpress_common::envs::MX_AUTH_TOKEN_PATH, + token_path.to_str().expect("token path"), + ); + } + + // Allowlisted caller: RPCs succeed and the verified token is served from cache. + let port = free_port(); + let (kube_client, calls) = + fake_kube_client_for_review(token_review_for("system:serviceaccount:vllm:worker")); + let (shutdown, handle) = start_auth_server(port, kube_client, security_config()); + let mut client = connect_client(port).await; + client.health_check().await.expect("first health_check"); + client.health_check().await.expect("second health_check"); + assert_eq!(calls.load(Ordering::SeqCst), 1); + stop_and_join(shutdown, handle).await; + + // Authenticated but non-allowlisted caller: PERMISSION_DENIED. + let port = free_port(); + let (kube_client, _calls) = + fake_kube_client_for_review(token_review_for("system:serviceaccount:other:intruder")); + let (shutdown, handle) = start_auth_server(port, kube_client, security_config()); + let mut client = connect_client(port).await; + let error = client + .health_check() + .await + .expect_err("non-allowlisted SA should be denied"); + assert_eq!( + grpc_code(&error), + Some(tonic::Code::PermissionDenied), + "expected PERMISSION_DENIED, got: {error}" + ); + stop_and_join(shutdown, handle).await; + + // Token the reviewer rejects: UNAUTHENTICATED. + let port = free_port(); + let (kube_client, _calls) = fake_kube_client_for_review(unauthenticated_review()); + let (shutdown, handle) = start_auth_server(port, kube_client, security_config()); + let mut client = connect_client(port).await; + let error = client + .health_check() + .await + .expect_err("unauthenticated caller should be rejected"); + assert_eq!( + grpc_code(&error), + Some(tonic::Code::Unauthenticated), + "expected UNAUTHENTICATED, got: {error}" + ); + stop_and_join(shutdown, handle).await; +}