chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[smg] use already published auth crate for better compilation speed (#17723)

Browse Source
This commit is contained in:
Simo Lin
2026-01-25 17:52:42 -05:00
committed by GitHub
parent 344eeaee90
commit bf139af37a
8 changed files with 2 additions and 2184 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
sgl-model-gateway/Cargo.toml
View File

@@ -71,7 +71,6 @@ metrics-exporter-prometheus = "0.17.0"
uuid = { version = "1.10", features = ["v4", "serde"] }
ulid = "1.2.1"
parking_lot = "0.12.4"
rayon = "1.10"
thiserror = "2.0.12"
@@ -81,19 +80,14 @@ url = "2.5.4"
validator = { version = "0.20.0", features = ["derive"] }
tokio-stream = { version = "0.1", features = ["sync"] }
anyhow = "1.0"
tokenizers = { version = "0.22.0" }
tiktoken-rs = { version = "0.7.0" }
reasoning-parser = "1.0.0"
openai-protocol = { version = "1.0.0", features = ["axum"] }
tool-parser = "1.0.0"
llm-tokenizer = "1.0.0"
minijinja = { version = "2.0", features = ["unstable_machinery", "json", "builtins"] }
minijinja-contrib = { version = "2.0", features = ["pycompat"] }
smg-auth = "1.0.0"
rustls = { version = "0.23", default-features = false, features = ["ring", "std"] }
tokio-rustls = { version = "0.26", default-features = false, features = ["ring", "logging", "tls12"] }
rustls-pemfile = "2.2"
openssl = "0.10.73"
hf-hub = { version = "0.4.3", features = ["tokio"] }
rmcp = { version = "0.8.3", features = ["client", "server",
"transport-child-process",
"transport-sse-client-reqwest",
@@ -106,7 +100,6 @@ 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"
num-bigint = "0.4"
image = { version = "0.25.4", default-features = false, features = ["png", "jpeg", "gif", "bmp", "ico", "tiff", "webp"] }
@@ -143,7 +136,6 @@ aho-corasick = "1.1.4"
[build-dependencies]
tonic-prost-build = "0.14.2"
prost-build = "0.14.1"
chrono = { version = "0.4", features = ["clock"] }
toml = "0.9"

357
sgl-model-gateway/src/auth/audit.rs
View File

@@ -1,357 +0,0 @@
//! 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: "smg::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"));
}
}

345
sgl-model-gateway/src/auth/config.rs
View File

@@ -1,345 +0,0 @@
//! 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"));
}
}

517
sgl-model-gateway/src/auth/jwks.rs
View File

@@ -1,517 +0,0 @@
//! 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(crate) 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(crate) 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.
#[allow(dead_code)]
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.
#[allow(dead_code)]
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))));
}
}

536
sgl-model-gateway/src/auth/jwt.rs
View File

@@ -1,536 +0,0 @@
//! 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(crate) 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(crate) enum Audience {
Single(String),
Multiple(Vec<String>),
#[default]
None,
}
impl Audience {
#[allow(dead_code)]
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
#[allow(dead_code)]
pub(crate) 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.
#[allow(dead_code)]
pub(crate) 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(crate) 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.
#[allow(dead_code)]
pub(crate) 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"));
}
}

399
sgl-model-gateway/src/auth/middleware.rs
View File

@@ -1,399 +0,0 @@
//! 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());
}
}

20
sgl-model-gateway/src/auth/mod.rs
View File

@@ -1,20 +0,0 @@
//! 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,
};

2
sgl-model-gateway/src/lib.rs
View File

@@ -1,5 +1,5 @@
pub mod app_context;
pub mod auth;
pub use smg_auth as auth;
pub mod config;
pub mod core;
pub mod data_connector;