From 5e2cda6158e670e64b926a9985d65826c537ac82 Mon Sep 17 00:00:00 2001 From: Simo Lin Date: Sun, 7 Dec 2025 16:05:53 -0800 Subject: [PATCH] [model-gateway] Fixed WASM Security Vulnerability - Execution Timeout (#14588) --- sgl-model-gateway/src/wasm/errors.rs | 4 +- sgl-model-gateway/src/wasm/runtime.rs | 74 ++++++++++++++++++++++----- 2 files changed, 64 insertions(+), 14 deletions(-) diff --git a/sgl-model-gateway/src/wasm/errors.rs b/sgl-model-gateway/src/wasm/errors.rs index c5b9187ed..561ee75b1 100644 --- a/sgl-model-gateway/src/wasm/errors.rs +++ b/sgl-model-gateway/src/wasm/errors.rs @@ -103,8 +103,8 @@ pub enum WasmRuntimeError { #[error("function not found: {0}")] FunctionNotFound(String), - #[error("execution timeout")] - Timeout, + #[error("execution timeout after {0}ms")] + Timeout(u64), #[error("execution failed: {0}")] CallFailed(String), diff --git a/sgl-model-gateway/src/wasm/runtime.rs b/sgl-model-gateway/src/wasm/runtime.rs index 2223b0968..58c8ffc74 100644 --- a/sgl-model-gateway/src/wasm/runtime.rs +++ b/sgl-model-gateway/src/wasm/runtime.rs @@ -3,19 +3,27 @@ //! Manages WASM component execution using wasmtime with async support. //! Provides a thread pool for concurrent WASM execution and metrics tracking. -use std::sync::{ - atomic::{AtomicU64, Ordering}, - Arc, +use std::{ + sync::{ + atomic::{AtomicU64, Ordering}, + Arc, + }, + time::Duration, }; use tokio::sync::oneshot; use tracing::{debug, error, info}; use wasmtime::{ component::{Component, Linker, ResourceTable}, - Config, Engine, Store, + Config, Engine, Store, UpdateDeadline, }; use wasmtime_wasi::WasiCtx; +/// Epoch increment interval in milliseconds. +/// Epochs are used for cooperative timeout enforcement in WASM execution. +/// A smaller interval gives finer-grained timeout control but slightly more overhead. +const EPOCH_INTERVAL_MS: u64 = 100; + use crate::wasm::{ config::WasmRuntimeConfig, errors::{Result, WasmError, WasmRuntimeError}, @@ -150,6 +158,17 @@ impl WasmRuntime { } } +/// Maps a wasmtime error to a WasmError, detecting epoch interruption (timeout) traps. +fn map_wasm_error(e: wasmtime::Error, timeout_ms: u64) -> WasmError { + // Use proper trap code detection instead of brittle string matching. + // Wasmtime uses Trap::Interrupt for epoch-based interruptions. + if e.downcast_ref::() == Some(&wasmtime::Trap::Interrupt) { + WasmError::from(WasmRuntimeError::Timeout(timeout_ms)) + } else { + WasmError::from(WasmRuntimeError::CallFailed(e.to_string())) + } +} + impl WasmThreadPool { pub fn new(config: WasmRuntimeConfig) -> Result { let (sender, receiver) = async_channel::unbounded(); @@ -163,7 +182,7 @@ impl WasmThreadPool { let num_workers = config.thread_pool_size.clamp(1, max_workers); info!( - target: "sglang_router_rs::wasm::runtime", + target: "sgl_model_gateway::wasm::runtime", "Initializing WASM runtime with {} workers", num_workers ); @@ -178,7 +197,7 @@ impl WasmThreadPool { Ok(rt) => rt, Err(e) => { error!( - target: "sglang_router_rs::wasm::runtime", + target: "sgl_model_gateway::wasm::runtime", worker_id = worker_id, "Failed to create tokio runtime: {}", e @@ -220,7 +239,7 @@ impl WasmThreadPool { config: WasmRuntimeConfig, ) { debug!( - target: "sglang_router_rs::wasm::runtime", + target: "sgl_model_gateway::wasm::runtime", worker_id = worker_id, thread_id = ?std::thread::current().id(), "Worker started" @@ -230,12 +249,13 @@ impl WasmThreadPool { wasmtime_config.async_stack_size(config.max_stack_size); wasmtime_config.async_support(true); wasmtime_config.wasm_component_model(true); // Enable component model + wasmtime_config.epoch_interruption(true); // Enable epoch-based timeout interruption let engine = match Engine::new(&wasmtime_config) { Ok(engine) => engine, Err(e) => { error!( - target: "sglang_router_rs::wasm::runtime", + target: "sgl_model_gateway::wasm::runtime", worker_id = worker_id, "Failed to create engine: {}", e @@ -244,15 +264,36 @@ impl WasmThreadPool { } }; + // Start epoch incrementer for timeout enforcement. + // The engine's epoch counter is incremented periodically, and each Store + // can set a deadline (number of epochs). When the deadline is reached, + // WASM execution is interrupted with a trap. + let engine_for_epoch = engine.clone(); + let epoch_handle = tokio::spawn(async move { + let mut interval = tokio::time::interval(Duration::from_millis(EPOCH_INTERVAL_MS)); + loop { + interval.tick().await; + engine_for_epoch.increment_epoch(); + } + }); + + debug!( + target: "sgl_model_gateway::wasm::runtime", + worker_id = worker_id, + epoch_interval_ms = EPOCH_INTERVAL_MS, + "Epoch incrementer started for timeout enforcement" + ); + loop { let task = match receiver.recv().await { Ok(task) => task, Err(_) => { debug!( - target: "sglang_router_rs::wasm::runtime", + target: "sgl_model_gateway::wasm::runtime", worker_id = worker_id, "Worker shutting down" ); + epoch_handle.abort(); // Stop the epoch incrementer break; // channel closed, exit loop } }; @@ -284,7 +325,7 @@ impl WasmThreadPool { wasm_bytes: Vec, attach_point: WasmModuleAttachPoint, input: WasmComponentInput, - _config: &WasmRuntimeConfig, + config: &WasmRuntimeConfig, ) -> Result { // Compile component from bytes // Note: The WASM file must be in component format (not plain WASM module) @@ -309,6 +350,15 @@ impl WasmThreadPool { }, ); + // Set epoch deadline for timeout enforcement. + // The deadline is the number of epoch ticks before execution is interrupted. + // With EPOCH_INTERVAL_MS=100ms and max_execution_time_ms=1000ms, deadline=10 epochs. + let deadline_epochs = (config.max_execution_time_ms / EPOCH_INTERVAL_MS).max(1); + store.set_epoch_deadline(deadline_epochs); + + // Configure what happens when the deadline is reached during async yields + store.epoch_deadline_callback(|_store| Ok(UpdateDeadline::Yield(1))); + let output = match attach_point { WasmModuleAttachPoint::Middleware(MiddlewareAttachPoint::OnRequest) => { let request = match input { @@ -333,7 +383,7 @@ impl WasmThreadPool { .sgl_model_gateway_middleware_on_request() .call_on_request(&mut store, &request) .await - .map_err(|e| WasmError::from(WasmRuntimeError::CallFailed(e.to_string())))?; + .map_err(|e| map_wasm_error(e, config.max_execution_time_ms))?; WasmComponentOutput::MiddlewareAction(action_result) } @@ -361,7 +411,7 @@ impl WasmThreadPool { .sgl_model_gateway_middleware_on_response() .call_on_response(&mut store, &response) .await - .map_err(|e| WasmError::from(WasmRuntimeError::CallFailed(e.to_string())))?; + .map_err(|e| map_wasm_error(e, config.max_execution_time_ms))?; WasmComponentOutput::MiddlewareAction(action_result) }