[model-gateway] add JWT/OIDC authentication for control plane APIs (#15850)
This commit is contained in:
Simo Lin
@@ -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]]
|
||||
|
||||
@@ -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 <jwt-token>`
|
||||
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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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: <ip>:<port>",
|
||||
)
|
||||
# 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
|
||||
|
||||
@@ -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<String>,
|
||||
#[pyo3(get, set)]
|
||||
pub role_mapping: HashMap<String, String>,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PyJwtConfig {
|
||||
#[new]
|
||||
#[pyo3(signature = (
|
||||
issuer,
|
||||
audience,
|
||||
jwks_uri = None,
|
||||
role_mapping = HashMap::new(),
|
||||
))]
|
||||
fn new(
|
||||
issuer: String,
|
||||
audience: String,
|
||||
jwks_uri: Option<String>,
|
||||
role_mapping: HashMap<String, String>,
|
||||
) -> 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<PyJwtConfig>,
|
||||
#[pyo3(get, set)]
|
||||
pub api_keys: Vec<PyApiKeyEntry>,
|
||||
#[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<PyJwtConfig>,
|
||||
api_keys: Vec<PyApiKeyEntry>,
|
||||
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<String>,
|
||||
enable_trace: bool,
|
||||
otlp_traces_endpoint: String,
|
||||
control_plane_auth: Option<PyControlPlaneAuthConfig>,
|
||||
}
|
||||
|
||||
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<String>,
|
||||
enable_trace: bool,
|
||||
otlp_traces_endpoint: String,
|
||||
control_plane_auth: Option<PyControlPlaneAuthConfig>,
|
||||
) -> PyResult<Self> {
|
||||
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::<PolicyType>()?;
|
||||
m.add_class::<BackendType>()?;
|
||||
m.add_class::<HistoryBackendType>()?;
|
||||
m.add_class::<PyRole>()?;
|
||||
m.add_class::<PyApiKeyEntry>()?;
|
||||
m.add_class::<PyJwtConfig>()?;
|
||||
m.add_class::<PyControlPlaneAuthConfig>()?;
|
||||
m.add_class::<PyOracleConfig>()?;
|
||||
m.add_class::<PyPostgresConfig>()?;
|
||||
m.add_class::<Router>()?;
|
||||
|
||||
@@ -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<Utc>,
|
||||
|
||||
/// 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<String>,
|
||||
|
||||
/// Operation outcome
|
||||
pub outcome: AuditOutcome,
|
||||
|
||||
/// Request ID for correlation
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub request_id: Option<String>,
|
||||
|
||||
/// Additional details or error message
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub details: Option<String>,
|
||||
}
|
||||
|
||||
/// 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"));
|
||||
}
|
||||
}
|
||||
@@ -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<Self, Self::Err> {
|
||||
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<String>,
|
||||
|
||||
/// 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<String, Role>,
|
||||
|
||||
/// 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<String>, audience: impl Into<String>) -> 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<String>) -> 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<String>, 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<String>,
|
||||
name: impl Into<String>,
|
||||
key: impl Into<String>,
|
||||
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<JwtConfig>,
|
||||
|
||||
/// 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<ApiKeyEntry>,
|
||||
|
||||
/// 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::<Role>().unwrap(), Role::Admin);
|
||||
assert_eq!("ADMIN".parse::<Role>().unwrap(), Role::Admin);
|
||||
assert_eq!("user".parse::<Role>().unwrap(), Role::User);
|
||||
assert_eq!("USER".parse::<Role>().unwrap(), Role::User);
|
||||
assert!("invalid".parse::<Role>().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"));
|
||||
}
|
||||
}
|
||||
@@ -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<Url, JwksError> {
|
||||
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<Option<CachedJwks>>,
|
||||
/// 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<String>, ttl: Duration) -> Result<Self, JwksError> {
|
||||
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<Self, JwksError> {
|
||||
// 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<JwkSet, JwksError> {
|
||||
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<JwkSet, JwksError> {
|
||||
// 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<Jwk, JwksError> {
|
||||
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))));
|
||||
}
|
||||
}
|
||||
@@ -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<String>,
|
||||
|
||||
/// Issuer
|
||||
pub iss: Option<String>,
|
||||
|
||||
/// Audience (can be string or array)
|
||||
#[serde(default)]
|
||||
pub aud: Audience,
|
||||
|
||||
/// Expiration time
|
||||
pub exp: Option<u64>,
|
||||
|
||||
/// Issued at
|
||||
pub iat: Option<u64>,
|
||||
|
||||
/// Not before
|
||||
pub nbf: Option<u64>,
|
||||
|
||||
/// JWT ID
|
||||
pub jti: Option<String>,
|
||||
|
||||
/// Email (common claim)
|
||||
pub email: Option<String>,
|
||||
|
||||
/// Name (common claim)
|
||||
pub name: Option<String>,
|
||||
|
||||
/// Preferred username (OIDC claim)
|
||||
pub preferred_username: Option<String>,
|
||||
|
||||
/// All other claims
|
||||
#[serde(flatten)]
|
||||
pub extra: HashMap<String, serde_json::Value>,
|
||||
}
|
||||
|
||||
/// 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<String>),
|
||||
#[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<String>,
|
||||
|
||||
/// Display name if present
|
||||
pub name: Option<String>,
|
||||
|
||||
/// 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<JwksProvider>,
|
||||
|
||||
/// JWT configuration
|
||||
config: JwtConfig,
|
||||
|
||||
/// Pre-configured validation settings
|
||||
validation: Validation,
|
||||
|
||||
/// JTI cache for replay protection (optional)
|
||||
/// Maps JTI -> expiration time
|
||||
jti_cache: Option<Mutex<LruCache<String, JtiCacheEntry>>>,
|
||||
|
||||
/// 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<JwksProvider>) -> 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<JwksProvider>,
|
||||
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, JwtValidatorError> {
|
||||
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<Self, JwtValidatorError> {
|
||||
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<ValidatedToken, JwtValidatorError> {
|
||||
// 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<StandardClaims> = 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<Role, JwtValidatorError> {
|
||||
// Try to get the role claim value
|
||||
let role_value = claims.extra.get(&self.config.role_claim);
|
||||
|
||||
let role_strings: Vec<String> = 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::<Role>() {
|
||||
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<DecodingKey, JwtValidatorError> {
|
||||
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<Algorithm, JwtValidatorError> {
|
||||
// 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<JwksProvider> {
|
||||
&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"));
|
||||
}
|
||||
}
|
||||
@@ -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<String>,
|
||||
|
||||
/// 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<B> PrincipalExt for http::Request<B> {
|
||||
fn principal(&self) -> Option<&Principal> {
|
||||
self.extensions().get::<Principal>()
|
||||
}
|
||||
}
|
||||
|
||||
/// 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<Arc<JwtValidator>>,
|
||||
|
||||
/// Audit logger
|
||||
pub audit_logger: AuditLogger,
|
||||
}
|
||||
|
||||
impl ControlPlaneAuthState {
|
||||
/// Create a new control plane auth state.
|
||||
pub fn new(config: ControlPlaneAuthConfig, jwt_validator: Option<Arc<JwtValidator>>) -> 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<Self, Box<dyn std::error::Error + Send + Sync>> {
|
||||
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<Self> {
|
||||
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<Response> {
|
||||
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<ControlPlaneAuthState>,
|
||||
mut request: Request<Body>,
|
||||
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::<RequestId>().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());
|
||||
}
|
||||
}
|
||||
@@ -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,
|
||||
};
|
||||
@@ -1,4 +1,5 @@
|
||||
pub mod app_context;
|
||||
pub mod auth;
|
||||
pub mod config;
|
||||
pub mod core;
|
||||
pub mod data_connector;
|
||||
|
||||
@@ -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<String>,
|
||||
|
||||
// ==================== 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<String>,
|
||||
|
||||
/// Expected JWT audience claim (usually the client ID or API identifier)
|
||||
/// Example: api://sgl-gateway
|
||||
#[arg(long, env = "JWT_AUDIENCE")]
|
||||
jwt_audience: Option<String>,
|
||||
|
||||
/// Explicit JWKS URI (if not provided, discovered from issuer)
|
||||
#[arg(long, env = "JWT_JWKS_URI")]
|
||||
jwt_jwks_uri: Option<String>,
|
||||
|
||||
/// 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<String>,
|
||||
|
||||
/// 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<String>,
|
||||
|
||||
/// 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<ApiKeyEntry> {
|
||||
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<String, Role> = 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<ApiKeyEntry> = 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,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -524,11 +524,14 @@ pub struct ServerConfig {
|
||||
pub request_timeout_secs: u64,
|
||||
pub request_id_headers: Option<Vec<String>>,
|
||||
pub shutdown_grace_period_secs: u64,
|
||||
/// Control plane authentication configuration
|
||||
pub control_plane_auth: Option<crate::auth::ControlPlaneAuthConfig>,
|
||||
}
|
||||
|
||||
pub fn build_app(
|
||||
app_state: Arc<AppState>,
|
||||
auth_config: AuthConfig,
|
||||
control_plane_auth_state: Option<crate::auth::ControlPlaneAuthState>,
|
||||
max_payload_size: usize,
|
||||
request_id_headers: Vec<String>,
|
||||
cors_allowed_origins: Vec<String>,
|
||||
@@ -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<Arc<AppState>>| {
|
||||
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<dyn std::error::Err
|
||||
api_key: config.router_config.api_key.clone(),
|
||||
};
|
||||
|
||||
// Initialize control plane authentication if configured
|
||||
let control_plane_auth_state =
|
||||
crate::auth::ControlPlaneAuthState::try_init(config.control_plane_auth.as_ref()).await;
|
||||
|
||||
let app = build_app(
|
||||
app_state,
|
||||
auth_config,
|
||||
control_plane_auth_state,
|
||||
config.max_payload_size,
|
||||
request_id_headers,
|
||||
config.router_config.cors_allowed_origins.clone(),
|
||||
|
||||
@@ -0,0 +1,939 @@
|
||||
//! Integration tests for control plane authentication.
|
||||
//!
|
||||
//! Tests the full authentication flow including:
|
||||
//! - JWT token validation with mock JWKS server
|
||||
//! - API key authentication
|
||||
//! - Role-based access control
|
||||
//! - Token expiration and replay protection
|
||||
|
||||
use std::{
|
||||
net::SocketAddr,
|
||||
sync::LazyLock,
|
||||
time::{Duration, SystemTime, UNIX_EPOCH},
|
||||
};
|
||||
|
||||
use axum::{routing::get, Json, Router};
|
||||
use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine};
|
||||
use jsonwebtoken::{encode, EncodingKey, Header};
|
||||
use rsa::{traits::PublicKeyParts, RsaPrivateKey};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::json;
|
||||
use sgl_model_gateway::auth::{
|
||||
ApiKeyEntry, ControlPlaneAuthConfig, ControlPlaneAuthState, JwtConfig, Role,
|
||||
};
|
||||
use tokio::net::TcpListener;
|
||||
|
||||
const TEST_KEY_ID: &str = "test-key-1";
|
||||
|
||||
/// Test RSA key pair - generates matching public/private key components
|
||||
struct TestKeyPair {
|
||||
private_key_pem: String,
|
||||
n_base64url: String,
|
||||
e_base64url: String,
|
||||
}
|
||||
|
||||
impl TestKeyPair {
|
||||
fn generate() -> 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<TestKeyPair> = 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<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
name: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
roles: Option<Vec<String>>,
|
||||
}
|
||||
|
||||
/// 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::<Role>().unwrap(), Role::Admin);
|
||||
assert_eq!("Admin".parse::<Role>().unwrap(), Role::Admin);
|
||||
assert_eq!("ADMIN".parse::<Role>().unwrap(), Role::Admin);
|
||||
assert_eq!("user".parse::<Role>().unwrap(), Role::User);
|
||||
assert_eq!("User".parse::<Role>().unwrap(), Role::User);
|
||||
assert_eq!("USER".parse::<Role>().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<String>, // 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());
|
||||
}
|
||||
@@ -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(),
|
||||
|
||||
@@ -207,6 +207,7 @@ async fn create_test_app_with_wasm() -> (axum::Router, Arc<AppContext>, 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
|
||||
|
||||
Reference in New Issue
Block a user