chore: vendor sglang v0.5.10 snapshot

third_party/sglang/sgl-model-gateway/examples/wasm/wasm-guest-ratelimit/src/lib.rs

@@ -0,0 +1,155 @@
+//! WASM Guest Rate Limit Example for sgl-model-gateway
+//!
+//! This example demonstrates rate limiting middleware
+//! for sgl-model-gateway using the WebAssembly Component Model.
+//!
+//! Features:
+//! - Rate limiting based on API Key or IP address
+//! - Fixed time window (e.g., 60 requests per minute)
+//! - Returns 429 Too Many Requests when limit exceeded
+//!
+//! Note: This is a simplified implementation. Since WASM components are stateless,
+//! each instance maintains its own counters. For production use, consider
+//! implementing rate limiting at the host/router level with shared state.
+
+wit_bindgen::generate!({
+    path: "../../../src/wasm/interface",
+    world: "sgl-model-gateway",
+});
+
+use std::cell::RefCell;
+
+use exports::sgl::model_gateway::{
+    middleware_on_request::Guest as OnRequestGuest,
+    middleware_on_response::Guest as OnResponseGuest,
+};
+use sgl::model_gateway::middleware_types::{Action, Request, Response};
+
+/// Main middleware implementation
+struct Middleware;
+
+// Rate limit configuration
+const RATE_LIMIT_REQUESTS: u64 = 60; // Maximum requests per window
+const RATE_LIMIT_WINDOW_MS: u64 = 60_000; // Time window in milliseconds (1 minute)
+
+// Simple in-memory counter (per WASM instance)
+// In a real implementation, this would be shared across all instances
+// This is a simplified example for demonstration purposes
+struct RateLimitState {
+    requests: Vec<(String, u64)>, // (identifier, timestamp_ms)
+}
+
+impl RateLimitState {
+    fn new() -> Self {
+        Self {
+            requests: Vec::new(),
+        }
+    }
+
+    // Clean up old entries outside the time window
+    fn cleanup(&mut self, current_time_ms: u64) {
+        let cutoff = current_time_ms.saturating_sub(RATE_LIMIT_WINDOW_MS);
+        self.requests.retain(|(_, timestamp)| *timestamp > cutoff);
+    }
+
+    // Check if identifier has exceeded rate limit
+    fn check_limit(&mut self, identifier: &str, current_time_ms: u64) -> bool {
+        self.cleanup(current_time_ms);
+
+        // Count requests in current window for this identifier
+        let count = self
+            .requests
+            .iter()
+            .filter(|(id, timestamp)| {
+                id == identifier
+                    && *timestamp > current_time_ms.saturating_sub(RATE_LIMIT_WINDOW_MS)
+            })
+            .count() as u64;
+
+        if count >= RATE_LIMIT_REQUESTS {
+            return false; // Limit exceeded
+        }
+
+        // Add new request
+        self.requests
+            .push((identifier.to_string(), current_time_ms));
+        true // Within limit
+    }
+}
+
+// Thread-local state (per WASM instance thread)
+// Using thread_local! is safer than static mut as it avoids unsafe blocks
+// and provides separate state for each thread automatically
+thread_local! {
+    static RATE_LIMIT_STATE: RefCell<RateLimitState> = RefCell::new(RateLimitState::new());
+}
+
+fn get_identifier(req: &Request) -> String {
+    // Helper function to find header value
+    let find_header_value =
+        |headers: &[sgl::model_gateway::middleware_types::Header], name: &str| -> Option<String> {
+            headers
+                .iter()
+                .find(|h| h.name.eq_ignore_ascii_case(name))
+                .map(|h| h.value.clone())
+        };
+
+    // Prefer API Key as identifier (more stable than IP)
+    if let Some(auth_header) = find_header_value(&req.headers, "authorization") {
+        if auth_header.starts_with("Bearer ") {
+            return format!("api_key:{}", &auth_header[7..]);
+        } else if auth_header.starts_with("ApiKey ") {
+            return format!("api_key:{}", &auth_header[7..]);
+        }
+    }
+
+    if let Some(api_key) = find_header_value(&req.headers, "x-api-key") {
+        return format!("api_key:{}", api_key);
+    }
+
+    // Fall back to IP address from forwarded headers
+    if let Some(forwarded_for) = find_header_value(&req.headers, "x-forwarded-for") {
+        // Take first IP from comma-separated list
+        let ip = forwarded_for.split(',').next().unwrap_or("").trim();
+        if !ip.is_empty() {
+            return format!("ip:{}", ip);
+        }
+    }
+
+    if let Some(real_ip) = find_header_value(&req.headers, "x-real-ip") {
+        return format!("ip:{}", real_ip);
+    }
+
+    // Last resort: use request ID (not ideal, but better than nothing)
+    format!("req_id:{}", req.request_id)
+}
+
+// Implement on-request interface
+impl OnRequestGuest for Middleware {
+    fn on_request(req: Request) -> Action {
+        let identifier = get_identifier(&req);
+        let current_time_ms = req.now_epoch_ms;
+
+        // Access thread-local state safely without unsafe blocks
+        // Each thread gets its own RateLimitState instance
+        RATE_LIMIT_STATE.with(|state| {
+            let mut state = state.borrow_mut();
+            if !state.check_limit(&identifier, current_time_ms) {
+                // Rate limit exceeded
+                return Action::Reject(429);
+            }
+            // Within rate limit, continue processing
+            Action::Continue
+        })
+    }
+}
+
+// Implement on-response interface (empty - not used for rate limiting)
+impl OnResponseGuest for Middleware {
+    fn on_response(_resp: Response) -> Action {
+        Action::Continue
+    }
+}
+
+// Export the component
+export!(Middleware);