[model-gateway] add JWT/OIDC authentication for control plane APIs (#15850)

sgl-model-gateway/bindings/python/src/lib.rs

@@ -31,6 +31,152 @@ pub enum HistoryBackendType {
     Postgres,
 }
 
+#[pyclass(eq)]
+#[derive(Clone, PartialEq, Debug, Default)]
+pub enum PyRole {
+    Admin,
+    #[default]
+    User,
+}
+
+impl PyRole {
+    pub fn to_auth_role(&self) -> auth::Role {
+        match self {
+            PyRole::Admin => auth::Role::Admin,
+            PyRole::User => auth::Role::User,
+        }
+    }
+}
+
+#[pyclass]
+#[derive(Clone, Debug, PartialEq)]
+pub struct PyApiKeyEntry {
+    #[pyo3(get, set)]
+    pub id: String,
+    #[pyo3(get, set)]
+    pub name: String,
+    #[pyo3(get, set)]
+    pub key: String,
+    #[pyo3(get, set)]
+    pub role: PyRole,
+}
+
+#[pymethods]
+impl PyApiKeyEntry {
+    #[new]
+    #[pyo3(signature = (id, name, key, role = PyRole::User))]
+    fn new(id: String, name: String, key: String, role: PyRole) -> Self {
+        PyApiKeyEntry { id, name, key, role }
+    }
+}
+
+impl PyApiKeyEntry {
+    pub fn to_auth_api_key_entry(&self) -> auth::ApiKeyEntry {
+        auth::ApiKeyEntry::new(&self.id, &self.name, &self.key, self.role.to_auth_role())
+    }
+}
+
+#[pyclass]
+#[derive(Clone, Debug, PartialEq)]
+pub struct PyJwtConfig {
+    #[pyo3(get, set)]
+    pub issuer: String,
+    #[pyo3(get, set)]
+    pub audience: String,
+    #[pyo3(get, set)]
+    pub jwks_uri: Option<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>()?;