feat(gateway): Add server-side TLS support (#15052)

This commit is contained in:
ratish
2025-12-15 03:37:55 +04:00
committed by GitHub
parent 99cb2ed988
commit 0e4108ba29
9 changed files with 312 additions and 8 deletions

View File

@@ -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",

View File

@@ -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):

View File

@@ -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)

View File

@@ -226,6 +226,8 @@ struct Router {
client_cert_path: Option<String>,
client_key_path: Option<String>,
ca_cert_paths: Vec<String>,
server_cert_path: Option<String>,
server_key_path: Option<String>,
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<String>,
client_key_path: Option<String>,
ca_cert_paths: Vec<String>,
server_cert_path: Option<String>,
server_key_path: Option<String>,
enable_trace: bool,
otlp_traces_endpoint: String,
) -> PyResult<Self> {
@@ -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,
})

View File

@@ -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)

View File

@@ -14,6 +14,8 @@ pub struct RouterConfigBuilder {
client_cert_path: Option<String>,
client_key_path: Option<String>,
ca_cert_paths: Vec<String>,
server_cert_path: Option<String>,
server_key_path: Option<String>,
mcp_config_path: Option<String>,
}
@@ -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<S1: Into<String>, S2: Into<String>>(
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<impl Into<String>>,
key_path: Option<impl Into<String>>,
) -> 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<Self> {
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<Self> {
if let Some(mcp_config_path) = &self.mcp_config_path {

View File

@@ -64,6 +64,12 @@ pub struct RouterConfig {
pub tool_call_parser: Option<String>,
#[serde(default)]
pub tokenizer_cache: TokenizerCacheConfig,
/// Server TLS certificate (PEM)
#[serde(skip)]
pub server_cert: Option<Vec<u8>>,
/// Server TLS private key (PEM)
#[serde(skip)]
pub server_key: Option<Vec<u8>>,
/// Combined certificate + key in PEM format, loaded from client_cert_path and client_key_path during config creation
#[serde(skip)]
pub client_identity: Option<Vec<u8>>,
@@ -523,6 +529,8 @@ impl Default for RouterConfig {
ca_certificates: vec![],
mcp_config: None,
enable_wasm: false,
server_cert: None,
server_key: None,
}
}
}

View File

@@ -360,6 +360,12 @@ struct CliArgs {
#[arg(long, default_value = "localhost:4317")]
otlp_traces_endpoint: String,
#[arg(long)]
tls_cert_path: Option<String>,
#[arg(long)]
tls_key_path: Option<String>,
}
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()
}

View File

@@ -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<dyn std::error::Err
// TcpListener::bind accepts &str and handles IPv4/IPv6 via ToSocketAddrs
let bind_addr = format!("{}:{}", config.host, config.port);
info!("Starting server on {}", bind_addr);
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<dyn std::error::Error>)?;
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<dyn std::error::Error>)?;
} 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<dyn std::error::Error>)?;
}
Ok(())
}