feat(gateway): Add server-side TLS support (#15052)
This commit is contained in:
@@ -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",
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
})
|
||||
|
||||
153
sgl-model-gateway/py_test/e2e_http/test_tls.py
Normal file
153
sgl-model-gateway/py_test/e2e_http/test_tls.py
Normal 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)
|
||||
@@ -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 {
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
|
||||
@@ -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(())
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user