From 0e4108ba29f43eb1d2e5bdd40f2f9bd14bc4e59f Mon Sep 17 00:00:00 2001 From: ratish <114130421+Ratish1@users.noreply.github.com> Date: Mon, 15 Dec 2025 03:37:55 +0400 Subject: [PATCH] feat(gateway): Add server-side TLS support (#15052) --- sgl-model-gateway/Cargo.toml | 3 + .../bindings/python/sglang_router/router.py | 2 + .../python/sglang_router/router_args.py | 28 ++++ sgl-model-gateway/bindings/python/src/lib.rs | 12 ++ .../py_test/e2e_http/test_tls.py | 153 ++++++++++++++++++ sgl-model-gateway/src/config/builder.rs | 58 +++++++ sgl-model-gateway/src/config/types.rs | 8 + sgl-model-gateway/src/main.rs | 9 +- sgl-model-gateway/src/server.rs | 47 +++++- 9 files changed, 312 insertions(+), 8 deletions(-) create mode 100644 sgl-model-gateway/py_test/e2e_http/test_tls.py diff --git a/sgl-model-gateway/Cargo.toml b/sgl-model-gateway/Cargo.toml index fa0e32880..58348fe0d 100644 --- a/sgl-model-gateway/Cargo.toml +++ b/sgl-model-gateway/Cargo.toml @@ -32,6 +32,7 @@ path = "src/main.rs" [dependencies] clap = { version = "4", features = ["derive", "env"] } axum = { version = "0.8.4", features = ["macros", "ws", "tracing"] } +axum-server = { version = "0.7.3", default_features = false, features = ["tls-rustls"] } tower = { version = "0.5", features = ["full"] } tower-http = { version = "0.6", features = ["trace", "compression-gzip", "cors", "timeout", "limit", "request-id", "util"] } serde = { version = "1.0", features = ["derive"] } @@ -81,6 +82,8 @@ tiktoken-rs = { version = "0.7.0" } minijinja = { version = "2.0", features = ["unstable_machinery", "json", "builtins"] } minijinja-contrib = { version = "2.0", features = ["pycompat"] } 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", diff --git a/sgl-model-gateway/bindings/python/sglang_router/router.py b/sgl-model-gateway/bindings/python/sglang_router/router.py index 05506e1cd..4abe19d29 100644 --- a/sgl-model-gateway/bindings/python/sglang_router/router.py +++ b/sgl-model-gateway/bindings/python/sglang_router/router.py @@ -135,6 +135,8 @@ class Router: health_check_endpoint: Health check endpoint path. Default: '/health' model_path: Model path for loading tokenizer (HuggingFace model ID or local path). Default: None tokenizer_path: Explicit tokenizer path (overrides model_path tokenizer if provided). Default: None + server_cert_path: Path to server TLS certificate (PEM format). Default: None + server_key_path: Path to server TLS private key (PEM format). Default: None """ def __init__(self, router: _Router): diff --git a/sgl-model-gateway/bindings/python/sglang_router/router_args.py b/sgl-model-gateway/bindings/python/sglang_router/router_args.py index 17b32c5f0..8ac4b5d04 100644 --- a/sgl-model-gateway/bindings/python/sglang_router/router_args.py +++ b/sgl-model-gateway/bindings/python/sglang_router/router_args.py @@ -118,6 +118,9 @@ class RouterArgs: client_cert_path: Optional[str] = None client_key_path: Optional[str] = None ca_cert_paths: List[str] = dataclasses.field(default_factory=list) + # Server TLS configuration + server_cert_path: Optional[str] = None + server_key_path: Optional[str] = None # Trace enable_trace: bool = False otlp_traces_endpoint: str = "localhost:4317" @@ -644,6 +647,19 @@ class RouterArgs: default=[], help="Path(s) to CA certificate(s) for verifying worker TLS certificates. Can specify multiple CAs.", ) + # Server TLS configuration + parser.add_argument( + f"--{prefix}tls-cert-path", + type=str, + default=None, + help="Path to server TLS certificate (PEM format)", + ) + parser.add_argument( + f"--{prefix}tls-key-path", + type=str, + default=None, + help="Path to server TLS private key (PEM format)", + ) parser.add_argument( f"--{prefix}enable-trace", action="store_true", @@ -678,6 +694,18 @@ class RouterArgs: elif attr.name in cli_args_dict: args_dict[attr.name] = cli_args_dict[attr.name] + # Special handling for CLI args with dashes vs dataclass fields with underscores + # e.g. --tls-cert-path maps to tls_cert_path in args namespace, but we might want server_cert_path in dataclass + # Wait, dataclass fields are server_cert_path/server_key_path + # CLI args are tls_cert_path/tls_key_path + # We need to manually map them if names don't match + + # Map tls args to server cert/key path + if f"{prefix}tls_cert_path" in cli_args_dict: + args_dict["server_cert_path"] = cli_args_dict[f"{prefix}tls_cert_path"] + if f"{prefix}tls_key_path" in cli_args_dict: + args_dict["server_key_path"] = cli_args_dict[f"{prefix}tls_key_path"] + # parse special arguments and remove "--prefill" and "--decode" from cli_args_dict args_dict["prefill_urls"] = cls._parse_prefill_urls( cli_args_dict.get(f"{prefix}prefill", None) diff --git a/sgl-model-gateway/bindings/python/src/lib.rs b/sgl-model-gateway/bindings/python/src/lib.rs index f33caf581..3e7a25b10 100644 --- a/sgl-model-gateway/bindings/python/src/lib.rs +++ b/sgl-model-gateway/bindings/python/src/lib.rs @@ -226,6 +226,8 @@ struct Router { client_cert_path: Option, client_key_path: Option, ca_cert_paths: Vec, + server_cert_path: Option, + server_key_path: Option, enable_trace: bool, otlp_traces_endpoint: String, } @@ -407,6 +409,10 @@ impl Router { self.client_key_path.as_ref(), ) .add_ca_certificates(self.ca_cert_paths.clone()) + .maybe_server_cert_and_key( + self.server_cert_path.as_ref(), + self.server_key_path.as_ref(), + ) .build() } } @@ -488,6 +494,8 @@ impl Router { client_cert_path = None, client_key_path = None, ca_cert_paths = vec![], + server_cert_path = None, + server_key_path = None, enable_trace = false, otlp_traces_endpoint = String::from("localhost:4317"), ))] @@ -566,6 +574,8 @@ impl Router { client_cert_path: Option, client_key_path: Option, ca_cert_paths: Vec, + server_cert_path: Option, + server_key_path: Option, enable_trace: bool, otlp_traces_endpoint: String, ) -> PyResult { @@ -658,6 +668,8 @@ impl Router { client_cert_path, client_key_path, ca_cert_paths, + server_cert_path, + server_key_path, enable_trace, otlp_traces_endpoint, }) diff --git a/sgl-model-gateway/py_test/e2e_http/test_tls.py b/sgl-model-gateway/py_test/e2e_http/test_tls.py new file mode 100644 index 000000000..8a1820b09 --- /dev/null +++ b/sgl-model-gateway/py_test/e2e_http/test_tls.py @@ -0,0 +1,153 @@ +import datetime +import os +import socket +import subprocess +import sys +import time +from contextlib import closing + +import requests +from cryptography import x509 +from cryptography.hazmat.primitives import hashes, serialization +from cryptography.hazmat.primitives.asymmetric import rsa +from cryptography.x509.oid import NameOID +from urllib3.exceptions import InsecureRequestWarning + +# Suppress insecure request warnings due to self-signed cert +requests.packages.urllib3.disable_warnings(InsecureRequestWarning) + + +def find_free_port() -> int: + with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s: + s.bind(("", 0)) + s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + return s.getsockname()[1] + + +def generate_self_signed_cert(cert_path: str, key_path: str) -> None: + """Generate a self-signed certificate and private key for localhost.""" + key = rsa.generate_private_key(public_exponent=65537, key_size=2048) + + subject = issuer = x509.Name( + [ + x509.NameAttribute(NameOID.COUNTRY_NAME, "US"), + x509.NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, "California"), + x509.NameAttribute(NameOID.LOCALITY_NAME, "San Francisco"), + x509.NameAttribute(NameOID.ORGANIZATION_NAME, "SGLang Test"), + x509.NameAttribute(NameOID.COMMON_NAME, "localhost"), + ] + ) + + cert = ( + x509.CertificateBuilder() + .subject_name(subject) + .issuer_name(issuer) + .public_key(key.public_key()) + .serial_number(x509.random_serial_number()) + .not_valid_before(datetime.datetime.utcnow()) + .not_valid_after(datetime.datetime.utcnow() + datetime.timedelta(days=10)) + .add_extension( + x509.SubjectAlternativeName([x509.DNSName("localhost")]), critical=False + ) + .sign(key, hashes.SHA256()) + ) + + with open(key_path, "wb") as f: + f.write( + key.private_bytes( + encoding=serialization.Encoding.PEM, + format=serialization.PrivateFormat.TraditionalOpenSSL, + encryption_algorithm=serialization.NoEncryption(), + ) + ) + + with open(cert_path, "wb") as f: + f.write(cert.public_bytes(serialization.Encoding.PEM)) + + +def test_tls_server() -> None: + """End-to-end test for TLS-enabled router startup and basic endpoints.""" + cert_path = "cert.pem" + key_path = "key.pem" + generate_self_signed_cert(cert_path, key_path) + + port = find_free_port() + + cmd = [ + sys.executable, + "-m", + "sglang_router.launch_router", + "--worker-urls", + "http://127.0.0.1:9999", # Dummy worker + "--host", + "127.0.0.1", + "--port", + str(port), + "--tls-cert-path", + cert_path, + "--tls-key-path", + key_path, + "--log-level", + "info", + ] + + proc = subprocess.Popen( + cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True + ) + + try: + # Wait for server to start and respond to health check + start_time = time.time() + while time.time() - start_time < 15: + try: + response = requests.get( + f"https://localhost:{port}/health", verify=False, timeout=2 + ) + if response.status_code == 200: + break + except requests.RequestException: + pass + + if proc.poll() is not None: + stdout, stderr = proc.communicate() + raise RuntimeError( + f"Router process died early.\nSTDOUT:\n{stdout}\nSTDERR:\n{stderr}" + ) + + time.sleep(0.5) + else: + raise TimeoutError("Server did not become healthy within 15 seconds") + + # Verify basic endpoints work over TLS + models_resp = requests.get( + f"https://localhost:{port}/v1/models", verify=False, timeout=2 + ) + assert models_resp.status_code in ( + 200, + 503, + ) # 503 expected with no healthy workers + + # Minimal generate request (should be rejected or queued) + gen_payload = {"model": "dummy", "prompt": "test", "max_new_tokens": 1} + gen_resp = requests.post( + f"https://localhost:{port}/generate", + json=gen_payload, + verify=False, + timeout=2, + ) + assert gen_resp.status_code in ( + 200, + 400, + 503, + ) # Various valid responses with dummy worker + + finally: + proc.terminate() + try: + proc.wait(timeout=5) + except subprocess.TimeoutExpired: + proc.kill() + + for path in (cert_path, key_path): + if os.path.exists(path): + os.remove(path) diff --git a/sgl-model-gateway/src/config/builder.rs b/sgl-model-gateway/src/config/builder.rs index ac3591222..37180d57d 100644 --- a/sgl-model-gateway/src/config/builder.rs +++ b/sgl-model-gateway/src/config/builder.rs @@ -14,6 +14,8 @@ pub struct RouterConfigBuilder { client_cert_path: Option, client_key_path: Option, ca_cert_paths: Vec, + server_cert_path: Option, + server_key_path: Option, mcp_config_path: Option, } @@ -29,6 +31,8 @@ impl RouterConfigBuilder { client_cert_path: None, client_key_path: None, ca_cert_paths: Vec::new(), + server_cert_path: None, + server_key_path: None, mcp_config_path: None, } } @@ -582,6 +586,30 @@ impl RouterConfigBuilder { self } + // ==================== Server TLS ==================== + + /// Both paths must be provided together. Files read during build() + pub fn server_cert_and_key, S2: Into>( + mut self, + cert_path: S1, + key_path: S2, + ) -> Self { + self.server_cert_path = Some(cert_path.into()); + self.server_key_path = Some(key_path.into()); + self + } + + /// Files read during build() + pub fn maybe_server_cert_and_key( + mut self, + cert_path: Option>, + key_path: Option>, + ) -> Self { + self.server_cert_path = cert_path.map(|p| p.into()); + self.server_key_path = key_path.map(|p| p.into()); + self + } + // ==================== MCP ==================== /// Config file loaded during build() @@ -610,6 +638,9 @@ impl RouterConfigBuilder { // Read mTLS certificates from paths if provided self = self.read_mtls_certificates()?; + // Read Server TLS certificates from paths if provided + self = self.read_server_certificates()?; + // Read MCP config from path if provided self = self.read_mcp_config()?; @@ -672,6 +703,33 @@ impl RouterConfigBuilder { Ok(self) } + /// Internal method to read Server TLS certificates from paths + fn read_server_certificates(mut self) -> ConfigResult { + match (&self.server_cert_path, &self.server_key_path) { + (Some(cert_path), Some(key_path)) => { + let cert = std::fs::read(cert_path).map_err(|e| ConfigError::ValidationFailed { + reason: format!( + "Failed to read server certificate from {}: {}", + cert_path, e + ), + })?; + let key = std::fs::read(key_path).map_err(|e| ConfigError::ValidationFailed { + reason: format!("Failed to read server key from {}: {}", key_path, e), + })?; + self.config.server_cert = Some(cert); + self.config.server_key = Some(key); + } + (None, None) => {} + _ => { + return Err(ConfigError::ValidationFailed { + reason: "Both --tls-cert-path and --tls-key-path must be specified together" + .to_string(), + }); + } + } + Ok(self) + } + /// Internal method to read MCP config from path fn read_mcp_config(mut self) -> ConfigResult { if let Some(mcp_config_path) = &self.mcp_config_path { diff --git a/sgl-model-gateway/src/config/types.rs b/sgl-model-gateway/src/config/types.rs index b12802e29..a933f1a41 100644 --- a/sgl-model-gateway/src/config/types.rs +++ b/sgl-model-gateway/src/config/types.rs @@ -64,6 +64,12 @@ pub struct RouterConfig { pub tool_call_parser: Option, #[serde(default)] pub tokenizer_cache: TokenizerCacheConfig, + /// Server TLS certificate (PEM) + #[serde(skip)] + pub server_cert: Option>, + /// Server TLS private key (PEM) + #[serde(skip)] + pub server_key: Option>, /// Combined certificate + key in PEM format, loaded from client_cert_path and client_key_path during config creation #[serde(skip)] pub client_identity: Option>, @@ -523,6 +529,8 @@ impl Default for RouterConfig { ca_certificates: vec![], mcp_config: None, enable_wasm: false, + server_cert: None, + server_key: None, } } } diff --git a/sgl-model-gateway/src/main.rs b/sgl-model-gateway/src/main.rs index e0c8f6bd9..d497d3712 100644 --- a/sgl-model-gateway/src/main.rs +++ b/sgl-model-gateway/src/main.rs @@ -360,6 +360,12 @@ struct CliArgs { #[arg(long, default_value = "localhost:4317")] otlp_traces_endpoint: String, + + #[arg(long)] + tls_cert_path: Option, + + #[arg(long)] + tls_key_path: Option, } enum OracleConnectSource { @@ -659,7 +665,8 @@ impl CliArgs { .retries(!self.disable_retries) .circuit_breaker(!self.disable_circuit_breaker) .enable_wasm(self.enable_wasm) - .igw(self.enable_igw); + .igw(self.enable_igw) + .maybe_server_cert_and_key(self.tls_cert_path.as_ref(), self.tls_key_path.as_ref()); builder.build() } diff --git a/sgl-model-gateway/src/server.rs b/sgl-model-gateway/src/server.rs index 3133a3f8f..65e816583 100644 --- a/sgl-model-gateway/src/server.rs +++ b/sgl-model-gateway/src/server.rs @@ -14,6 +14,7 @@ use axum::{ routing::{delete, get, post}, serve, Json, Router, }; +use rustls::crypto::ring; use serde::Deserialize; use serde_json::{json, Value}; use tokio::{net::TcpListener, signal, spawn}; @@ -983,13 +984,45 @@ pub async fn startup(config: ServerConfig) -> Result<(), Box)?; + + if let (Some(cert), Some(key)) = ( + &config.router_config.server_cert, + &config.router_config.server_key, + ) { + info!("TLS enabled"); + ring::default_provider() + .install_default() + .map_err(|e| format!("Failed to install rustls ring provider: {e:?}"))?; + + let tls_config = axum_server::tls_rustls::RustlsConfig::from_pem(cert.clone(), key.clone()) + .await + .map_err(|e| format!("Failed to create TLS config: {}", e))?; + + let addr: std::net::SocketAddr = bind_addr + .parse() + .map_err(|e| format!("Invalid address: {}", e))?; + + let handle = axum_server::Handle::new(); + let handle_clone = handle.clone(); + spawn(async move { + shutdown_signal().await; + handle_clone.graceful_shutdown(None); + }); + + axum_server::bind_rustls(addr, tls_config) + .handle(handle) + .serve(app.into_make_service()) + .await + .map_err(|e| Box::new(e) as Box)?; + } else { + let listener = TcpListener::bind(&bind_addr) + .await + .map_err(|e| format!("Failed to bind to {}: {}", bind_addr, e))?; + serve(listener, app) + .with_graceful_shutdown(shutdown_signal()) + .await + .map_err(|e| Box::new(e) as Box)?; + } Ok(()) }