From 4edee6954a9a0390598678d782f1860d9c105d27 Mon Sep 17 00:00:00 2001 From: Simo Lin Date: Fri, 26 Dec 2025 21:30:00 -0500 Subject: [PATCH] [model-gateway] add JWT/OIDC authentication for control plane APIs (#15850) --- sgl-model-gateway/Cargo.toml | 2 + sgl-model-gateway/README.md | 116 +++ .../bindings/python/sglang_router/router.py | 68 ++ .../python/sglang_router/router_args.py | 109 ++ sgl-model-gateway/bindings/python/src/lib.rs | 158 +++ sgl-model-gateway/src/auth/audit.rs | 357 +++++++ sgl-model-gateway/src/auth/config.rs | 345 +++++++ sgl-model-gateway/src/auth/jwks.rs | 515 ++++++++++ sgl-model-gateway/src/auth/jwt.rs | 532 ++++++++++ sgl-model-gateway/src/auth/middleware.rs | 399 ++++++++ sgl-model-gateway/src/auth/mod.rs | 20 + sgl-model-gateway/src/lib.rs | 1 + sgl-model-gateway/src/main.rs | 145 +++ sgl-model-gateway/src/server.rs | 39 +- .../tests/auth_integration_test.rs | 939 ++++++++++++++++++ sgl-model-gateway/tests/common/test_app.rs | 2 + sgl-model-gateway/tests/wasm_test.rs | 1 + 17 files changed, 3738 insertions(+), 10 deletions(-) create mode 100644 sgl-model-gateway/src/auth/audit.rs create mode 100644 sgl-model-gateway/src/auth/config.rs create mode 100644 sgl-model-gateway/src/auth/jwks.rs create mode 100644 sgl-model-gateway/src/auth/jwt.rs create mode 100644 sgl-model-gateway/src/auth/middleware.rs create mode 100644 sgl-model-gateway/src/auth/mod.rs create mode 100644 sgl-model-gateway/tests/auth_integration_test.rs diff --git a/sgl-model-gateway/Cargo.toml b/sgl-model-gateway/Cargo.toml index d5d2d2f60..eaa68340a 100644 --- a/sgl-model-gateway/Cargo.toml +++ b/sgl-model-gateway/Cargo.toml @@ -97,6 +97,7 @@ rmcp = { version = "0.8.3", features = ["client", "server", serde_yaml = "0.9" oracle = { version = "0.6.3", features = ["chrono"] } subtle = "2.6" +jsonwebtoken = { version = "9.3", default-features = false, features = ["use_pem"] } rustpython-parser = "0.4.0" num-traits = "0.2" image = { version = "0.25.4", default-features = false, features = ["png", "jpeg", "gif", "bmp", "ico", "tiff", "webp"] } @@ -144,6 +145,7 @@ npyz = { version = "0.8", features = ["npz"] } # For reading numpy .npz files i opentelemetry-proto = { version = "0.27", features = ["gen-tonic"] } tonic-v12 = { version = "0.12.3", package = "tonic" } serial_test = "3.0" +rsa = { version = "0.9", features = ["sha2"] } [[bench]] diff --git a/sgl-model-gateway/README.md b/sgl-model-gateway/README.md index f911749d6..60be371d1 100644 --- a/sgl-model-gateway/README.md +++ b/sgl-model-gateway/README.md @@ -825,6 +825,122 @@ python3 -m sglang_router.launch_router \ --api-key "secure-api-key" ``` +### Control Plane Authentication + +The gateway supports role-based access control (RBAC) for control plane APIs (worker management, tokenizer registration, cache operations). Two authentication methods are available: + +#### Authentication Methods + +| Method | Use Case | Configuration | +|--------|----------|---------------| +| **API Keys** | Service accounts, internal services | `--control-plane-api-keys` | +| **JWT/OIDC** | User authentication via Identity Provider | `--jwt-issuer`, `--jwt-audience` | + +Both methods can be used together. Requests are authenticated in order: API key → JWT token. + +#### Roles + +| Role | Access | +|------|--------| +| `admin` | Full access to all control plane APIs (workers, tokenizers, cache, etc.) | +| `user` | Inference/data plane APIs only (chat completions, embeddings, etc.) | + +#### API Key Authentication + +Static API keys for service accounts and automation: + +```bash +python3 -m sglang_router.launch_router \ + --worker-urls http://worker1:8000 \ + --control-plane-api-keys 'svc1:CI Pipeline:admin:secret-key-123' \ + 'svc2:Monitoring:user:readonly-key-456' \ + --control-plane-audit-enabled +``` + +**Format:** `id:name:role:key` +- `id` - Unique identifier for the key +- `name` - Human-readable description +- `role` - Either `admin` or `user` +- `key` - The secret key (stored as SHA-256 hash internally) + +**Usage:** +```bash +curl -H "Authorization: Bearer secret-key-123" \ + http://localhost:30000/workers +``` + +#### JWT/OIDC Authentication + +Authenticate users via an external Identity Provider (Azure AD, Okta, Auth0, Keycloak, etc.): + +```bash +python3 -m sglang_router.launch_router \ + --worker-urls http://worker1:8000 \ + --jwt-issuer "https://login.microsoftonline.com/{tenant-id}/v2.0" \ + --jwt-audience "api://my-gateway-client-id" \ + --jwt-jwks-uri "https://login.microsoftonline.com/{tenant-id}/discovery/v2.0/keys" \ + --jwt-role-mapping 'Gateway.Admins=admin' 'Gateway.Users=user' \ + --control-plane-audit-enabled +``` + +| Parameter | Description | +|-----------|-------------| +| `--jwt-issuer` | OIDC issuer URL. Used to validate the `iss` claim and discover JWKS endpoint via `.well-known/openid-configuration`. | +| `--jwt-audience` | Expected audience (`aud` claim). Typically your application's client ID or API identifier (e.g., `api://client-id`). | +| `--jwt-jwks-uri` | (Optional) Explicit JWKS URI. If omitted, discovered automatically from the issuer's OIDC configuration. | +| `--jwt-role-mapping` | Map IDP group/role names to gateway roles. Format: `idp_role=gateway_role`. | + +**How it works:** +1. User authenticates with Identity Provider (OAuth2/OIDC flow) +2. IDP issues a JWT token +3. User sends token to gateway: `Authorization: Bearer ` +4. Gateway validates the JWT: + - Verifies signature against JWKS + - Checks `iss` matches `--jwt-issuer` + - Checks `aud` matches `--jwt-audience` + - Validates expiration and other standard claims + - Extracts role from `roles` claim (or `groups` as fallback) + - Maps IDP role to gateway role via `--jwt-role-mapping` + +**Example Azure AD Configuration:** +```bash +# Azure AD issues tokens with: +# iss: https://login.microsoftonline.com/{tenant}/v2.0 +# aud: api://your-client-id (or the client ID itself) +# roles: ["Gateway.Admins"] or groups: ["group-id"] + +python3 -m sglang_router.launch_router \ + --jwt-issuer "https://login.microsoftonline.com/your-tenant-id/v2.0" \ + --jwt-audience "api://your-client-id" \ + --jwt-role-mapping 'Gateway.Admins=admin' 'Gateway.Users=user' +``` + +#### Audit Logging + +Enable `--control-plane-audit-enabled` to log all control plane operations with: +- Timestamp +- Principal (API key ID or JWT subject) +- Role +- Action performed +- Success/failure status + +#### Combined Authentication Example + +Use both API keys and JWT for different use cases: + +```bash +python3 -m sglang_router.launch_router \ + --worker-urls http://worker1:8000 \ + # API keys for service accounts + --control-plane-api-keys 'ci:CI/CD Pipeline:admin:ci-secret' \ + # JWT for human users via Azure AD + --jwt-issuer "https://login.microsoftonline.com/{tenant}/v2.0" \ + --jwt-audience "api://gateway" \ + --jwt-role-mapping 'Platform.Admins=admin' 'Platform.Users=user' \ + # Enable audit logging + --control-plane-audit-enabled +``` + ## Development & Testing ```bash # Build Rust components (debug mode, fast) diff --git a/sgl-model-gateway/bindings/python/sglang_router/router.py b/sgl-model-gateway/bindings/python/sglang_router/router.py index ae0bc04aa..e572134ae 100644 --- a/sgl-model-gateway/bindings/python/sglang_router/router.py +++ b/sgl-model-gateway/bindings/python/sglang_router/router.py @@ -5,8 +5,12 @@ from sglang_router.sglang_router_rs import ( BackendType, HistoryBackendType, PolicyType, + PyApiKeyEntry, + PyControlPlaneAuthConfig, + PyJwtConfig, PyOracleConfig, PyPostgresConfig, + PyRole, ) from sglang_router.sglang_router_rs import Router as _Router @@ -61,6 +65,60 @@ def history_backend_from_str(backend_str: Optional[str]) -> HistoryBackendType: raise ValueError(f"Unknown history backend: {backend_str}") +def role_from_str(role_str: str) -> PyRole: + """Convert role string to PyRole enum.""" + if role_str.lower() == "admin": + return PyRole.Admin + return PyRole.User + + +def build_control_plane_auth_config( + args_dict: dict, +) -> Optional[PyControlPlaneAuthConfig]: + """Build control plane auth config from args dict.""" + api_keys = args_dict.get("control_plane_api_keys", []) + jwt_issuer = args_dict.get("jwt_issuer") + jwt_audience = args_dict.get("jwt_audience") + audit_enabled = args_dict.get("control_plane_audit_enabled", False) + + # Check if any auth is configured + has_api_keys = bool(api_keys) + has_jwt = jwt_issuer is not None and jwt_audience is not None + + if not has_api_keys and not has_jwt: + return None + + # Build API key entries + py_api_keys = [] + for key_tuple in api_keys: + # Tuple format: (id, name, key, role) + key_id, name, key, role = key_tuple + py_api_keys.append( + PyApiKeyEntry( + id=key_id, + name=name, + key=key, + role=role_from_str(role), + ) + ) + + # Build JWT config if present + jwt_config = None + if has_jwt: + jwt_config = PyJwtConfig( + issuer=jwt_issuer, + audience=jwt_audience, + jwks_uri=args_dict.get("jwt_jwks_uri"), + role_mapping=args_dict.get("jwt_role_mapping", {}), + ) + + return PyControlPlaneAuthConfig( + jwt=jwt_config, + api_keys=py_api_keys, + audit_enabled=audit_enabled, + ) + + class Router: """ A high-performance router for distributing requests across worker nodes. @@ -202,6 +260,9 @@ class Router: ) args_dict["postgres_config"] = postgres_config + # Build control plane auth config + args_dict["control_plane_auth"] = build_control_plane_auth_config(args_dict) + # Remove fields that shouldn't be passed to Rust Router constructor fields_to_remove = [ "mini_lb", @@ -215,6 +276,13 @@ class Router: "oracle_pool_timeout_secs", "postgres_db_url", "postgres_pool_max", + # Control plane auth fields (converted to control_plane_auth) + "control_plane_api_keys", + "control_plane_audit_enabled", + "jwt_issuer", + "jwt_audience", + "jwt_jwks_uri", + "jwt_role_mapping", ] for field in fields_to_remove: args_dict.pop(field, None) diff --git a/sgl-model-gateway/bindings/python/sglang_router/router_args.py b/sgl-model-gateway/bindings/python/sglang_router/router_args.py index 03d206546..171ce8205 100644 --- a/sgl-model-gateway/bindings/python/sglang_router/router_args.py +++ b/sgl-model-gateway/bindings/python/sglang_router/router_args.py @@ -126,6 +126,15 @@ class RouterArgs: # Trace enable_trace: bool = False otlp_traces_endpoint: str = "localhost:4317" + # Control plane authentication + # API keys for control plane auth (list of tuples: id, name, key, role) + control_plane_api_keys: List[tuple] = dataclasses.field(default_factory=list) + control_plane_audit_enabled: bool = False + # JWT/OIDC configuration for control plane auth + jwt_issuer: Optional[str] = None + jwt_audience: Optional[str] = None + jwt_jwks_uri: Optional[str] = None + jwt_role_mapping: Dict[str, str] = dataclasses.field(default_factory=dict) @staticmethod def add_cli_args( @@ -686,6 +695,47 @@ class RouterArgs: default="localhost:4317", help="Config opentelemetry collector endpoint if --enable-trace is set. format: :", ) + # Control plane authentication + parser.add_argument( + f"--{prefix}control-plane-api-keys", + type=str, + nargs="*", + default=[], + help="API keys for control plane authentication. Format: 'id:name:role:key' where role is 'admin' or 'user'. " + "Example: --control-plane-api-keys 'key1:Service Account:admin:secret123' 'key2:Read Only:user:secret456'", + ) + parser.add_argument( + f"--{prefix}control-plane-audit-enabled", + action="store_true", + default=False, + help="Enable audit logging for control plane operations", + ) + parser.add_argument( + f"--{prefix}jwt-issuer", + type=str, + default=None, + help="OIDC issuer URL for JWT authentication (e.g., https://login.microsoftonline.com/{tenant}/v2.0)", + ) + parser.add_argument( + f"--{prefix}jwt-audience", + type=str, + default=None, + help="Expected audience claim for JWT tokens (usually the client ID or API identifier)", + ) + parser.add_argument( + f"--{prefix}jwt-jwks-uri", + type=str, + default=None, + help="Explicit JWKS URI. If not provided, discovered from issuer via .well-known/openid-configuration", + ) + parser.add_argument( + f"--{prefix}jwt-role-mapping", + type=str, + nargs="*", + default=[], + help="Mapping from IDP role/group names to gateway roles. Format: 'idp_role=gateway_role'. " + "Example: --jwt-role-mapping 'Gateway.Admin=admin' 'Gateway.User=user'", + ) @classmethod def from_cli_args( @@ -741,6 +791,16 @@ class RouterArgs: # Mooncake-specific annotation args_dict["bootstrap_port_annotation"] = "sglang.ai/bootstrap-port" + # Parse control plane API keys + args_dict["control_plane_api_keys"] = cls._parse_control_plane_api_keys( + cli_args_dict.get(f"{prefix}control_plane_api_keys", []) + ) + + # Parse JWT role mapping + args_dict["jwt_role_mapping"] = cls._parse_jwt_role_mapping( + cli_args_dict.get(f"{prefix}jwt_role_mapping", []) + ) + return cls(**args_dict) def _validate_router_args(self): @@ -830,3 +890,52 @@ class RouterArgs: # decode_list is a list of single-element lists due to nargs=1 return [url[0] for url in decode_list] + + @staticmethod + def _parse_control_plane_api_keys(api_keys_list): + """Parse control plane API keys from --control-plane-api-keys arguments. + + Format: id:name:role:key + Example: --control-plane-api-keys 'key1:Service Account:admin:secret123' + """ + if not api_keys_list: + return [] + + parsed_keys = [] + for key_str in api_keys_list: + parts = key_str.split(":", 3) # Split into at most 4 parts + if len(parts) != 4: + raise ValueError( + f"Invalid API key format: '{key_str}'. Expected 'id:name:role:key'" + ) + key_id, name, role, key = parts + role_lower = role.lower() + if role_lower not in ("admin", "user"): + raise ValueError(f"Invalid role: '{role}'. Must be 'admin' or 'user'") + parsed_keys.append((key_id, name, key, role_lower)) + return parsed_keys + + @staticmethod + def _parse_jwt_role_mapping(role_mapping_list): + """Parse JWT role mapping from --jwt-role-mapping arguments. + + Format: idp_role=gateway_role + Example: --jwt-role-mapping 'Gateway.Admin=admin' 'Gateway.User=user' + """ + if not role_mapping_list: + return {} + + mapping = {} + for mapping_str in role_mapping_list: + if "=" not in mapping_str: + raise ValueError( + f"Invalid role mapping format: '{mapping_str}'. Expected 'idp_role=gateway_role'" + ) + idp_role, gateway_role = mapping_str.split("=", 1) + gateway_role_lower = gateway_role.lower() + if gateway_role_lower not in ("admin", "user"): + raise ValueError( + f"Invalid gateway role: '{gateway_role}'. Must be 'admin' or 'user'" + ) + mapping[idp_role] = gateway_role_lower + return mapping diff --git a/sgl-model-gateway/bindings/python/src/lib.rs b/sgl-model-gateway/bindings/python/src/lib.rs index 0a159b46e..b27a5a29e 100644 --- a/sgl-model-gateway/bindings/python/src/lib.rs +++ b/sgl-model-gateway/bindings/python/src/lib.rs @@ -31,6 +31,152 @@ pub enum HistoryBackendType { Postgres, } +#[pyclass(eq)] +#[derive(Clone, PartialEq, Debug, Default)] +pub enum PyRole { + Admin, + #[default] + User, +} + +impl PyRole { + pub fn to_auth_role(&self) -> auth::Role { + match self { + PyRole::Admin => auth::Role::Admin, + PyRole::User => auth::Role::User, + } + } +} + +#[pyclass] +#[derive(Clone, Debug, PartialEq)] +pub struct PyApiKeyEntry { + #[pyo3(get, set)] + pub id: String, + #[pyo3(get, set)] + pub name: String, + #[pyo3(get, set)] + pub key: String, + #[pyo3(get, set)] + pub role: PyRole, +} + +#[pymethods] +impl PyApiKeyEntry { + #[new] + #[pyo3(signature = (id, name, key, role = PyRole::User))] + fn new(id: String, name: String, key: String, role: PyRole) -> Self { + PyApiKeyEntry { id, name, key, role } + } +} + +impl PyApiKeyEntry { + pub fn to_auth_api_key_entry(&self) -> auth::ApiKeyEntry { + auth::ApiKeyEntry::new(&self.id, &self.name, &self.key, self.role.to_auth_role()) + } +} + +#[pyclass] +#[derive(Clone, Debug, PartialEq)] +pub struct PyJwtConfig { + #[pyo3(get, set)] + pub issuer: String, + #[pyo3(get, set)] + pub audience: String, + #[pyo3(get, set)] + pub jwks_uri: Option, + #[pyo3(get, set)] + pub role_mapping: HashMap, +} + +#[pymethods] +impl PyJwtConfig { + #[new] + #[pyo3(signature = ( + issuer, + audience, + jwks_uri = None, + role_mapping = HashMap::new(), + ))] + fn new( + issuer: String, + audience: String, + jwks_uri: Option, + role_mapping: HashMap, + ) -> Self { + PyJwtConfig { + issuer, + audience, + jwks_uri, + role_mapping, + } + } +} + +impl PyJwtConfig { + pub fn to_auth_jwt_config(&self) -> auth::JwtConfig { + let mut config = auth::JwtConfig::new(&self.issuer, &self.audience); + + // Conditionally set JWKS URI + if let Some(ref uri) = self.jwks_uri { + config = config.with_jwks_uri(uri); + } + + // Add role mappings + for (idp_role, gateway_role) in &self.role_mapping { + let role = match gateway_role.to_lowercase().as_str() { + "admin" => auth::Role::Admin, + _ => auth::Role::User, + }; + config = config.with_role_mapping(idp_role, role); + } + + config + } +} + +#[pyclass] +#[derive(Clone, Debug, Default, PartialEq)] +pub struct PyControlPlaneAuthConfig { + #[pyo3(get, set)] + pub jwt: Option, + #[pyo3(get, set)] + pub api_keys: Vec, + #[pyo3(get, set)] + pub audit_enabled: bool, +} + +#[pymethods] +impl PyControlPlaneAuthConfig { + #[new] + #[pyo3(signature = ( + jwt = None, + api_keys = vec![], + audit_enabled = true, + ))] + fn new( + jwt: Option, + api_keys: Vec, + audit_enabled: bool, + ) -> Self { + PyControlPlaneAuthConfig { + jwt, + api_keys, + audit_enabled, + } + } +} + +impl PyControlPlaneAuthConfig { + pub fn to_auth_control_plane_config(&self) -> auth::ControlPlaneAuthConfig { + auth::ControlPlaneAuthConfig { + jwt: self.jwt.as_ref().map(|j| j.to_auth_jwt_config()), + api_keys: self.api_keys.iter().map(|k| k.to_auth_api_key_entry()).collect(), + audit_enabled: self.audit_enabled, + } + } +} + #[pyclass] #[derive(Clone, PartialEq)] pub struct PyOracleConfig { @@ -232,6 +378,7 @@ struct Router { server_key_path: Option, enable_trace: bool, otlp_traces_endpoint: String, + control_plane_auth: Option, } impl Router { @@ -502,6 +649,7 @@ impl Router { server_key_path = None, enable_trace = false, otlp_traces_endpoint = String::from("localhost:4317"), + control_plane_auth = None, ))] #[allow(clippy::too_many_arguments)] fn new( @@ -583,6 +731,7 @@ impl Router { server_key_path: Option, enable_trace: bool, otlp_traces_endpoint: String, + control_plane_auth: Option, ) -> PyResult { let mut all_urls = worker_urls.clone(); @@ -678,6 +827,7 @@ impl Router { server_key_path, enable_trace, otlp_traces_endpoint, + control_plane_auth, }) } @@ -736,6 +886,10 @@ impl Router { request_timeout_secs: self.request_timeout_secs, request_id_headers: self.request_id_headers.clone(), shutdown_grace_period_secs: self.shutdown_grace_period_secs, + control_plane_auth: self + .control_plane_auth + .as_ref() + .map(|c| c.to_auth_control_plane_config()), }) .await .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string())) @@ -772,6 +926,10 @@ fn sglang_router_rs(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/sgl-model-gateway/src/auth/audit.rs b/sgl-model-gateway/src/auth/audit.rs new file mode 100644 index 000000000..77aebe3ca --- /dev/null +++ b/sgl-model-gateway/src/auth/audit.rs @@ -0,0 +1,357 @@ +//! Audit logging for control plane operations. +//! +//! Provides structured audit events for security monitoring and compliance. +//! +//! Security features: +//! - Input sanitization to prevent log injection attacks +//! - Structured logging for safe parsing + +use chrono::{DateTime, Utc}; +use serde::Serialize; +use tracing::{info, span, Level}; + +use super::config::Role; + +/// Maximum length for sanitized strings to prevent log flooding +const MAX_SANITIZED_LENGTH: usize = 1024; + +/// Sanitize a string for safe logging. +/// Removes/escapes control characters and newlines that could be used for log injection. +fn sanitize_for_log(input: &str) -> String { + let mut result = String::with_capacity(input.len().min(MAX_SANITIZED_LENGTH)); + + for ch in input.chars().take(MAX_SANITIZED_LENGTH) { + match ch { + // Replace newlines and carriage returns with escaped versions + '\n' => result.push_str("\\n"), + '\r' => result.push_str("\\r"), + // Replace tabs with escaped version + '\t' => result.push_str("\\t"), + // Escape backslashes + '\\' => result.push_str("\\\\"), + // Remove other control characters (ASCII 0-31, 127) + c if c.is_control() => { + result.push_str(&format!("\\x{:02x}", c as u32)); + } + // Keep printable characters as-is + c => result.push(c), + } + } + + // Indicate truncation if input was too long + if input.len() > MAX_SANITIZED_LENGTH { + result.push_str("...[TRUNCATED]"); + } + + result +} + +/// Outcome of an audited operation. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] +#[serde(rename_all = "lowercase")] +pub enum AuditOutcome { + /// Operation completed successfully + Success, + /// Operation denied (authorization failure) + Denied, +} + +impl std::fmt::Display for AuditOutcome { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + AuditOutcome::Success => write!(f, "success"), + AuditOutcome::Denied => write!(f, "denied"), + } + } +} + +/// Audit event for control plane operations. +#[derive(Debug, Clone, Serialize)] +pub struct AuditEvent { + /// Timestamp of the event + pub timestamp: DateTime, + + /// Principal who performed the action (subject ID or API key ID) + pub principal: String, + + /// Authentication method used (jwt, api_key) + pub auth_method: String, + + /// Role of the principal + pub role: Role, + + /// HTTP method + pub method: String, + + /// Request path + pub path: String, + + /// Resource being accessed (e.g., worker ID, wasm module ID) + #[serde(skip_serializing_if = "Option::is_none")] + pub resource: Option, + + /// Operation outcome + pub outcome: AuditOutcome, + + /// Request ID for correlation + #[serde(skip_serializing_if = "Option::is_none")] + pub request_id: Option, + + /// Additional details or error message + #[serde(skip_serializing_if = "Option::is_none")] + pub details: Option, +} + +/// Context for audit logging containing request and principal information. +#[derive(Debug, Clone)] +pub struct AuditContext<'a> { + /// Principal who performed the action + pub principal: &'a str, + /// Authentication method used (jwt, api_key) + pub auth_method: &'a str, + /// Role of the principal + pub role: Role, + /// HTTP method + pub method: &'a str, + /// Request path + pub path: &'a str, + /// Request ID for correlation + pub request_id: Option<&'a str>, +} + +impl<'a> AuditContext<'a> { + /// Create a new audit context. + pub fn new( + principal: &'a str, + auth_method: &'a str, + role: Role, + method: &'a str, + path: &'a str, + request_id: Option<&'a str>, + ) -> Self { + Self { + principal, + auth_method, + role, + method, + path, + request_id, + } + } +} + +impl AuditEvent { + /// Create a new audit event from context. + fn from_context( + ctx: &AuditContext<'_>, + outcome: AuditOutcome, + resource: Option<&str>, + details: Option<&str>, + ) -> Self { + Self { + timestamp: Utc::now(), + principal: sanitize_for_log(ctx.principal), + auth_method: sanitize_for_log(ctx.auth_method), + role: ctx.role, + method: sanitize_for_log(ctx.method), + path: sanitize_for_log(ctx.path), + resource: resource.map(sanitize_for_log), + outcome, + request_id: ctx.request_id.map(sanitize_for_log), + details: details.map(sanitize_for_log), + } + } + + /// Create an audit event for unauthenticated requests. + fn unauthenticated(method: &str, path: &str, reason: &str, request_id: Option<&str>) -> Self { + Self { + timestamp: Utc::now(), + principal: "unauthenticated".to_string(), + auth_method: "none".to_string(), + role: Role::User, + method: sanitize_for_log(method), + path: sanitize_for_log(path), + resource: None, + outcome: AuditOutcome::Denied, + request_id: request_id.map(sanitize_for_log), + details: Some(sanitize_for_log(reason)), + } + } +} + +/// Audit logger that emits structured audit events. +#[derive(Clone, Default)] +pub struct AuditLogger { + enabled: bool, +} + +impl AuditLogger { + /// Create a new audit logger. + pub fn new(enabled: bool) -> Self { + Self { enabled } + } + + /// Check if audit logging is enabled. + pub fn is_enabled(&self) -> bool { + self.enabled + } + + /// Log an audit event. + pub fn log(&self, event: &AuditEvent) { + if !self.enabled { + return; + } + + // Create a span for structured logging + let _span = span!( + Level::INFO, + "audit", + principal = %event.principal, + auth_method = %event.auth_method, + role = %event.role, + method = %event.method, + path = %event.path, + outcome = %event.outcome, + ) + .entered(); + + // Log the event + info!( + target: "sgl_model_gateway::audit", + timestamp = %event.timestamp.to_rfc3339(), + principal = %event.principal, + auth_method = %event.auth_method, + role = %event.role, + method = %event.method, + path = %event.path, + resource = ?event.resource, + outcome = %event.outcome, + request_id = ?event.request_id, + details = ?event.details, + "control_plane_audit" + ); + } + + /// Log a successful operation. + pub fn log_success(&self, ctx: &AuditContext<'_>, resource: Option<&str>) { + self.log(&AuditEvent::from_context( + ctx, + AuditOutcome::Success, + resource, + None, + )); + } + + /// Log a denied operation (authorization failure). + pub fn log_denied(&self, ctx: &AuditContext<'_>, reason: &str) { + self.log(&AuditEvent::from_context( + ctx, + AuditOutcome::Denied, + None, + Some(reason), + )); + } + + /// Log an authentication failure (before principal is known). + pub fn log_auth_failure( + &self, + method: &str, + path: &str, + reason: &str, + request_id: Option<&str>, + ) { + self.log(&AuditEvent::unauthenticated( + method, path, reason, request_id, + )); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_audit_event_from_context() { + let ctx = AuditContext::new( + "user@example.com", + "jwt", + Role::Admin, + "POST", + "/workers", + Some("req-abc"), + ); + let event = AuditEvent::from_context(&ctx, AuditOutcome::Success, Some("worker-123"), None); + + assert_eq!(event.principal, "user@example.com"); + assert_eq!(event.auth_method, "jwt"); + assert_eq!(event.role, Role::Admin); + assert_eq!(event.method, "POST"); + assert_eq!(event.path, "/workers"); + assert_eq!(event.resource, Some("worker-123".to_string())); + assert_eq!(event.outcome, AuditOutcome::Success); + assert_eq!(event.request_id, Some("req-abc".to_string())); + } + + #[test] + fn test_audit_outcome_display() { + assert_eq!(AuditOutcome::Success.to_string(), "success"); + assert_eq!(AuditOutcome::Denied.to_string(), "denied"); + } + + #[test] + fn test_sanitize_for_log_normal_input() { + assert_eq!(sanitize_for_log("normal string"), "normal string"); + assert_eq!(sanitize_for_log("user@example.com"), "user@example.com"); + assert_eq!(sanitize_for_log("/api/v1/workers"), "/api/v1/workers"); + } + + #[test] + fn test_sanitize_for_log_newlines() { + assert_eq!(sanitize_for_log("line1\nline2"), "line1\\nline2"); + assert_eq!(sanitize_for_log("line1\r\nline2"), "line1\\r\\nline2"); + } + + #[test] + fn test_sanitize_for_log_control_chars() { + assert_eq!(sanitize_for_log("test\x00null"), "test\\x00null"); + assert_eq!(sanitize_for_log("test\x1bescape"), "test\\x1bescape"); + } + + #[test] + fn test_sanitize_for_log_backslashes() { + assert_eq!(sanitize_for_log("path\\to\\file"), "path\\\\to\\\\file"); + } + + #[test] + fn test_sanitize_for_log_truncation() { + let long_string = "a".repeat(2000); + let sanitized = sanitize_for_log(&long_string); + assert!(sanitized.len() < 2000); + assert!(sanitized.ends_with("...[TRUNCATED]")); + } + + #[test] + fn test_audit_event_sanitizes_inputs() { + let ctx = AuditContext::new( + "user\ninjected", + "jwt", + Role::User, + "GET", + "/workers\r\nfake_log_entry", + None, + ); + let event = AuditEvent::from_context( + &ctx, + AuditOutcome::Denied, + None, + Some("error\x00with\x1bnull"), + ); + + // Verify sanitization was applied + assert!(!event.principal.contains('\n')); + assert!(event.principal.contains("\\n")); + assert!(!event.path.contains('\r')); + assert!(event.path.contains("\\r")); + assert!(event.details.as_ref().unwrap().contains("\\x00")); + } +} diff --git a/sgl-model-gateway/src/auth/config.rs b/sgl-model-gateway/src/auth/config.rs new file mode 100644 index 000000000..774896d19 --- /dev/null +++ b/sgl-model-gateway/src/auth/config.rs @@ -0,0 +1,345 @@ +//! Configuration types for control plane authentication. +//! +//! Security features: +//! - API keys are hashed at load time (never stored in plaintext in memory) +//! - Constant-time comparison for API key validation + +use std::collections::HashMap; + +use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; +use subtle::ConstantTimeEq; + +/// Role-based access control for control plane APIs. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)] +#[serde(rename_all = "lowercase")] +pub enum Role { + /// Full access to all control plane APIs (workers, wasm, tokenizers, etc.) + Admin, + /// Access to inference/data plane APIs only (default for backward compatibility) + #[default] + User, +} + +impl Role { + /// Check if this role has admin privileges for control plane APIs. + pub fn is_admin(&self) -> bool { + matches!(self, Role::Admin) + } +} + +impl std::fmt::Display for Role { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Role::Admin => write!(f, "admin"), + Role::User => write!(f, "user"), + } + } +} + +impl std::str::FromStr for Role { + type Err = String; + + fn from_str(s: &str) -> Result { + match s.to_lowercase().as_str() { + "admin" => Ok(Role::Admin), + "user" => Ok(Role::User), + _ => Err(format!("Invalid role: {}. Valid roles: admin, user", s)), + } + } +} + +/// JWT/OIDC configuration for external identity provider integration. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct JwtConfig { + /// OIDC issuer URL (e.g., "https://login.microsoftonline.com/{tenant}/v2.0") + /// Used to discover JWKS endpoint via .well-known/openid-configuration + pub issuer: String, + + /// Expected audience claim (usually the client ID or API identifier) + pub audience: String, + + /// Optional explicit JWKS URI. If not provided, discovered from issuer. + #[serde(skip_serializing_if = "Option::is_none")] + pub jwks_uri: Option, + + /// Claim name containing the role (default: "role" or "roles") + #[serde(default = "default_role_claim")] + pub role_claim: String, + + /// Mapping from IDP role/group names to gateway roles. + /// Example: {"Gateway.Admin": "admin", "Gateway.User": "user"} + #[serde(default)] + pub role_mapping: HashMap, + + /// Clock skew tolerance in seconds (default: 30) + #[serde(default = "default_leeway_secs")] + pub leeway_secs: u64, + + /// JWKS cache TTL in seconds (default: 3600 = 1 hour) + #[serde(default = "default_jwks_cache_ttl_secs")] + pub jwks_cache_ttl_secs: u64, +} + +fn default_role_claim() -> String { + "roles".to_string() +} + +fn default_leeway_secs() -> u64 { + 30 +} + +fn default_jwks_cache_ttl_secs() -> u64 { + 3600 +} + +impl JwtConfig { + /// Create a new JWT config with required fields. + pub fn new(issuer: impl Into, audience: impl Into) -> Self { + Self { + issuer: issuer.into(), + audience: audience.into(), + jwks_uri: None, + role_claim: default_role_claim(), + role_mapping: HashMap::new(), + leeway_secs: default_leeway_secs(), + jwks_cache_ttl_secs: default_jwks_cache_ttl_secs(), + } + } + + /// Set explicit JWKS URI instead of using OIDC discovery. + pub fn with_jwks_uri(mut self, jwks_uri: impl Into) -> Self { + self.jwks_uri = Some(jwks_uri.into()); + self + } + + /// Add a role mapping from IDP role to gateway role. + pub fn with_role_mapping(mut self, idp_role: impl Into, gateway_role: Role) -> Self { + self.role_mapping.insert(idp_role.into(), gateway_role); + self + } +} + +/// Hash an API key using SHA-256. +/// Returns the hash as a fixed-size byte array. +fn hash_api_key(key: &str) -> [u8; 32] { + let mut hasher = Sha256::new(); + hasher.update(key.as_bytes()); + hasher.finalize().into() +} + +/// API key entry for service account authentication. +/// The key is hashed at construction time and never stored in plaintext. +#[derive(Clone, Serialize, Deserialize)] +pub struct ApiKeyEntry { + /// Unique identifier for this key + pub id: String, + + /// Human-readable name/description + pub name: String, + + /// SHA-256 hash of the API key (never store plaintext) + #[serde(skip)] + key_hash: [u8; 32], + + /// Role assigned to this API key + #[serde(default)] + pub role: Role, +} + +impl std::fmt::Debug for ApiKeyEntry { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ApiKeyEntry") + .field("id", &self.id) + .field("name", &self.name) + .field("key_hash", &"[REDACTED]") + .field("role", &self.role) + .finish() + } +} + +impl ApiKeyEntry { + /// Create a new API key entry. + /// The key is immediately hashed and the plaintext is not stored. + pub fn new( + id: impl Into, + name: impl Into, + key: impl Into, + role: Role, + ) -> Self { + let key_str = key.into(); + Self { + id: id.into(), + name: name.into(), + key_hash: hash_api_key(&key_str), + role, + } + } + + /// Check if the provided key matches this entry. + /// Uses constant-time comparison to prevent timing attacks. + pub fn verify(&self, key: &str) -> bool { + let provided_hash = hash_api_key(key); + self.key_hash.ct_eq(&provided_hash).into() + } +} + +/// Complete control plane authentication configuration. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct ControlPlaneAuthConfig { + /// JWT/OIDC configuration for external IDP. + #[serde(skip_serializing_if = "Option::is_none")] + pub jwt: Option, + + /// API keys for service accounts (parsed from CLI). + /// Format: "role:key" where role is "admin" or "user" + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub api_keys: Vec, + + /// Enable audit logging for control plane operations + #[serde(default = "default_audit_enabled")] + pub audit_enabled: bool, +} + +fn default_audit_enabled() -> bool { + true +} + +impl ControlPlaneAuthConfig { + /// Check if any authentication method is configured. + pub fn is_enabled(&self) -> bool { + self.jwt.is_some() || !self.api_keys.is_empty() + } + + /// Check if JWT authentication is configured. + pub fn has_jwt(&self) -> bool { + self.jwt.is_some() + } + + /// Check if any API keys are configured. + pub fn has_api_keys(&self) -> bool { + !self.api_keys.is_empty() + } + + /// Find an API key by its value and return the entry if found. + /// Uses constant-time hash comparison to prevent timing attacks. + pub fn find_api_key(&self, key: &str) -> Option<&ApiKeyEntry> { + // Iterate through all keys to prevent timing leaks about key existence + // We use a variable to track the match to ensure constant-time behavior + let mut found: Option<&ApiKeyEntry> = None; + for entry in &self.api_keys { + if entry.verify(key) { + found = Some(entry); + // Don't break early to maintain constant-time behavior + } + } + found + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_role_parsing() { + assert_eq!("admin".parse::().unwrap(), Role::Admin); + assert_eq!("ADMIN".parse::().unwrap(), Role::Admin); + assert_eq!("user".parse::().unwrap(), Role::User); + assert_eq!("USER".parse::().unwrap(), Role::User); + assert!("invalid".parse::().is_err()); + } + + #[test] + fn test_role_display() { + assert_eq!(Role::Admin.to_string(), "admin"); + assert_eq!(Role::User.to_string(), "user"); + } + + #[test] + fn test_role_is_admin() { + assert!(Role::Admin.is_admin()); + assert!(!Role::User.is_admin()); + } + + #[test] + fn test_jwt_config_builder() { + let config = JwtConfig::new("https://issuer.example.com", "api://my-app") + .with_jwks_uri("https://issuer.example.com/.well-known/jwks.json") + .with_role_mapping("Gateway.Admin", Role::Admin) + .with_role_mapping("Gateway.User", Role::User); + + assert_eq!(config.issuer, "https://issuer.example.com"); + assert_eq!(config.audience, "api://my-app"); + assert!(config.jwks_uri.is_some()); + assert_eq!(config.role_mapping.len(), 2); + } + + #[test] + fn test_control_plane_auth_config() { + let mut config = ControlPlaneAuthConfig::default(); + assert!(!config.is_enabled()); + assert!(!config.has_jwt()); + assert!(!config.has_api_keys()); + + config.jwt = Some(JwtConfig::new("https://issuer.example.com", "api://test")); + assert!(config.is_enabled()); + assert!(config.has_jwt()); + + config.api_keys.push(ApiKeyEntry::new( + "test-key", + "Test Key", + "secret123", + Role::Admin, + )); + assert!(config.has_api_keys()); + } + + #[test] + fn test_find_api_key() { + let config = ControlPlaneAuthConfig { + jwt: None, + api_keys: vec![ + ApiKeyEntry::new("key1", "Key 1", "secret1", Role::Admin), + ApiKeyEntry::new("key2", "Key 2", "secret2", Role::User), + ], + audit_enabled: true, + }; + + let found = config.find_api_key("secret1"); + assert!(found.is_some()); + assert_eq!(found.unwrap().id, "key1"); + assert_eq!(found.unwrap().role, Role::Admin); + + let found = config.find_api_key("secret2"); + assert!(found.is_some()); + assert_eq!(found.unwrap().role, Role::User); + + assert!(config.find_api_key("invalid").is_none()); + } + + #[test] + fn test_api_key_hashing() { + let entry = ApiKeyEntry::new("test", "Test Key", "my-secret-key", Role::Admin); + + // Verify the correct key works + assert!(entry.verify("my-secret-key")); + + // Verify wrong keys don't work + assert!(!entry.verify("wrong-key")); + assert!(!entry.verify("my-secret-ke")); // prefix + assert!(!entry.verify("my-secret-keyy")); // suffix + assert!(!entry.verify("")); // empty + } + + #[test] + fn test_api_key_debug_redacts_hash() { + let entry = ApiKeyEntry::new("test", "Test Key", "secret", Role::Admin); + let debug_str = format!("{:?}", entry); + + // Should contain REDACTED, not the actual hash + assert!(debug_str.contains("[REDACTED]")); + // Should not contain the key or hash bytes + assert!(!debug_str.contains("secret")); + } +} diff --git a/sgl-model-gateway/src/auth/jwks.rs b/sgl-model-gateway/src/auth/jwks.rs new file mode 100644 index 000000000..f7d34aedd --- /dev/null +++ b/sgl-model-gateway/src/auth/jwks.rs @@ -0,0 +1,515 @@ +//! JWKS (JSON Web Key Set) fetching and caching. +//! +//! Handles: +//! - OIDC discovery to find JWKS endpoint +//! - Fetching and parsing JWKS +//! - Caching with TTL and automatic refresh +//! +//! Security features: +//! - SSRF protection: only HTTPS URLs allowed, private IPs blocked +//! - Response size limits to prevent DoS +//! - Redirect prevention + +use std::{ + net::{IpAddr, Ipv4Addr, Ipv6Addr}, + time::{Duration, Instant}, +}; + +use jsonwebtoken::jwk::{Jwk, JwkSet}; +use parking_lot::RwLock; +use tracing::{debug, info, warn}; +use url::Url; + +/// Maximum allowed JWKS response size (1 MB) +const MAX_JWKS_RESPONSE_SIZE: u64 = 1024 * 1024; + +/// Error types for JWKS operations. +#[derive(Debug, thiserror::Error)] +pub enum JwksError { + #[error("Failed to fetch OIDC discovery document: {0}")] + DiscoveryFetch(String), + + #[error("Failed to parse OIDC discovery document: {0}")] + DiscoveryParse(String), + + #[error("JWKS URI not found in discovery document")] + JwksUriNotFound, + + #[error("Failed to fetch JWKS: {0}")] + JwksFetch(String), + + #[error("Failed to parse JWKS: {0}")] + JwksParse(String), + + #[error("Key not found for kid: {0}")] + KeyNotFound(String), + + #[error("Invalid URL: {0}")] + InvalidUrl(String), + + #[error("SSRF protection: {0}")] + SsrfBlocked(String), + + #[error("Response too large: {0} bytes (max: {1})")] + ResponseTooLarge(u64, u64), + + #[error("Failed to create HTTP client: {0}")] + HttpClientError(String), +} + +/// Check if an IP address is private/internal (SSRF protection). +fn is_private_ip(ip: &IpAddr) -> bool { + match ip { + IpAddr::V4(ipv4) => { + ipv4.is_private() // 10.0.0.0/8, 172.16.0.0/12, 192.168.0.0/16 + || ipv4.is_loopback() // 127.0.0.0/8 + || ipv4.is_link_local() // 169.254.0.0/16 + || ipv4.is_broadcast() // 255.255.255.255 + || ipv4.is_unspecified() // 0.0.0.0 + || is_shared_address(ipv4) // 100.64.0.0/10 (CGNAT) + || is_documentation_v4(ipv4) // 192.0.2.0/24, 198.51.100.0/24, 203.0.113.0/24 + || is_cloud_metadata(ipv4) // 169.254.169.254 + } + IpAddr::V6(ipv6) => { + ipv6.is_loopback() // ::1 + || ipv6.is_unspecified() // :: + || is_unique_local(ipv6) // fc00::/7 + || is_link_local_v6(ipv6) // fe80::/10 + } + } +} + +/// Check for CGNAT shared address space (100.64.0.0/10) +fn is_shared_address(ip: &Ipv4Addr) -> bool { + let octets = ip.octets(); + octets[0] == 100 && (octets[1] & 0xC0) == 64 +} + +/// Check for IPv4 documentation addresses +fn is_documentation_v4(ip: &Ipv4Addr) -> bool { + let octets = ip.octets(); + // 192.0.2.0/24 (TEST-NET-1) + (octets[0] == 192 && octets[1] == 0 && octets[2] == 2) + // 198.51.100.0/24 (TEST-NET-2) + || (octets[0] == 198 && octets[1] == 51 && octets[2] == 100) + // 203.0.113.0/24 (TEST-NET-3) + || (octets[0] == 203 && octets[1] == 0 && octets[2] == 113) +} + +/// Check for cloud metadata endpoint (169.254.169.254) +fn is_cloud_metadata(ip: &Ipv4Addr) -> bool { + ip.octets() == [169, 254, 169, 254] +} + +/// Check for IPv6 unique local addresses (fc00::/7) +fn is_unique_local(ip: &Ipv6Addr) -> bool { + let segments = ip.segments(); + (segments[0] & 0xfe00) == 0xfc00 +} + +/// Check for IPv6 link-local addresses (fe80::/10) +fn is_link_local_v6(ip: &Ipv6Addr) -> bool { + let segments = ip.segments(); + (segments[0] & 0xffc0) == 0xfe80 +} + +/// Validate a URL for SSRF protection. +/// Returns the validated URL or an error. +/// +/// For testing purposes, HTTP is allowed for localhost/127.0.0.1 only. +/// In production, only HTTPS should be used. +pub fn validate_url(url_str: &str) -> Result { + let url = Url::parse(url_str) + .map_err(|e| JwksError::InvalidUrl(format!("Failed to parse URL: {}", e)))?; + + // Check if this is a localhost URL (allowed for testing with HTTP) + let is_localhost = match url.host_str() { + Some("localhost") | Some("127.0.0.1") => true, + Some(host) => host == "::1" || host == "[::1]", + None => false, + }; + + // Only allow HTTPS (except for localhost in tests) + match url.scheme() { + "https" => {} + "http" => { + if !is_localhost { + return Err(JwksError::SsrfBlocked( + "Only HTTPS URLs are allowed for JWKS endpoints".to_string(), + )); + } + // HTTP localhost is allowed - skip further checks and return early + return Ok(url); + } + scheme => { + return Err(JwksError::SsrfBlocked(format!( + "Invalid URL scheme '{}'. Only HTTPS is allowed", + scheme + ))); + } + } + + // For HTTPS URLs, check for private/internal addresses + if let Some(host) = url.host() { + match host { + url::Host::Ipv4(ip) => { + if is_private_ip(&IpAddr::V4(ip)) { + return Err(JwksError::SsrfBlocked(format!( + "Private/internal IP addresses are not allowed: {}", + ip + ))); + } + } + url::Host::Ipv6(ip) => { + if is_private_ip(&IpAddr::V6(ip)) { + return Err(JwksError::SsrfBlocked(format!( + "Private/internal IP addresses are not allowed: {}", + ip + ))); + } + } + url::Host::Domain(domain) => { + // Block common internal hostnames (except localhost which was handled above) + let lower = domain.to_lowercase(); + if lower == "metadata" + || lower == "metadata.google.internal" + || lower.ends_with(".internal") + || lower.ends_with(".local") + { + return Err(JwksError::SsrfBlocked(format!( + "Internal hostnames are not allowed: {}", + domain + ))); + } + } + } + } + + Ok(url) +} + +/// OIDC discovery document (subset of fields we need). +#[derive(Debug, serde::Deserialize)] +struct OidcDiscovery { + jwks_uri: String, + #[allow(dead_code)] + issuer: String, +} + +/// Cached JWKS with expiration tracking. +struct CachedJwks { + jwks: JwkSet, + fetched_at: Instant, + ttl: Duration, +} + +impl CachedJwks { + fn is_expired(&self) -> bool { + self.fetched_at.elapsed() > self.ttl + } +} + +/// JWKS provider with caching and automatic refresh. +pub struct JwksProvider { + /// HTTP client for fetching JWKS + client: reqwest::Client, + /// JWKS endpoint URL (validated) + jwks_uri: String, + /// Cached JWKS + cache: RwLock>, + /// Cache TTL + ttl: Duration, +} + +impl JwksProvider { + /// Create a new JWKS provider with explicit JWKS URI. + /// The URL is validated for SSRF protection. + pub fn new(jwks_uri: impl Into, ttl: Duration) -> Result { + let jwks_uri = jwks_uri.into(); + + // Validate URL for SSRF protection + validate_url(&jwks_uri)?; + + let client = reqwest::Client::builder() + .timeout(Duration::from_secs(10)) + .redirect(reqwest::redirect::Policy::none()) // Prevent SSRF via redirects + .build() + .map_err(|e| JwksError::HttpClientError(e.to_string()))?; + + Ok(Self { + client, + jwks_uri, + cache: RwLock::new(None), + ttl, + }) + } + + /// Create a new JWKS provider using OIDC discovery. + /// Both issuer and discovered JWKS URI are validated for SSRF protection. + pub async fn from_issuer(issuer: &str, ttl: Duration) -> Result { + // Normalize and validate issuer URL + let issuer = issuer.trim_end_matches('/'); + let discovery_url = format!("{}/.well-known/openid-configuration", issuer); + + // Validate discovery URL for SSRF protection + validate_url(&discovery_url)?; + + let client = reqwest::Client::builder() + .timeout(Duration::from_secs(10)) + .redirect(reqwest::redirect::Policy::none()) + .build() + .map_err(|e| JwksError::HttpClientError(e.to_string()))?; + + info!("Fetching OIDC discovery from: {}", discovery_url); + + let response = client + .get(&discovery_url) + .send() + .await + .map_err(|e| JwksError::DiscoveryFetch(e.to_string()))?; + + if !response.status().is_success() { + return Err(JwksError::DiscoveryFetch(format!( + "HTTP {}", + response.status() + ))); + } + + // Check response size before parsing + if let Some(content_length) = response.content_length() { + if content_length > MAX_JWKS_RESPONSE_SIZE { + return Err(JwksError::ResponseTooLarge( + content_length, + MAX_JWKS_RESPONSE_SIZE, + )); + } + } + + let discovery: OidcDiscovery = response + .json() + .await + .map_err(|e| JwksError::DiscoveryParse(e.to_string()))?; + + // Validate discovered JWKS URI for SSRF protection + validate_url(&discovery.jwks_uri)?; + + info!("Discovered JWKS URI: {}", discovery.jwks_uri); + + Ok(Self { + client, + jwks_uri: discovery.jwks_uri, + cache: RwLock::new(None), + ttl, + }) + } + + /// Get the JWKS URI. + pub fn jwks_uri(&self) -> &str { + &self.jwks_uri + } + + /// Fetch JWKS from the endpoint. + /// Response size is limited to prevent DoS attacks. + async fn fetch_jwks(&self) -> Result { + debug!("Fetching JWKS from: {}", self.jwks_uri); + + let response = self + .client + .get(&self.jwks_uri) + .send() + .await + .map_err(|e| JwksError::JwksFetch(e.to_string()))?; + + if !response.status().is_success() { + return Err(JwksError::JwksFetch(format!("HTTP {}", response.status()))); + } + + // Check response size before reading body + if let Some(content_length) = response.content_length() { + if content_length > MAX_JWKS_RESPONSE_SIZE { + return Err(JwksError::ResponseTooLarge( + content_length, + MAX_JWKS_RESPONSE_SIZE, + )); + } + } + + // Read body with size limit + let bytes = response + .bytes() + .await + .map_err(|e| JwksError::JwksFetch(e.to_string()))?; + + if bytes.len() as u64 > MAX_JWKS_RESPONSE_SIZE { + return Err(JwksError::ResponseTooLarge( + bytes.len() as u64, + MAX_JWKS_RESPONSE_SIZE, + )); + } + + let jwks: JwkSet = + serde_json::from_slice(&bytes).map_err(|e| JwksError::JwksParse(e.to_string()))?; + + debug!("Fetched JWKS with {} keys", jwks.keys.len()); + Ok(jwks) + } + + /// Get the cached JWKS, refreshing if expired or not present. + pub async fn get_jwks(&self) -> Result { + // Check cache first + { + let cache = self.cache.read(); + if let Some(cached) = cache.as_ref() { + if !cached.is_expired() { + return Ok(cached.jwks.clone()); + } + } + } + + // Cache miss or expired, fetch new JWKS + let jwks = self.fetch_jwks().await?; + + // Update cache + { + let mut cache = self.cache.write(); + *cache = Some(CachedJwks { + jwks: jwks.clone(), + fetched_at: Instant::now(), + ttl: self.ttl, + }); + } + + Ok(jwks) + } + + /// Get a specific key by kid (key ID). + pub async fn get_key(&self, kid: &str) -> Result { + let jwks = self.get_jwks().await?; + + // First try to find by kid + if let Some(key) = jwks.find(kid) { + return Ok(key.clone()); + } + + // Key not found - try refreshing the cache in case keys were rotated + warn!("Key {} not found in cached JWKS, refreshing...", kid); + let jwks = self.fetch_jwks().await?; + + // Update cache + { + let mut cache = self.cache.write(); + *cache = Some(CachedJwks { + jwks: jwks.clone(), + fetched_at: Instant::now(), + ttl: self.ttl, + }); + } + + jwks.find(kid) + .cloned() + .ok_or_else(|| JwksError::KeyNotFound(kid.to_string())) + } + + /// Force refresh the JWKS cache. + pub async fn refresh(&self) -> Result<(), JwksError> { + let jwks = self.fetch_jwks().await?; + let mut cache = self.cache.write(); + *cache = Some(CachedJwks { + jwks, + fetched_at: Instant::now(), + ttl: self.ttl, + }); + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_cached_jwks_expiration() { + let jwks = JwkSet { keys: vec![] }; + let cached = CachedJwks { + jwks, + fetched_at: Instant::now() - Duration::from_secs(100), + ttl: Duration::from_secs(60), + }; + assert!(cached.is_expired()); + + let jwks = JwkSet { keys: vec![] }; + let cached = CachedJwks { + jwks, + fetched_at: Instant::now(), + ttl: Duration::from_secs(60), + }; + assert!(!cached.is_expired()); + } + + #[test] + fn test_validate_url_https_required() { + // Valid HTTPS URL + assert!(validate_url("https://example.com/.well-known/jwks.json").is_ok()); + + // HTTP not allowed for non-localhost + assert!(validate_url("http://example.com/.well-known/jwks.json").is_err()); + + // HTTP allowed for localhost (testing) + assert!(validate_url("http://localhost:8080/jwks").is_ok()); + assert!(validate_url("http://127.0.0.1:8080/jwks").is_ok()); + + // Invalid schemes + assert!(validate_url("ftp://example.com/jwks").is_err()); + assert!(validate_url("file:///etc/passwd").is_err()); + } + + #[test] + fn test_validate_url_blocks_private_ips() { + // Private IPv4 ranges + assert!(validate_url("https://10.0.0.1/jwks").is_err()); + assert!(validate_url("https://172.16.0.1/jwks").is_err()); + assert!(validate_url("https://192.168.1.1/jwks").is_err()); + + // Loopback + assert!(validate_url("https://127.0.0.1/jwks").is_err()); + + // Link-local + assert!(validate_url("https://169.254.1.1/jwks").is_err()); + + // Cloud metadata endpoint + assert!(validate_url("https://169.254.169.254/jwks").is_err()); + + // CGNAT + assert!(validate_url("https://100.64.0.1/jwks").is_err()); + } + + #[test] + fn test_validate_url_blocks_internal_hostnames() { + assert!(validate_url("https://metadata/jwks").is_err()); + assert!(validate_url("https://metadata.google.internal/jwks").is_err()); + assert!(validate_url("https://internal.example.internal/jwks").is_err()); + assert!(validate_url("https://printer.local/jwks").is_err()); + } + + #[test] + fn test_is_private_ip() { + use std::net::Ipv4Addr; + + // Private ranges + assert!(is_private_ip(&IpAddr::V4(Ipv4Addr::new(10, 0, 0, 1)))); + assert!(is_private_ip(&IpAddr::V4(Ipv4Addr::new(172, 16, 0, 1)))); + assert!(is_private_ip(&IpAddr::V4(Ipv4Addr::new(192, 168, 1, 1)))); + + // Loopback + assert!(is_private_ip(&IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)))); + + // Cloud metadata + assert!(is_private_ip(&IpAddr::V4(Ipv4Addr::new( + 169, 254, 169, 254 + )))); + + // Public IPs should not be blocked + assert!(!is_private_ip(&IpAddr::V4(Ipv4Addr::new(8, 8, 8, 8)))); + assert!(!is_private_ip(&IpAddr::V4(Ipv4Addr::new(1, 1, 1, 1)))); + } +} diff --git a/sgl-model-gateway/src/auth/jwt.rs b/sgl-model-gateway/src/auth/jwt.rs new file mode 100644 index 000000000..d43d0352a --- /dev/null +++ b/sgl-model-gateway/src/auth/jwt.rs @@ -0,0 +1,532 @@ +//! JWT validation for control plane authentication. +//! +//! Supports: +//! - RS256, RS384, RS512 (RSA) +//! - ES256, ES384 (ECDSA) +//! - Audience and issuer validation +//! - Role extraction from claims +//! +//! Security features: +//! - Algorithm verification (token alg must match key algorithm) +//! - Optional JTI tracking for replay protection +//! - SSRF protection via JWKS provider + +use std::{ + collections::HashMap, + num::NonZeroUsize, + sync::Arc, + time::{Duration, Instant}, +}; + +use jsonwebtoken::{ + decode, decode_header, + jwk::{AlgorithmParameters, Jwk}, + Algorithm, DecodingKey, TokenData, Validation, +}; +use lru::LruCache; +use parking_lot::Mutex; +use serde::{Deserialize, Serialize}; +use tracing::{debug, warn}; + +use super::{ + config::{JwtConfig, Role}, + jwks::{JwksError, JwksProvider}, +}; + +/// Default size for JTI cache (number of tokens to track) +const DEFAULT_JTI_CACHE_SIZE: usize = 10_000; + +/// Error types for JWT validation. +#[derive(Debug, thiserror::Error)] +pub enum JwtValidatorError { + #[error("Invalid token format")] + InvalidFormat, + + #[error("Token header missing 'kid' claim")] + MissingKid, + + #[error("Failed to get signing key: {0}")] + KeyError(#[from] JwksError), + + #[error("Unsupported algorithm: {0:?}")] + UnsupportedAlgorithm(Algorithm), + + #[error("Token validation failed: {0}")] + ValidationFailed(String), + + #[error("Failed to decode token: {0}")] + DecodeFailed(#[from] jsonwebtoken::errors::Error), + + #[error("Failed to extract role from claims")] + RoleExtractionFailed, + + #[error("No role mapping found for: {0}")] + NoRoleMapping(String), + + #[error("Algorithm mismatch: token uses {token_alg:?} but key requires {key_alg:?}")] + AlgorithmMismatch { + token_alg: Algorithm, + key_alg: Algorithm, + }, + + #[error("Token replay detected: JTI '{0}' has already been used")] + TokenReplay(String), +} + +/// Standard JWT claims we extract. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct StandardClaims { + /// Subject (user ID) + pub sub: Option, + + /// Issuer + pub iss: Option, + + /// Audience (can be string or array) + #[serde(default)] + pub aud: Audience, + + /// Expiration time + pub exp: Option, + + /// Issued at + pub iat: Option, + + /// Not before + pub nbf: Option, + + /// JWT ID + pub jti: Option, + + /// Email (common claim) + pub email: Option, + + /// Name (common claim) + pub name: Option, + + /// Preferred username (OIDC claim) + pub preferred_username: Option, + + /// All other claims + #[serde(flatten)] + pub extra: HashMap, +} + +/// Audience claim can be a single string or an array. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +#[serde(untagged)] +pub enum Audience { + Single(String), + Multiple(Vec), + #[default] + None, +} + +impl Audience { + pub fn contains(&self, aud: &str) -> bool { + match self { + Audience::Single(s) => s == aud, + Audience::Multiple(v) => v.iter().any(|s| s == aud), + Audience::None => false, + } + } +} + +/// Validated token with extracted claims and role. +#[derive(Debug, Clone)] +pub struct ValidatedToken { + /// Subject (user ID) + pub subject: String, + + /// Issuer + pub issuer: String, + + /// Assigned role + pub role: Role, + + /// Email if present + pub email: Option, + + /// Display name if present + pub name: Option, + + /// Full claims for additional processing + pub claims: StandardClaims, +} + +/// JTI (JWT ID) cache entry with expiration tracking. +struct JtiCacheEntry { + /// When the token expires (for cleanup) + expires_at: Instant, +} + +/// JWT validator with JWKS integration. +pub struct JwtValidator { + /// JWKS provider for key fetching + jwks_provider: Arc, + + /// JWT configuration + config: JwtConfig, + + /// Pre-configured validation settings + validation: Validation, + + /// JTI cache for replay protection (optional) + /// Maps JTI -> expiration time + jti_cache: Option>>, + + /// Whether to enable JTI replay protection + enable_jti_check: bool, +} + +impl JwtValidator { + /// Create a new JWT validator with explicit JWKS URI. + pub fn new(config: JwtConfig, jwks_provider: Arc) -> Self { + Self::new_with_options(config, jwks_provider, false) + } + + /// Create a new JWT validator with optional JTI replay protection. + pub fn new_with_options( + config: JwtConfig, + jwks_provider: Arc, + enable_jti_check: bool, + ) -> Self { + let mut validation = Validation::default(); + + // Set audience + validation.set_audience(&[&config.audience]); + + // Set issuer + validation.set_issuer(&[&config.issuer]); + + // Set leeway for clock skew + validation.leeway = config.leeway_secs; + + // We'll set the algorithm per-token based on the key + validation.algorithms = vec![ + Algorithm::RS256, + Algorithm::RS384, + Algorithm::RS512, + Algorithm::ES256, + Algorithm::ES384, + ]; + + let jti_cache = if enable_jti_check { + Some(Mutex::new(LruCache::new( + NonZeroUsize::new(DEFAULT_JTI_CACHE_SIZE).unwrap(), + ))) + } else { + None + }; + + Self { + jwks_provider, + config, + validation, + jti_cache, + enable_jti_check, + } + } + + /// Create a new JWT validator using OIDC discovery. + pub async fn from_config(config: JwtConfig) -> Result { + Self::from_config_with_options(config, false).await + } + + /// Create a new JWT validator with optional JTI replay protection. + pub async fn from_config_with_options( + config: JwtConfig, + enable_jti_check: bool, + ) -> Result { + let ttl = Duration::from_secs(config.jwks_cache_ttl_secs); + + let jwks_provider = if let Some(jwks_uri) = &config.jwks_uri { + Arc::new(JwksProvider::new(jwks_uri.clone(), ttl)?) + } else { + Arc::new(JwksProvider::from_issuer(&config.issuer, ttl).await?) + }; + + Ok(Self::new_with_options( + config, + jwks_provider, + enable_jti_check, + )) + } + + /// Validate a JWT token and extract claims. + pub async fn validate(&self, token: &str) -> Result { + // Decode header to get kid and algorithm + let header = decode_header(token)?; + let kid = header.kid.ok_or(JwtValidatorError::MissingKid)?; + let token_algorithm = header.alg; + + debug!( + "Validating JWT with kid: {}, alg: {:?}", + kid, token_algorithm + ); + + // Get the signing key + let jwk = self.jwks_provider.get_key(&kid).await?; + + // Determine algorithm from JWK + let key_algorithm = Self::jwk_to_algorithm(&jwk)?; + + // SECURITY: Verify token algorithm matches key algorithm + // This prevents algorithm confusion attacks + if token_algorithm != key_algorithm { + warn!( + "Algorithm mismatch: token uses {:?} but key requires {:?}", + token_algorithm, key_algorithm + ); + return Err(JwtValidatorError::AlgorithmMismatch { + token_alg: token_algorithm, + key_alg: key_algorithm, + }); + } + + // Create decoding key from JWK + let decoding_key = Self::jwk_to_decoding_key(&jwk)?; + + // Create validation with specific algorithm + let mut validation = self.validation.clone(); + validation.algorithms = vec![key_algorithm]; + + // Decode and validate token + let token_data: TokenData = decode(token, &decoding_key, &validation)?; + + let claims = token_data.claims; + + // Check JTI for replay protection if enabled + if self.enable_jti_check { + if let Some(jti) = &claims.jti { + self.check_jti_replay(jti, &claims)?; + } + } + + // Extract subject + let subject = claims + .sub + .clone() + .or_else(|| claims.email.clone()) + .or_else(|| claims.preferred_username.clone()) + .unwrap_or_else(|| "unknown".to_string()); + + // Extract issuer + let issuer = claims + .iss + .clone() + .unwrap_or_else(|| self.config.issuer.clone()); + + // Extract role + let role = self.extract_role(&claims)?; + + debug!( + "JWT validated: subject={}, issuer={}, role={:?}", + subject, issuer, role + ); + + Ok(ValidatedToken { + subject, + issuer, + role, + email: claims.email.clone(), + name: claims.name.clone(), + claims, + }) + } + + /// Check if a JTI has been used before (replay protection). + fn check_jti_replay( + &self, + jti: &str, + claims: &StandardClaims, + ) -> Result<(), JwtValidatorError> { + let Some(cache) = &self.jti_cache else { + return Ok(()); + }; + + let mut cache = cache.lock(); + + // Clean up expired entries first (lazy cleanup) + let now = Instant::now(); + + // Check if JTI exists and is still valid + if let Some(entry) = cache.get(jti) { + if entry.expires_at > now { + // Token is still valid but JTI was already used + return Err(JwtValidatorError::TokenReplay(jti.to_string())); + } + // Entry expired, remove it + cache.pop(jti); + } + + // Calculate expiration time from claims + let expires_at = if let Some(exp) = claims.exp { + // exp is Unix timestamp + let exp_duration = Duration::from_secs(exp); + let now_unix = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default(); + + if exp_duration > now_unix { + now + (exp_duration - now_unix) + } else { + // Already expired, but we still track it briefly + now + Duration::from_secs(60) + } + } else { + // No exp claim, use default TTL + now + Duration::from_secs(3600) + }; + + // Record this JTI + cache.put(jti.to_string(), JtiCacheEntry { expires_at }); + + Ok(()) + } + + /// Extract role from claims using configured role claim and mapping. + fn extract_role(&self, claims: &StandardClaims) -> Result { + // Try to get the role claim value + let role_value = claims.extra.get(&self.config.role_claim); + + let role_strings: Vec = match role_value { + Some(serde_json::Value::String(s)) => vec![s.clone()], + Some(serde_json::Value::Array(arr)) => arr + .iter() + .filter_map(|v| v.as_str().map(String::from)) + .collect(), + None => { + // Try alternate claim names + let alternates = ["role", "roles", "groups", "group"]; + let mut found = Vec::new(); + for alt in alternates { + if let Some(v) = claims.extra.get(alt) { + match v { + serde_json::Value::String(s) => found.push(s.clone()), + serde_json::Value::Array(arr) => found + .extend(arr.iter().filter_map(|v| v.as_str().map(String::from))), + _ => {} + } + } + } + found + } + _ => Vec::new(), + }; + + // If no role mapping configured, check for direct "admin" or "user" values + if self.config.role_mapping.is_empty() { + for role_str in &role_strings { + if let Ok(role) = role_str.parse::() { + return Ok(role); + } + } + // Default to User if no explicit role found + warn!("No role found in JWT claims, defaulting to User"); + return Ok(Role::User); + } + + // Use role mapping + for role_str in &role_strings { + if let Some(role) = self.config.role_mapping.get(role_str) { + return Ok(*role); + } + } + + // Check if any mapped role is admin - if so, we need explicit mapping + // Otherwise default to User for safety + warn!( + "No matching role mapping found for {:?}, defaulting to User", + role_strings + ); + Ok(Role::User) + } + + /// Convert a JWK to a DecodingKey. + fn jwk_to_decoding_key(jwk: &Jwk) -> Result { + match &jwk.algorithm { + AlgorithmParameters::RSA(rsa) => Ok(DecodingKey::from_rsa_components(&rsa.n, &rsa.e)?), + AlgorithmParameters::EllipticCurve(ec) => { + Ok(DecodingKey::from_ec_components(&ec.x, &ec.y)?) + } + AlgorithmParameters::OctetKey(_) => { + Err(JwtValidatorError::UnsupportedAlgorithm(Algorithm::HS256)) + } + AlgorithmParameters::OctetKeyPair(_) => { + Err(JwtValidatorError::UnsupportedAlgorithm(Algorithm::EdDSA)) + } + } + } + + /// Determine the algorithm from a JWK. + fn jwk_to_algorithm(jwk: &Jwk) -> Result { + // First check if algorithm is explicitly specified + if let Some(alg) = &jwk.common.key_algorithm { + return Ok(match alg { + jsonwebtoken::jwk::KeyAlgorithm::RS256 => Algorithm::RS256, + jsonwebtoken::jwk::KeyAlgorithm::RS384 => Algorithm::RS384, + jsonwebtoken::jwk::KeyAlgorithm::RS512 => Algorithm::RS512, + jsonwebtoken::jwk::KeyAlgorithm::ES256 => Algorithm::ES256, + jsonwebtoken::jwk::KeyAlgorithm::ES384 => Algorithm::ES384, + other => { + return Err(JwtValidatorError::ValidationFailed(format!( + "Unsupported key algorithm: {:?}", + other + ))) + } + }); + } + + // Infer from key type + match &jwk.algorithm { + AlgorithmParameters::RSA(_) => Ok(Algorithm::RS256), // Default RSA to RS256 + AlgorithmParameters::EllipticCurve(ec) => { + use jsonwebtoken::jwk::EllipticCurve; + match ec.curve { + EllipticCurve::P256 => Ok(Algorithm::ES256), + EllipticCurve::P384 => Ok(Algorithm::ES384), + // Other curves not supported for ECDSA + _ => Err(JwtValidatorError::ValidationFailed(format!( + "Unsupported EC curve: {:?}", + ec.curve + ))), + } + } + _ => Err(JwtValidatorError::ValidationFailed( + "Cannot determine algorithm from key".to_string(), + )), + } + } + + /// Get a reference to the JWKS provider. + pub fn jwks_provider(&self) -> &Arc { + &self.jwks_provider + } + + /// Get a reference to the JWT config. + pub fn config(&self) -> &JwtConfig { + &self.config + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_audience_contains() { + let single = Audience::Single("api://test".to_string()); + assert!(single.contains("api://test")); + assert!(!single.contains("other")); + + let multiple = + Audience::Multiple(vec!["api://test".to_string(), "api://other".to_string()]); + assert!(multiple.contains("api://test")); + assert!(multiple.contains("api://other")); + assert!(!multiple.contains("unknown")); + + let none = Audience::None; + assert!(!none.contains("anything")); + } +} diff --git a/sgl-model-gateway/src/auth/middleware.rs b/sgl-model-gateway/src/auth/middleware.rs new file mode 100644 index 000000000..5e82e1b3e --- /dev/null +++ b/sgl-model-gateway/src/auth/middleware.rs @@ -0,0 +1,399 @@ +//! Control plane authentication middleware. +//! +//! Provides middleware for authenticating and authorizing access to control plane APIs. +//! Supports both JWT/OIDC tokens and API keys. + +use std::sync::Arc; + +use axum::{ + body::Body, + extract::{Request, State}, + http::{header, StatusCode}, + middleware::Next, + response::{IntoResponse, Response}, +}; +use tracing::{debug, error, info, warn}; + +use super::{ + audit::{AuditContext, AuditLogger}, + config::{ControlPlaneAuthConfig, Role}, + jwt::JwtValidator, +}; +use crate::middleware::RequestId; + +/// Authenticated principal information. +#[derive(Debug, Clone)] +pub struct Principal { + /// Subject identifier (user ID, email, or API key ID) + pub id: String, + + /// Display name if available + pub name: Option, + + /// Authentication method used + pub auth_method: AuthMethod, + + /// Assigned role + pub role: Role, +} + +/// Authentication method used to authenticate the principal. +#[derive(Debug, Clone)] +pub enum AuthMethod { + /// JWT/OIDC token from external IDP + Jwt { issuer: String }, + /// API key for service accounts + ApiKey { key_id: String }, +} + +impl std::fmt::Display for AuthMethod { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + AuthMethod::Jwt { issuer } => write!(f, "jwt:{}", issuer), + AuthMethod::ApiKey { key_id } => write!(f, "api_key:{}", key_id), + } + } +} + +/// Extension trait for extracting Principal from request extensions. +pub trait PrincipalExt { + fn principal(&self) -> Option<&Principal>; +} + +impl PrincipalExt for http::Request { + fn principal(&self) -> Option<&Principal> { + self.extensions().get::() + } +} + +/// State for the control plane authentication middleware. +#[derive(Clone)] +pub struct ControlPlaneAuthState { + /// Authentication configuration + pub config: ControlPlaneAuthConfig, + + /// JWT validator (if JWT auth is configured) + pub jwt_validator: Option>, + + /// Audit logger + pub audit_logger: AuditLogger, +} + +impl ControlPlaneAuthState { + /// Create a new control plane auth state. + pub fn new(config: ControlPlaneAuthConfig, jwt_validator: Option>) -> Self { + let audit_logger = AuditLogger::new(config.audit_enabled); + Self { + config, + jwt_validator, + audit_logger, + } + } + + /// Create from config, initializing JWT validator if needed. + pub async fn from_config( + config: ControlPlaneAuthConfig, + ) -> Result> { + let jwt_validator = if let Some(jwt_config) = &config.jwt { + Some(Arc::new( + JwtValidator::from_config(jwt_config.clone()).await?, + )) + } else { + None + }; + + Ok(Self::new(config, jwt_validator)) + } + + /// Try to initialize control plane auth from config. + /// + /// Returns `Some(state)` if auth is configured and initialized successfully. + /// Returns `None` if auth is not configured or initialization fails (with error logged). + pub async fn try_init(config: Option<&ControlPlaneAuthConfig>) -> Option { + let config = config.filter(|c| c.is_enabled())?; + + info!("Initializing control plane authentication..."); + match Self::from_config(config.clone()).await { + Ok(state) => { + if config.has_jwt() { + info!("Control plane JWT/OIDC authentication enabled"); + } + if config.has_api_keys() { + info!( + "Control plane API key authentication enabled ({} keys)", + config.api_keys.len() + ); + } + if config.audit_enabled { + info!("Control plane audit logging enabled"); + } + Some(state) + } + Err(e) => { + error!( + "Failed to initialize control plane auth: {}. Falling back to simple API key auth.", + e + ); + None + } + } + } + + /// Check if authentication is required. + pub fn is_auth_required(&self) -> bool { + self.config.is_enabled() + } +} + +/// Check admin role and log denial if not admin. +/// Returns Some(Response) if denied, None if allowed. +fn check_admin_role( + principal_id: &str, + auth_method: &str, + role: Role, + method: &str, + path: &str, + request_id: Option<&str>, + audit_logger: &AuditLogger, +) -> Option { + if role.is_admin() { + return None; + } + + warn!( + "{} {} has role {:?} but admin is required for control plane access", + auth_method, principal_id, role + ); + let ctx = AuditContext::new(principal_id, auth_method, role, method, path, request_id); + audit_logger.log_denied(&ctx, "Admin role required for control plane access"); + + Some( + ( + StatusCode::FORBIDDEN, + "Admin role required for control plane access", + ) + .into_response(), + ) +} + +/// Log successful authentication. +fn log_auth_success( + principal: &Principal, + auth_method: &str, + method: &str, + path: &str, + request_id: Option<&str>, + audit_logger: &AuditLogger, +) { + debug!( + "{} authentication successful for {} with role {:?}", + auth_method, principal.id, principal.role + ); + let ctx = AuditContext::new( + &principal.id, + auth_method, + principal.role, + method, + path, + request_id, + ); + audit_logger.log_success(&ctx, None); +} + +/// Control plane authentication middleware. +/// +/// This middleware: +/// 1. Extracts the Bearer token from the Authorization header +/// 2. Attempts JWT validation first (if configured) +/// 3. Falls back to API key validation (if configured) +/// 4. Checks if the authenticated principal has admin role +/// 5. Logs audit events for control plane access +/// +/// Returns 401 Unauthorized if authentication fails. +/// Returns 403 Forbidden if the user doesn't have admin role. +pub async fn control_plane_auth_middleware( + State(auth_state): State, + mut request: Request, + next: Next, +) -> Response { + // If no authentication is configured, allow through (backward compatibility) + if !auth_state.is_auth_required() { + return next.run(request).await; + } + + let method = request.method().to_string(); + let path = request.uri().path().to_string(); + let request_id = request.extensions().get::().map(|r| r.0.clone()); + + // Extract Bearer token from Authorization header + let token = request + .headers() + .get(header::AUTHORIZATION) + .and_then(|h| h.to_str().ok()) + .and_then(|h| h.strip_prefix("Bearer ")); + + let Some(token) = token else { + debug!("Missing or invalid Authorization header for control plane API"); + auth_state.audit_logger.log_auth_failure( + &method, + &path, + "Missing or invalid Authorization header", + request_id.as_deref(), + ); + return ( + StatusCode::UNAUTHORIZED, + [("WWW-Authenticate", "Bearer realm=\"control-plane\"")], + "Missing or invalid Authorization header", + ) + .into_response(); + }; + + // Try JWT validation first + if let Some(jwt_validator) = &auth_state.jwt_validator { + match jwt_validator.validate(token).await { + Ok(validated_token) => { + if let Some(resp) = check_admin_role( + &validated_token.subject, + "jwt", + validated_token.role, + &method, + &path, + request_id.as_deref(), + &auth_state.audit_logger, + ) { + return resp; + } + + let principal = Principal { + id: validated_token.subject.clone(), + name: validated_token.name.clone(), + auth_method: AuthMethod::Jwt { + issuer: validated_token.issuer.clone(), + }, + role: validated_token.role, + }; + + log_auth_success( + &principal, + "jwt", + &method, + &path, + request_id.as_deref(), + &auth_state.audit_logger, + ); + request.extensions_mut().insert(principal); + return next.run(request).await; + } + Err(e) => { + // If the token looks like a JWT (3 parts separated by dots), it's likely + // an invalid JWT, not an API key. Fail fast with a specific error + // instead of silently falling back. This provides better feedback. + if token.split('.').count() == 3 { + warn!("Invalid JWT provided: {}. Not falling back to API key.", e); + auth_state.audit_logger.log_auth_failure( + &method, + &path, + &format!("Invalid JWT: {}", e), + request_id.as_deref(), + ); + return ( + StatusCode::UNAUTHORIZED, + [("WWW-Authenticate", "Bearer realm=\"control-plane\"")], + format!("Invalid JWT: {}", e), + ) + .into_response(); + } + debug!("JWT validation failed: {}, trying API key", e); + } + } + } + + // Try API key validation + if let Some(api_key_entry) = auth_state.config.find_api_key(token) { + if let Some(resp) = check_admin_role( + &api_key_entry.id, + "api_key", + api_key_entry.role, + &method, + &path, + request_id.as_deref(), + &auth_state.audit_logger, + ) { + return resp; + } + + let principal = Principal { + id: api_key_entry.id.clone(), + name: Some(api_key_entry.name.clone()), + auth_method: AuthMethod::ApiKey { + key_id: api_key_entry.id.clone(), + }, + role: api_key_entry.role, + }; + + log_auth_success( + &principal, + "api_key", + &method, + &path, + request_id.as_deref(), + &auth_state.audit_logger, + ); + request.extensions_mut().insert(principal); + return next.run(request).await; + } + + // Authentication failed + debug!("Control plane authentication failed: invalid token"); + auth_state.audit_logger.log_auth_failure( + &method, + &path, + "Invalid token", + request_id.as_deref(), + ); + + ( + StatusCode::UNAUTHORIZED, + [("WWW-Authenticate", "Bearer realm=\"control-plane\"")], + "Invalid authentication token", + ) + .into_response() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_auth_method_display() { + let jwt = AuthMethod::Jwt { + issuer: "https://example.com".to_string(), + }; + assert_eq!(jwt.to_string(), "jwt:https://example.com"); + + let api_key = AuthMethod::ApiKey { + key_id: "key-123".to_string(), + }; + assert_eq!(api_key.to_string(), "api_key:key-123"); + } + + #[test] + fn test_control_plane_auth_state_no_config() { + let config = ControlPlaneAuthConfig::default(); + let state = ControlPlaneAuthState::new(config, None); + assert!(!state.is_auth_required()); + } + + #[test] + fn test_control_plane_auth_state_with_api_keys() { + use super::super::config::ApiKeyEntry; + + let config = ControlPlaneAuthConfig { + jwt: None, + api_keys: vec![ApiKeyEntry::new("test", "Test Key", "secret", Role::Admin)], + audit_enabled: true, + }; + let state = ControlPlaneAuthState::new(config, None); + assert!(state.is_auth_required()); + } +} diff --git a/sgl-model-gateway/src/auth/mod.rs b/sgl-model-gateway/src/auth/mod.rs new file mode 100644 index 000000000..947a2e895 --- /dev/null +++ b/sgl-model-gateway/src/auth/mod.rs @@ -0,0 +1,20 @@ +//! Authentication and authorization module for control plane APIs. +//! +//! This module provides: +//! - JWT/OIDC authentication for external IDP integration +//! - API key authentication with role-based access +//! - Audit logging for control plane operations +//! - Middleware for securing admin and worker routes + +mod audit; +mod config; +mod jwks; +mod jwt; +mod middleware; + +pub use audit::{AuditEvent, AuditLogger, AuditOutcome}; +pub use config::{ApiKeyEntry, ControlPlaneAuthConfig, JwtConfig, Role}; +pub use jwt::{JwtValidator, JwtValidatorError}; +pub use middleware::{ + control_plane_auth_middleware, ControlPlaneAuthState, Principal, PrincipalExt, +}; diff --git a/sgl-model-gateway/src/lib.rs b/sgl-model-gateway/src/lib.rs index dac820881..a4ee2e714 100644 --- a/sgl-model-gateway/src/lib.rs +++ b/sgl-model-gateway/src/lib.rs @@ -1,4 +1,5 @@ pub mod app_context; +pub mod auth; pub mod config; pub mod core; pub mod data_connector; diff --git a/sgl-model-gateway/src/main.rs b/sgl-model-gateway/src/main.rs index 263e22b03..445a7f8cf 100644 --- a/sgl-model-gateway/src/main.rs +++ b/sgl-model-gateway/src/main.rs @@ -2,6 +2,7 @@ use std::collections::HashMap; use clap::{ArgAction, Parser, Subcommand, ValueEnum}; use sgl_model_gateway::{ + auth::{ApiKeyEntry, ControlPlaneAuthConfig, JwtConfig, Role}, config::{ CircuitBreakerConfig, ConfigError, ConfigResult, DiscoveryConfig, HealthCheckConfig, HistoryBackend, MetricsConfig, OracleConfig, PolicyConfig, PostgresConfig, RetryConfig, @@ -372,6 +373,40 @@ struct CliArgs { #[arg(long)] tls_key_path: Option, + + // ==================== Control Plane Authentication ==================== + /// JWT issuer URL for control plane authentication (OIDC discovery will be used) + /// Example: https://login.microsoftonline.com/{tenant}/v2.0 + #[arg(long, env = "JWT_ISSUER")] + jwt_issuer: Option, + + /// Expected JWT audience claim (usually the client ID or API identifier) + /// Example: api://sgl-gateway + #[arg(long, env = "JWT_AUDIENCE")] + jwt_audience: Option, + + /// Explicit JWKS URI (if not provided, discovered from issuer) + #[arg(long, env = "JWT_JWKS_URI")] + jwt_jwks_uri: Option, + + /// JWT claim name containing the role (default: "roles") + #[arg(long, default_value = "roles")] + jwt_role_claim: String, + + /// Role mapping from IDP role to gateway role (format: "idp_role=gateway_role") + /// Can be specified multiple times. Example: --jwt-role-mapping "Gateway.Admin=admin" + #[arg(long, action = ArgAction::Append)] + jwt_role_mapping: Vec, + + /// API keys for control plane access (format: "id:name:role:key") + /// Can be specified multiple times. + /// Example: --control-plane-api-keys "svc1:CI Pipeline:admin:secret123" + #[arg(long = "control-plane-api-keys", action = ArgAction::Append, env = "CONTROL_PLANE_API_KEYS")] + control_plane_api_keys: Vec, + + /// Disable audit logging for control plane operations + #[arg(long, default_value_t = false)] + disable_audit_logging: bool, } enum OracleConnectSource { @@ -379,7 +414,106 @@ enum OracleConnectSource { Wallet { path: String, alias: String }, } +/// Parse role mapping from CLI format "idp_role=gateway_role" +fn parse_role_mapping(mapping: &str) -> Option<(String, Role)> { + let parts: Vec<&str> = mapping.splitn(2, '=').collect(); + if parts.len() != 2 { + eprintln!( + "WARNING: Invalid role mapping format '{}'. Expected 'idp_role=gateway_role'", + mapping + ); + return None; + } + let idp_role = parts[0].to_string(); + let gateway_role = match parts[1].to_lowercase().as_str() { + "admin" => Role::Admin, + "user" => Role::User, + other => { + eprintln!( + "WARNING: Invalid gateway role '{}' in mapping. Valid roles: admin, user", + other + ); + return None; + } + }; + Some((idp_role, gateway_role)) +} + +/// Parse control plane API key from CLI format "id:name:role:key" +fn parse_control_plane_api_key(key_str: &str) -> Option { + let parts: Vec<&str> = key_str.splitn(4, ':').collect(); + if parts.len() != 4 { + eprintln!( + "WARNING: Invalid control-plane-api-key format '{}'. Expected 'id:name:role:key'", + key_str + ); + return None; + } + let id = parts[0]; + let name = parts[1]; + let role_str = parts[2]; + let key = parts[3]; + + let role = match role_str.to_lowercase().as_str() { + "admin" => Role::Admin, + "user" => Role::User, + other => { + eprintln!( + "WARNING: Invalid role '{}' in control-plane-api-key. Valid roles: admin, user", + other + ); + return None; + } + }; + + Some(ApiKeyEntry::new(id, name, key, role)) +} + impl CliArgs { + /// Build control plane authentication configuration from CLI args. + fn build_control_plane_auth_config(&self) -> ControlPlaneAuthConfig { + // Build JWT config if issuer and audience are provided + let jwt = match (&self.jwt_issuer, &self.jwt_audience) { + (Some(issuer), Some(audience)) => { + let role_mapping: HashMap = self + .jwt_role_mapping + .iter() + .filter_map(|m| parse_role_mapping(m)) + .collect(); + + let mut jwt_config = JwtConfig::new(issuer.clone(), audience.clone()); + jwt_config.role_claim = self.jwt_role_claim.clone(); + jwt_config.role_mapping = role_mapping; + if let Some(jwks_uri) = &self.jwt_jwks_uri { + jwt_config.jwks_uri = Some(jwks_uri.clone()); + } + Some(jwt_config) + } + (Some(_), None) => { + eprintln!("WARNING: --jwt-issuer provided but --jwt-audience is missing. JWT auth disabled."); + None + } + (None, Some(_)) => { + eprintln!("WARNING: --jwt-audience provided but --jwt-issuer is missing. JWT auth disabled."); + None + } + (None, None) => None, + }; + + // Build API keys from CLI args + let api_keys: Vec = self + .control_plane_api_keys + .iter() + .filter_map(|k| parse_control_plane_api_key(k)) + .collect(); + + ControlPlaneAuthConfig { + jwt, + api_keys, + audit_enabled: !self.disable_audit_logging, + } + } + fn determine_connection_mode(worker_urls: &[String]) -> ConnectionMode { for url in worker_urls { if url.starts_with("grpc://") || url.starts_with("grpcs://") { @@ -699,6 +833,16 @@ impl CliArgs { }, }); + // Build control plane auth config + let control_plane_auth = { + let config = self.build_control_plane_auth_config(); + if config.is_enabled() { + Some(config) + } else { + None + } + }; + ServerConfig { host: self.host.clone(), port: self.port, @@ -715,6 +859,7 @@ impl CliArgs { Some(self.request_id_headers.clone()) }, shutdown_grace_period_secs: self.shutdown_grace_period_secs, + control_plane_auth, } } } diff --git a/sgl-model-gateway/src/server.rs b/sgl-model-gateway/src/server.rs index 82149f9b7..95595e723 100644 --- a/sgl-model-gateway/src/server.rs +++ b/sgl-model-gateway/src/server.rs @@ -524,11 +524,14 @@ pub struct ServerConfig { pub request_timeout_secs: u64, pub request_id_headers: Option>, pub shutdown_grace_period_secs: u64, + /// Control plane authentication configuration + pub control_plane_auth: Option, } pub fn build_app( app_state: Arc, auth_config: AuthConfig, + control_plane_auth_state: Option, max_payload_size: usize, request_id_headers: Vec, cors_allowed_origins: Vec, @@ -593,6 +596,7 @@ pub fn build_app( .route("/get_model_info", get(get_model_info)) .route("/get_server_info", get(get_server_info)); + // Build admin routes with control plane auth if configured, otherwise use simple API key auth let admin_routes = Router::new() .route("/flush_cache", post(flush_cache)) .route("/get_loads", get(get_loads)) @@ -613,22 +617,32 @@ pub fn build_app( .route( "/v1/tokenizers/{tokenizer_id}/status", get(v1_tokenizers_status), - ) - .route_layer(axum::middleware::from_fn_with_state( - auth_config.clone(), - middleware::auth_middleware, - )); + ); + // Build worker routes let worker_routes = Router::new() .route("/workers", post(create_worker).get(list_workers_rest)) .route( "/workers/{worker_id}", get(get_worker).put(update_worker).delete(delete_worker), - ) - .route_layer(axum::middleware::from_fn_with_state( - auth_config.clone(), - middleware::auth_middleware, - )); + ); + + // Apply authentication middleware to control plane routes + let apply_control_plane_auth = |routes: Router>| { + if let Some(ref cp_state) = control_plane_auth_state { + routes.route_layer(axum::middleware::from_fn_with_state( + cp_state.clone(), + crate::auth::control_plane_auth_middleware, + )) + } else { + routes.route_layer(axum::middleware::from_fn_with_state( + auth_config.clone(), + middleware::auth_middleware, + )) + } + }; + let admin_routes = apply_control_plane_auth(admin_routes); + let worker_routes = apply_control_plane_auth(worker_routes); Router::new() .merge(protected_routes) @@ -881,9 +895,14 @@ pub async fn startup(config: ServerConfig) -> Result<(), Box Self { + // Generate a new 2048-bit RSA key pair for testing + use rsa::{pkcs8::EncodePrivateKey, rand_core::OsRng}; + + let mut rng = OsRng; + let private_key = + RsaPrivateKey::new(&mut rng, 2048).expect("Failed to generate RSA key pair"); + + // Get the public key components + let n_bytes = private_key.n().to_bytes_be(); + let e_bytes = private_key.e().to_bytes_be(); + + // Encode to base64url (no padding) for JWKS + let n_base64url = URL_SAFE_NO_PAD.encode(&n_bytes); + let e_base64url = URL_SAFE_NO_PAD.encode(&e_bytes); + + // Export private key as PEM + let private_key_pem = private_key + .to_pkcs8_pem(rsa::pkcs8::LineEnding::LF) + .expect("Failed to export private key") + .to_string(); + + Self { + private_key_pem, + n_base64url, + e_base64url, + } + } +} + +/// Lazily initialized test key pair (shared across all tests) +static TEST_KEYS: LazyLock = LazyLock::new(TestKeyPair::generate); + +/// Claims structure for test tokens +#[derive(Debug, Serialize, Deserialize)] +struct TestClaims { + sub: String, + iss: String, + aud: String, + exp: u64, + iat: u64, + #[serde(skip_serializing_if = "Option::is_none")] + jti: Option, + #[serde(skip_serializing_if = "Option::is_none")] + name: Option, + #[serde(skip_serializing_if = "Option::is_none")] + roles: Option>, +} + +/// Create a test JWT token using the generated test key pair +fn create_test_token(claims: &TestClaims) -> String { + let mut header = Header::new(jsonwebtoken::Algorithm::RS256); + header.kid = Some(TEST_KEY_ID.to_string()); + + let key = EncodingKey::from_rsa_pem(TEST_KEYS.private_key_pem.as_bytes()) + .expect("Failed to create encoding key"); + + encode(&header, claims, &key).expect("Failed to encode token") +} + +/// Create claims with default values +fn create_claims(sub: &str, roles: Vec<&str>) -> TestClaims { + let now = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_secs(); + + TestClaims { + sub: sub.to_string(), + iss: "https://test-issuer.example.com".to_string(), + aud: "test-gateway".to_string(), + exp: now + 3600, // 1 hour from now + iat: now, + jti: Some(uuid::Uuid::new_v4().to_string()), + name: Some("Test User".to_string()), + roles: Some(roles.into_iter().map(String::from).collect()), + } +} + +/// Start a mock JWKS server using the generated test key pair +async fn start_mock_jwks_server() -> (SocketAddr, tokio::task::JoinHandle<()>) { + // Use the generated key pair components + let jwks_response = json!({ + "keys": [{ + "kty": "RSA", + "use": "sig", + "kid": TEST_KEY_ID, + "alg": "RS256", + "n": TEST_KEYS.n_base64url, + "e": TEST_KEYS.e_base64url + }] + }); + + let app = Router::new().route( + "/.well-known/jwks.json", + get(move || { + let jwks = jwks_response.clone(); + async move { Json(jwks) } + }), + ); + + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + + let handle = tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + + // Give server time to start + tokio::time::sleep(Duration::from_millis(50)).await; + + (addr, handle) +} + +// ============================================================================ +// JWT Authentication Tests +// ============================================================================ + +#[tokio::test] +async fn test_jwt_valid_token_with_admin_role() { + let (addr, _server) = start_mock_jwks_server().await; + + let jwt_config = JwtConfig::new("https://test-issuer.example.com", "test-gateway") + .with_jwks_uri(format!( + "http://127.0.0.1:{}/.well-known/jwks.json", + addr.port() + )) + .with_role_mapping("admin", Role::Admin); + + let config = ControlPlaneAuthConfig { + jwt: Some(jwt_config), + api_keys: vec![], + audit_enabled: false, + }; + + let state = ControlPlaneAuthState::from_config(config) + .await + .expect("Failed to create auth state"); + + // Create a valid token with admin role + let claims = create_claims("admin@example.com", vec!["admin"]); + let token = create_test_token(&claims); + + // Validate the token + let jwt_validator = state.jwt_validator.as_ref().expect("JWT validator not set"); + let result = jwt_validator.validate(&token).await; + + assert!( + result.is_ok(), + "Token validation failed: {:?}", + result.err() + ); + let validated = result.unwrap(); + assert_eq!(validated.subject, "admin@example.com"); + assert_eq!(validated.role, Role::Admin); +} + +#[tokio::test] +async fn test_jwt_valid_token_with_user_role() { + let (addr, _server) = start_mock_jwks_server().await; + + let jwt_config = JwtConfig::new("https://test-issuer.example.com", "test-gateway") + .with_jwks_uri(format!( + "http://127.0.0.1:{}/.well-known/jwks.json", + addr.port() + )) + .with_role_mapping("admin", Role::Admin); + + let config = ControlPlaneAuthConfig { + jwt: Some(jwt_config), + api_keys: vec![], + audit_enabled: false, + }; + + let state = ControlPlaneAuthState::from_config(config) + .await + .expect("Failed to create auth state"); + + // Create a valid token with user role (not admin) + let claims = create_claims("user@example.com", vec!["user", "viewer"]); + let token = create_test_token(&claims); + + let jwt_validator = state.jwt_validator.as_ref().expect("JWT validator not set"); + let result = jwt_validator.validate(&token).await; + + assert!(result.is_ok()); + let validated = result.unwrap(); + assert_eq!(validated.subject, "user@example.com"); + assert_eq!(validated.role, Role::User); // Not admin +} + +#[tokio::test] +async fn test_jwt_expired_token() { + let (addr, _server) = start_mock_jwks_server().await; + + let jwt_config = JwtConfig::new("https://test-issuer.example.com", "test-gateway") + .with_jwks_uri(format!( + "http://127.0.0.1:{}/.well-known/jwks.json", + addr.port() + )); + + let config = ControlPlaneAuthConfig { + jwt: Some(jwt_config), + api_keys: vec![], + audit_enabled: false, + }; + + let state = ControlPlaneAuthState::from_config(config) + .await + .expect("Failed to create auth state"); + + // Create an expired token + let now = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_secs(); + + let claims = TestClaims { + sub: "user@example.com".to_string(), + iss: "https://test-issuer.example.com".to_string(), + aud: "test-gateway".to_string(), + exp: now - 3600, // Expired 1 hour ago + iat: now - 7200, + jti: None, + name: None, + roles: Some(vec!["admin".to_string()]), + }; + let token = create_test_token(&claims); + + let jwt_validator = state.jwt_validator.as_ref().expect("JWT validator not set"); + let result = jwt_validator.validate(&token).await; + + assert!(result.is_err(), "Expired token should fail validation"); + let err = result.unwrap_err(); + assert!( + err.to_string().contains("exp") || err.to_string().contains("ExpiredSignature"), + "Error should mention expiration: {}", + err + ); +} + +#[tokio::test] +async fn test_jwt_wrong_audience() { + let (addr, _server) = start_mock_jwks_server().await; + + let jwt_config = + JwtConfig::new("https://test-issuer.example.com", "correct-audience").with_jwks_uri( + format!("http://127.0.0.1:{}/.well-known/jwks.json", addr.port()), + ); + + let config = ControlPlaneAuthConfig { + jwt: Some(jwt_config), + api_keys: vec![], + audit_enabled: false, + }; + + let state = ControlPlaneAuthState::from_config(config) + .await + .expect("Failed to create auth state"); + + // Create a token with wrong audience + let now = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_secs(); + + let claims = TestClaims { + sub: "user@example.com".to_string(), + iss: "https://test-issuer.example.com".to_string(), + aud: "wrong-audience".to_string(), // Wrong audience + exp: now + 3600, + iat: now, + jti: None, + name: None, + roles: Some(vec!["admin".to_string()]), + }; + let token = create_test_token(&claims); + + let jwt_validator = state.jwt_validator.as_ref().expect("JWT validator not set"); + let result = jwt_validator.validate(&token).await; + + assert!(result.is_err(), "Wrong audience should fail validation"); +} + +#[tokio::test] +async fn test_jwt_wrong_issuer() { + let (addr, _server) = start_mock_jwks_server().await; + + let jwt_config = + JwtConfig::new("https://correct-issuer.example.com", "test-gateway").with_jwks_uri( + format!("http://127.0.0.1:{}/.well-known/jwks.json", addr.port()), + ); + + let config = ControlPlaneAuthConfig { + jwt: Some(jwt_config), + api_keys: vec![], + audit_enabled: false, + }; + + let state = ControlPlaneAuthState::from_config(config) + .await + .expect("Failed to create auth state"); + + // Create a token with wrong issuer + let claims = create_claims("user@example.com", vec!["admin"]); + // claims has issuer "https://test-issuer.example.com" which doesn't match + let token = create_test_token(&claims); + + let jwt_validator = state.jwt_validator.as_ref().expect("JWT validator not set"); + let result = jwt_validator.validate(&token).await; + + assert!(result.is_err(), "Wrong issuer should fail validation"); +} + +// ============================================================================ +// API Key Authentication Tests +// ============================================================================ + +#[tokio::test] +async fn test_api_key_valid_admin() { + let config = ControlPlaneAuthConfig { + jwt: None, + api_keys: vec![ApiKeyEntry::new( + "admin-key-1", + "Admin Service Account", + "sk-test-admin-key-12345", + Role::Admin, + )], + audit_enabled: false, + }; + + // Verify we can find the API key + let found = config.find_api_key("sk-test-admin-key-12345"); + assert!(found.is_some(), "API key should be found"); + + let entry = found.unwrap(); + assert_eq!(entry.id, "admin-key-1"); + assert_eq!(entry.role, Role::Admin); +} + +#[tokio::test] +async fn test_api_key_valid_user() { + let config = ControlPlaneAuthConfig { + jwt: None, + api_keys: vec![ApiKeyEntry::new( + "user-key-1", + "User Service Account", + "sk-test-user-key-12345", + Role::User, + )], + audit_enabled: false, + }; + + let found = config.find_api_key("sk-test-user-key-12345"); + assert!(found.is_some()); + + let entry = found.unwrap(); + assert_eq!(entry.role, Role::User); +} + +#[tokio::test] +async fn test_api_key_invalid() { + let config = ControlPlaneAuthConfig { + jwt: None, + api_keys: vec![ApiKeyEntry::new( + "admin-key-1", + "Admin Service Account", + "sk-correct-key", + Role::Admin, + )], + audit_enabled: false, + }; + + // Try with wrong key + let found = config.find_api_key("sk-wrong-key"); + assert!(found.is_none(), "Wrong API key should not be found"); +} + +#[tokio::test] +async fn test_api_key_timing_attack_resistance() { + // Verify that key comparison uses constant-time comparison + let config = ControlPlaneAuthConfig { + jwt: None, + api_keys: vec![ApiKeyEntry::new( + "key-1", + "Test Key", + "sk-abcdefghijklmnop", + Role::Admin, + )], + audit_enabled: false, + }; + + // These should all take roughly the same time due to constant-time comparison + // (We can't easily test timing in unit tests, but we verify the function works) + assert!(config.find_api_key("sk-abcdefghijklmnop").is_some()); + assert!(config.find_api_key("sk-abcdefghijklmnox").is_none()); // One char different + assert!(config.find_api_key("sk-xxxxxxxxxxxxxxxx").is_none()); // All different + assert!(config.find_api_key("sk-a").is_none()); // Much shorter +} + +// ============================================================================ +// Combined Auth Tests (JWT + API Key fallback) +// ============================================================================ + +#[tokio::test] +async fn test_combined_auth_jwt_and_api_keys() { + let (addr, _server) = start_mock_jwks_server().await; + + let jwt_config = JwtConfig::new("https://test-issuer.example.com", "test-gateway") + .with_jwks_uri(format!( + "http://127.0.0.1:{}/.well-known/jwks.json", + addr.port() + )) + .with_role_mapping("admin", Role::Admin); + + let config = ControlPlaneAuthConfig { + jwt: Some(jwt_config), + api_keys: vec![ApiKeyEntry::new( + "backup-key", + "Backup Admin Key", + "sk-backup-admin-key", + Role::Admin, + )], + audit_enabled: true, + }; + + let state = ControlPlaneAuthState::from_config(config.clone()) + .await + .expect("Failed to create auth state"); + + // Both JWT and API key should work + assert!(state.jwt_validator.is_some()); + assert!(config.find_api_key("sk-backup-admin-key").is_some()); + + // Verify JWT works + let claims = create_claims("jwt-user@example.com", vec!["admin"]); + let token = create_test_token(&claims); + let result = state.jwt_validator.as_ref().unwrap().validate(&token).await; + assert!(result.is_ok()); +} + +// ============================================================================ +// Role-Based Access Control Tests +// ============================================================================ + +#[tokio::test] +async fn test_role_admin_has_access() { + assert!(Role::Admin.is_admin()); +} + +#[tokio::test] +async fn test_role_user_no_admin_access() { + assert!(!Role::User.is_admin()); +} + +#[tokio::test] +async fn test_role_parsing() { + assert_eq!("admin".parse::().unwrap(), Role::Admin); + assert_eq!("Admin".parse::().unwrap(), Role::Admin); + assert_eq!("ADMIN".parse::().unwrap(), Role::Admin); + assert_eq!("user".parse::().unwrap(), Role::User); + assert_eq!("User".parse::().unwrap(), Role::User); + assert_eq!("USER".parse::().unwrap(), Role::User); +} + +// ============================================================================ +// try_init Helper Tests +// ============================================================================ + +#[tokio::test] +async fn test_try_init_with_no_config() { + let result = ControlPlaneAuthState::try_init(None).await; + assert!(result.is_none()); +} + +#[tokio::test] +async fn test_try_init_with_disabled_config() { + let config = ControlPlaneAuthConfig { + jwt: None, + api_keys: vec![], // No auth configured = disabled + audit_enabled: false, + }; + + let result = ControlPlaneAuthState::try_init(Some(&config)).await; + assert!(result.is_none()); +} + +#[tokio::test] +async fn test_try_init_with_api_keys_only() { + let config = ControlPlaneAuthConfig { + jwt: None, + api_keys: vec![ApiKeyEntry::new("key-1", "Test", "sk-test", Role::Admin)], + audit_enabled: false, + }; + + let result = ControlPlaneAuthState::try_init(Some(&config)).await; + assert!(result.is_some()); +} + +// ============================================================================ +// Audit Logging Tests +// ============================================================================ + +#[tokio::test] +async fn test_audit_logging_enabled() { + let config = ControlPlaneAuthConfig { + jwt: None, + api_keys: vec![ApiKeyEntry::new("key-1", "Test", "sk-test", Role::Admin)], + audit_enabled: true, + }; + + let state = ControlPlaneAuthState::from_config(config) + .await + .expect("Failed to create auth state"); + + assert!(state.audit_logger.is_enabled()); +} + +#[tokio::test] +async fn test_audit_logging_disabled() { + let config = ControlPlaneAuthConfig { + jwt: None, + api_keys: vec![ApiKeyEntry::new("key-1", "Test", "sk-test", Role::Admin)], + audit_enabled: false, + }; + + let state = ControlPlaneAuthState::from_config(config) + .await + .expect("Failed to create auth state"); + + assert!(!state.audit_logger.is_enabled()); +} + +// ============================================================================ +// JTI Replay Protection Tests +// ============================================================================ + +#[tokio::test] +async fn test_jwt_jti_replay_protection() { + use sgl_model_gateway::auth::JwtValidator; + + let (addr, _server) = start_mock_jwks_server().await; + + let jwt_config = JwtConfig::new("https://test-issuer.example.com", "test-gateway") + .with_jwks_uri(format!( + "http://127.0.0.1:{}/.well-known/jwks.json", + addr.port() + )) + .with_role_mapping("admin", Role::Admin); + + // Create validator with JTI replay protection enabled + let validator = JwtValidator::from_config_with_options(jwt_config, true) + .await + .expect("Failed to create JWT validator"); + + // Create a token with a specific JTI + let claims = create_claims("user@example.com", vec!["admin"]); + let token = create_test_token(&claims); + + // First validation should succeed + let result1 = validator.validate(&token).await; + assert!(result1.is_ok(), "First validation should succeed"); + + // Second validation with same token should fail (replay) + let result2 = validator.validate(&token).await; + assert!( + result2.is_err(), + "Second validation should fail (replay detected)" + ); + let err = result2.unwrap_err(); + assert!( + err.to_string().contains("replay") || err.to_string().contains("already been used"), + "Error should mention replay: {}", + err + ); +} + +#[tokio::test] +async fn test_jwt_different_tokens_no_replay() { + use sgl_model_gateway::auth::JwtValidator; + + let (addr, _server) = start_mock_jwks_server().await; + + let jwt_config = JwtConfig::new("https://test-issuer.example.com", "test-gateway") + .with_jwks_uri(format!( + "http://127.0.0.1:{}/.well-known/jwks.json", + addr.port() + )) + .with_role_mapping("admin", Role::Admin); + + // Create validator with JTI replay protection enabled + let validator = JwtValidator::from_config_with_options(jwt_config, true) + .await + .expect("Failed to create JWT validator"); + + // Create two different tokens (different JTIs) + let claims1 = create_claims("user1@example.com", vec!["admin"]); + let claims2 = create_claims("user2@example.com", vec!["admin"]); + let token1 = create_test_token(&claims1); + let token2 = create_test_token(&claims2); + + // Both should succeed since they have different JTIs + let result1 = validator.validate(&token1).await; + assert!(result1.is_ok(), "First token validation should succeed"); + + let result2 = validator.validate(&token2).await; + assert!( + result2.is_ok(), + "Second token validation should succeed (different JTI)" + ); +} + +// ============================================================================ +// Malformed Token Tests +// ============================================================================ + +#[tokio::test] +async fn test_jwt_malformed_token() { + let (addr, _server) = start_mock_jwks_server().await; + + let jwt_config = JwtConfig::new("https://test-issuer.example.com", "test-gateway") + .with_jwks_uri(format!( + "http://127.0.0.1:{}/.well-known/jwks.json", + addr.port() + )); + + let config = ControlPlaneAuthConfig { + jwt: Some(jwt_config), + api_keys: vec![], + audit_enabled: false, + }; + + let state = ControlPlaneAuthState::from_config(config) + .await + .expect("Failed to create auth state"); + + let jwt_validator = state.jwt_validator.as_ref().expect("JWT validator not set"); + + // Test various malformed tokens + let malformed_tokens = [ + "not-a-jwt", + "only.two.parts.is.not.enough", + "", + "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9", // Only header + "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.", // Header with dot + ".....", + ]; + + for token in malformed_tokens { + let result = jwt_validator.validate(token).await; + assert!( + result.is_err(), + "Malformed token '{}' should fail validation", + token + ); + } +} + +#[tokio::test] +async fn test_jwt_missing_kid_in_header() { + let (addr, _server) = start_mock_jwks_server().await; + + let jwt_config = JwtConfig::new("https://test-issuer.example.com", "test-gateway") + .with_jwks_uri(format!( + "http://127.0.0.1:{}/.well-known/jwks.json", + addr.port() + )); + + let config = ControlPlaneAuthConfig { + jwt: Some(jwt_config), + api_keys: vec![], + audit_enabled: false, + }; + + let state = ControlPlaneAuthState::from_config(config) + .await + .expect("Failed to create auth state"); + + // Create a token WITHOUT kid in header + let claims = create_claims("user@example.com", vec!["admin"]); + let mut header = Header::new(jsonwebtoken::Algorithm::RS256); + header.kid = None; // Explicitly no kid + + let key = EncodingKey::from_rsa_pem(TEST_KEYS.private_key_pem.as_bytes()) + .expect("Failed to create encoding key"); + let token_without_kid = encode(&header, &claims, &key).expect("Failed to encode token"); + + let jwt_validator = state.jwt_validator.as_ref().expect("JWT validator not set"); + let result = jwt_validator.validate(&token_without_kid).await; + + assert!(result.is_err(), "Token without kid should fail validation"); + let err = result.unwrap_err(); + assert!( + err.to_string().contains("kid") || err.to_string().contains("Missing"), + "Error should mention missing kid: {}", + err + ); +} + +// ============================================================================ +// Edge Cases and Security Tests +// ============================================================================ + +#[tokio::test] +async fn test_jwt_role_extraction_from_groups_claim() { + let (addr, _server) = start_mock_jwks_server().await; + + // JWT validator checks alternate claims ("groups", "roles", etc.) as fallback + let jwt_config = JwtConfig::new("https://test-issuer.example.com", "test-gateway") + .with_jwks_uri(format!( + "http://127.0.0.1:{}/.well-known/jwks.json", + addr.port() + )) + .with_role_mapping("administrators", Role::Admin); + + let config = ControlPlaneAuthConfig { + jwt: Some(jwt_config), + api_keys: vec![], + audit_enabled: false, + }; + + let state = ControlPlaneAuthState::from_config(config) + .await + .expect("Failed to create auth state"); + + // Create token with 'groups' claim instead of 'roles' + let now = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_secs(); + + #[derive(Serialize)] + struct GroupsClaims { + sub: String, + iss: String, + aud: String, + exp: u64, + iat: u64, + jti: String, + groups: Vec, // Use 'groups' instead of 'roles' + } + + let claims = GroupsClaims { + sub: "group-user@example.com".to_string(), + iss: "https://test-issuer.example.com".to_string(), + aud: "test-gateway".to_string(), + exp: now + 3600, + iat: now, + jti: uuid::Uuid::new_v4().to_string(), + groups: vec!["administrators".to_string(), "developers".to_string()], + }; + + let mut header = Header::new(jsonwebtoken::Algorithm::RS256); + header.kid = Some(TEST_KEY_ID.to_string()); + let key = EncodingKey::from_rsa_pem(TEST_KEYS.private_key_pem.as_bytes()) + .expect("Failed to create encoding key"); + let token = encode(&header, &claims, &key).expect("Failed to encode token"); + + let jwt_validator = state.jwt_validator.as_ref().expect("JWT validator not set"); + let result = jwt_validator.validate(&token).await; + + assert!( + result.is_ok(), + "Token with groups claim should validate: {:?}", + result.err() + ); + let validated = result.unwrap(); + assert_eq!( + validated.role, + Role::Admin, + "Should map 'administrators' to Admin role" + ); +} + +#[tokio::test] +async fn test_jwt_no_role_defaults_to_user() { + let (addr, _server) = start_mock_jwks_server().await; + + let jwt_config = JwtConfig::new("https://test-issuer.example.com", "test-gateway") + .with_jwks_uri(format!( + "http://127.0.0.1:{}/.well-known/jwks.json", + addr.port() + )); + // No role mapping configured + + let config = ControlPlaneAuthConfig { + jwt: Some(jwt_config), + api_keys: vec![], + audit_enabled: false, + }; + + let state = ControlPlaneAuthState::from_config(config) + .await + .expect("Failed to create auth state"); + + // Create token without any role claims + let now = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_secs(); + + #[derive(Serialize)] + struct MinimalClaims { + sub: String, + iss: String, + aud: String, + exp: u64, + iat: u64, + jti: String, + // No roles, groups, or role claims + } + + let claims = MinimalClaims { + sub: "minimal-user@example.com".to_string(), + iss: "https://test-issuer.example.com".to_string(), + aud: "test-gateway".to_string(), + exp: now + 3600, + iat: now, + jti: uuid::Uuid::new_v4().to_string(), + }; + + let mut header = Header::new(jsonwebtoken::Algorithm::RS256); + header.kid = Some(TEST_KEY_ID.to_string()); + let key = EncodingKey::from_rsa_pem(TEST_KEYS.private_key_pem.as_bytes()) + .expect("Failed to create encoding key"); + let token = encode(&header, &claims, &key).expect("Failed to encode token"); + + let jwt_validator = state.jwt_validator.as_ref().expect("JWT validator not set"); + let result = jwt_validator.validate(&token).await; + + assert!(result.is_ok(), "Token without role should still validate"); + let validated = result.unwrap(); + assert_eq!( + validated.role, + Role::User, + "Should default to User role when no role found" + ); +} + +#[tokio::test] +async fn test_multiple_api_keys() { + let config = ControlPlaneAuthConfig { + jwt: None, + api_keys: vec![ + ApiKeyEntry::new("admin-1", "Admin Key 1", "sk-admin-key-1", Role::Admin), + ApiKeyEntry::new("admin-2", "Admin Key 2", "sk-admin-key-2", Role::Admin), + ApiKeyEntry::new("user-1", "User Key 1", "sk-user-key-1", Role::User), + ], + audit_enabled: false, + }; + + // All keys should be findable + assert!(config.find_api_key("sk-admin-key-1").is_some()); + assert!(config.find_api_key("sk-admin-key-2").is_some()); + assert!(config.find_api_key("sk-user-key-1").is_some()); + + // Verify correct roles + assert_eq!( + config.find_api_key("sk-admin-key-1").unwrap().role, + Role::Admin + ); + assert_eq!( + config.find_api_key("sk-user-key-1").unwrap().role, + Role::User + ); + + // Non-existent key should not be found + assert!(config.find_api_key("sk-nonexistent").is_none()); +} + +#[tokio::test] +async fn test_config_is_enabled_checks() { + // Empty config is not enabled + let empty_config = ControlPlaneAuthConfig { + jwt: None, + api_keys: vec![], + audit_enabled: false, + }; + assert!(!empty_config.is_enabled()); + + // Config with only API keys is enabled + let api_key_config = ControlPlaneAuthConfig { + jwt: None, + api_keys: vec![ApiKeyEntry::new("key-1", "Test", "sk-test", Role::Admin)], + audit_enabled: false, + }; + assert!(api_key_config.is_enabled()); + + // Config with only JWT is enabled + let jwt_config = ControlPlaneAuthConfig { + jwt: Some(JwtConfig::new("https://issuer.example.com", "audience")), + api_keys: vec![], + audit_enabled: false, + }; + assert!(jwt_config.is_enabled()); + + // Config with both is enabled + let full_config = ControlPlaneAuthConfig { + jwt: Some(JwtConfig::new("https://issuer.example.com", "audience")), + api_keys: vec![ApiKeyEntry::new("key-1", "Test", "sk-test", Role::Admin)], + audit_enabled: true, + }; + assert!(full_config.is_enabled()); +} diff --git a/sgl-model-gateway/tests/common/test_app.rs b/sgl-model-gateway/tests/common/test_app.rs index 00acdbd09..995cb0c97 100644 --- a/sgl-model-gateway/tests/common/test_app.rs +++ b/sgl-model-gateway/tests/common/test_app.rs @@ -110,6 +110,7 @@ pub fn create_test_app( build_app( app_state, auth_config, + None, // No control plane auth for tests router_config.max_payload_size, request_id_headers, router_config.cors_allowed_origins.clone(), @@ -152,6 +153,7 @@ pub fn create_test_app_with_context( build_app( app_state, auth_config, + None, // No control plane auth for tests router_config.max_payload_size, request_id_headers, router_config.cors_allowed_origins.clone(), diff --git a/sgl-model-gateway/tests/wasm_test.rs b/sgl-model-gateway/tests/wasm_test.rs index cc965ceb1..25dd90e04 100644 --- a/sgl-model-gateway/tests/wasm_test.rs +++ b/sgl-model-gateway/tests/wasm_test.rs @@ -207,6 +207,7 @@ async fn create_test_app_with_wasm() -> (axum::Router, Arc, TempDir) let app = build_app( app_state, sgl_model_gateway::middleware::AuthConfig { api_key: None }, + None, // No control plane auth for tests 256 * 1024 * 1024, request_id_headers, vec![], // cors_allowed_origins