diff --git a/model_gateway/src/middleware.rs b/model_gateway/src/middleware.rs deleted file mode 100644 index 2c4bca6be4..0000000000 --- a/model_gateway/src/middleware.rs +++ /dev/null @@ -1,1110 +0,0 @@ -use std::{ - pin::Pin, - sync::Arc, - task::{Context, Poll}, - time::{Duration, Instant}, -}; - -use axum::{ - body::Body, - extract::{Request, State}, - http::{header, HeaderValue, StatusCode}, - middleware::Next, - response::{IntoResponse, Response}, - Json, -}; -use bytes::Bytes; -use http_body::Frame; -use rand::Rng; -use serde_json::json; -use sha2::{Digest, Sha256}; -use smg_data_connector::{with_request_context, RequestContext as StorageRequestContext}; -use subtle::ConstantTimeEq; -use tokio::sync::{mpsc, oneshot}; -use tower::{Layer, Service}; -use tower_http::trace::{MakeSpan, OnRequest, OnResponse, TraceLayer}; -use tracing::{debug, error, field::Empty, info, info_span, warn, Span}; -use tracing_opentelemetry::OpenTelemetrySpanExt; - -pub use crate::worker::token_bucket::TokenBucket; -use crate::{ - config::RouterConfig, - observability::{ - inflight_tracker::InFlightRequestTracker, - metrics::{method_to_static_str, metrics_labels, Metrics}, - otel_trace::extract_trace_context_http, - }, - routers::error::extract_error_code_from_response, - server::AppState, - wasm::{ - module::{MiddlewareAttachPoint, WasmModuleAttachPoint}, - spec::{ - apply_modify_action_to_headers, build_wasm_headers_from_axum_headers, - smg::gateway::middleware_types::{ - Action, Request as WasmRequest, Response as WasmResponse, - }, - }, - types::WasmComponentInput, - }, -}; - -fn extract_header_str(headers: &http::HeaderMap, name: &str) -> Option { - headers - .get(name) - .and_then(|value| value.to_str().ok()) - .map(str::trim) - .filter(|value| !value.is_empty()) - .map(ToOwned::to_owned) -} - -fn build_storage_request_context( - config: &RouterConfig, - headers: &http::HeaderMap, -) -> Option { - let mut ctx = StorageRequestContext::new(); - - for (header_name, context_key) in &config.storage_context_headers { - let header_name = header_name.trim(); - let context_key = context_key.trim(); - - if header_name.is_empty() || context_key.is_empty() { - continue; - } - - if let Some(value) = extract_header_str(headers, header_name) { - ctx.set(context_key, value); - } - } - - (!ctx.data().is_empty()).then_some(ctx) -} - -pub async fn storage_context_middleware( - State(state): State>, - request: Request, - next: Next, -) -> Response { - if state - .context - .router_config - .storage_context_headers - .is_empty() - { - return next.run(request).await; - } - - match build_storage_request_context(&state.context.router_config, request.headers()) { - Some(ctx) => with_request_context(ctx, next.run(request)).await, - None => next.run(request).await, - } -} - -/// A body wrapper that holds a token and returns it when the body is fully consumed or dropped. -/// This ensures that for streaming responses, the token is only returned after the entire -/// stream has been sent to the client. -pub struct TokenGuardBody { - inner: Body, - /// The token bucket to return tokens to. Uses Option so we can take() on drop. - token_bucket: Option>, - /// Number of tokens to return. - tokens: f64, -} - -impl TokenGuardBody { - /// Create a new TokenGuardBody that will return tokens when dropped. - pub fn new(inner: Body, token_bucket: Arc, tokens: f64) -> Self { - Self { - inner, - token_bucket: Some(token_bucket), - tokens, - } - } -} - -impl Drop for TokenGuardBody { - fn drop(&mut self) { - if let Some(bucket) = self.token_bucket.take() { - debug!( - "TokenGuardBody: stream ended, returning {} tokens to bucket", - self.tokens - ); - // Use lock-free sync return - no runtime needed, guaranteed token return - bucket.return_tokens_sync(self.tokens); - } - } -} - -impl http_body::Body for TokenGuardBody { - type Data = Bytes; - type Error = axum::Error; - - fn poll_frame( - self: Pin<&mut Self>, - cx: &mut Context<'_>, - ) -> Poll, Self::Error>>> { - // SAFETY: We never move the inner body, and Body is Unpin - // (it's a type alias for UnsyncBoxBody which is Unpin) - let this = self.get_mut(); - Pin::new(&mut this.inner).poll_frame(cx) - } - - fn is_end_stream(&self) -> bool { - self.inner.is_end_stream() - } - - fn size_hint(&self) -> http_body::SizeHint { - self.inner.size_hint() - } -} - -#[derive(Clone)] -pub struct AuthConfig { - /// Precomputed SHA-256 hash of the API key, used for constant-time comparison - /// that doesn't leak key length via timing. - api_key_hash: Option<[u8; 32]>, -} - -impl AuthConfig { - pub fn new(api_key: Option) -> Self { - Self { - api_key_hash: api_key.map(|k| Sha256::digest(k.as_bytes()).into()), - } - } -} - -/// Middleware to validate Bearer token against configured API key. -/// Only active when router has an API key configured. -pub async fn auth_middleware( - State(auth_config): State, - request: Request, - next: Next, -) -> Response { - if let Some(expected_hash) = &auth_config.api_key_hash { - let token = request - .headers() - .get(header::AUTHORIZATION) - .and_then(|h| h.to_str().ok()) - .and_then(|h| h.strip_prefix("Bearer ")); - - let authorized = token.is_some_and(|t| { - Sha256::digest(t.as_bytes()) - .as_slice() - .ct_eq(expected_hash) - .unwrap_u8() - == 1 - }); - if !authorized { - return StatusCode::UNAUTHORIZED.into_response(); - } - } - - next.run(request).await -} - -/// Alphanumeric characters for request ID generation (as bytes for O(1) indexing) -const REQUEST_ID_CHARS: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789"; - -/// Generate OpenAI-compatible request ID based on endpoint. -fn generate_request_id(path: &str) -> String { - let prefix = if path.contains("/chat/completions") { - "chatcmpl-" - } else if path.contains("/completions") { - "cmpl-" - } else if path.contains("/generate") { - "gnt-" - } else if path.contains("/responses") { - "resp-" - } else { - "req-" - }; - - // Generate a random string similar to OpenAI's format - // Use byte array indexing (O(1)) instead of chars().nth() (O(n)) - let mut rng = rand::rng(); - let random_part: String = (0..24) - .map(|_| { - let idx = rng.random_range(0..REQUEST_ID_CHARS.len()); - REQUEST_ID_CHARS[idx] as char - }) - .collect(); - - format!("{prefix}{random_part}") -} - -// Re-export RequestId from auth crate for backward compatibility -pub use smg_auth::RequestId; - -/// Tower Layer for request ID middleware -#[derive(Clone)] -pub struct RequestIdLayer { - headers: Arc>, -} - -impl RequestIdLayer { - pub fn new(headers: Vec) -> Self { - Self { - headers: Arc::new(headers), - } - } -} - -impl Layer for RequestIdLayer { - type Service = RequestIdMiddleware; - - fn layer(&self, inner: S) -> Self::Service { - RequestIdMiddleware { - inner, - headers: self.headers.clone(), - } - } -} - -/// Tower Service for request ID middleware -#[derive(Clone)] -pub struct RequestIdMiddleware { - inner: S, - headers: Arc>, -} - -impl Service for RequestIdMiddleware -where - S: Service + Send + 'static, - S::Future: Send + 'static, -{ - type Response = S::Response; - type Error = S::Error; - type Future = - Pin> + Send>>; - - fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { - self.inner.poll_ready(cx) - } - - fn call(&mut self, mut req: Request) -> Self::Future { - let headers = self.headers.clone(); - - // Extract request ID from headers or generate new one - let mut request_id = None; - - for header_name in headers.iter() { - if let Some(header_value) = req.headers().get(header_name) { - if let Ok(value) = header_value.to_str() { - request_id = Some(value.to_string()); - break; - } - } - } - - let request_id = request_id.unwrap_or_else(|| generate_request_id(req.uri().path())); - - // Insert request ID into request extensions for other middleware/handlers to use - req.extensions_mut().insert(RequestId(request_id.clone())); - - // Call the inner service - let future = self.inner.call(req); - - Box::pin(async move { - let mut response = future.await?; - - // Add request ID to response headers - response.headers_mut().insert( - "x-request-id", - HeaderValue::from_str(&request_id) - .unwrap_or_else(|_| HeaderValue::from_static("invalid-request-id")), - ); - - Ok(response) - }) - } -} - -/// Custom span maker that includes request ID -#[derive(Clone, Debug)] -pub struct RequestSpan; - -impl MakeSpan for RequestSpan { - fn make_span(&mut self, request: &Request) -> Span { - // Extract incoming W3C trace context (traceparent/tracestate) so that - // server-side spans become children of the caller's distributed trace. - let parent_cx = extract_trace_context_http(request.headers()); - - // Don't try to extract request ID here - it won't be available yet - // The RequestIdLayer runs after TraceLayer creates the span - let span = info_span!( - target: "smg::otel-trace", - "http_request", - method = %request.method(), - uri = %request.uri(), - version = ?request.version(), - request_id = Empty, // Will be set later - status_code = Empty, - latency = Empty, - error = Empty, - module = "smg" - ); - - span.set_parent(parent_cx); - span - } -} - -/// Custom on_request handler -#[derive(Clone, Debug)] -pub struct RequestLogger; - -impl OnRequest for RequestLogger { - fn on_request(&mut self, request: &Request, span: &Span) { - let _enter = span.enter(); - - // Try to get the request ID from extensions - // This will work if RequestIdLayer has already run - if let Some(request_id) = request.extensions().get::() { - span.record("request_id", request_id.0.as_str()); - } - - let method = method_to_static_str(request.method().as_str()); - let path = normalize_path_for_metrics(request.uri().path()); - Metrics::record_http_request(method, &path); - - // Log the request start - info!( - target: "smg::request", - "started processing request" - ); - } -} - -/// Custom on_response handler -#[derive(Clone, Debug, Default)] -pub struct ResponseLogger; - -impl OnResponse for ResponseLogger { - fn on_response(self, response: &Response, latency: Duration, span: &Span) { - let status = response.status(); - let status_code = status.as_u16(); - - let error_code = extract_error_code_from_response(response); - - // Layer 1: HTTP metrics - Metrics::record_http_response(status_code, error_code); - - // Record these in the span for structured logging/observability tools - span.record("status_code", status_code); - // Use microseconds as integer to avoid format! string allocation - span.record("latency", latency.as_micros() as u64); - - // Log the response completion - let _enter = span.enter(); - if status.is_server_error() { - error!( - target: "smg::response", - "request failed with server error" - ); - } else if status.is_client_error() { - warn!( - target: "smg::response", - "request failed with client error" - ); - } else { - info!( - target: "smg::response", - "finished processing request" - ); - } - } -} - -/// Create a configured TraceLayer for HTTP logging -/// Note: Actual request/response logging with request IDs is done in RequestIdService -pub fn create_logging_layer() -> TraceLayer< - tower_http::classify::SharedClassifier, - RequestSpan, - RequestLogger, - ResponseLogger, -> { - TraceLayer::new_for_http() - .make_span_with(RequestSpan) - .on_request(RequestLogger) - .on_response(ResponseLogger) -} - -/// Request queue entry -pub struct QueuedRequest { - /// Time when the request was queued - queued_at: Instant, - /// Channel to send the permit back when acquired - permit_tx: oneshot::Sender>, -} - -/// Queue processor that handles queued requests -pub struct QueueProcessor { - token_bucket: Arc, - queue_rx: mpsc::Receiver, - queue_timeout: Duration, -} - -impl QueueProcessor { - pub fn new( - token_bucket: Arc, - queue_rx: mpsc::Receiver, - queue_timeout: Duration, - ) -> Self { - Self { - token_bucket, - queue_rx, - queue_timeout, - } - } - - pub async fn run(mut self) { - debug!("Starting concurrency queue processor"); - - // Process requests in a single task to reduce overhead - while let Some(queued) = self.queue_rx.recv().await { - // Check timeout immediately - let elapsed = queued.queued_at.elapsed(); - if elapsed >= self.queue_timeout { - warn!("Request already timed out in queue"); - let _ = queued.permit_tx.send(Err(StatusCode::REQUEST_TIMEOUT)); - continue; - } - - let remaining_timeout = self.queue_timeout - elapsed; - - // Try to acquire token for this request - if self.token_bucket.try_acquire(1.0).is_ok() { - // Got token immediately - debug!("Queue: acquired token immediately for queued request"); - let _ = queued.permit_tx.send(Ok(())); - } else { - // Need to wait for token - let token_bucket = self.token_bucket.clone(); - - // Spawn task only when we actually need to wait - #[expect( - clippy::disallowed_methods, - reason = "fire-and-forget permit acquisition: task is bounded by remaining_timeout and communicates via oneshot; dropping the JoinHandle detaches the task but it self-terminates" - )] - tokio::spawn(async move { - if token_bucket - .acquire_timeout(1.0, remaining_timeout) - .await - .is_ok() - { - debug!("Queue: acquired token after waiting"); - let _ = queued.permit_tx.send(Ok(())); - } else { - warn!("Queue: request timed out waiting for token"); - let _ = queued.permit_tx.send(Err(StatusCode::REQUEST_TIMEOUT)); - } - }); - } - } - - warn!("Concurrency queue processor shutting down"); - } -} - -/// State for the concurrency limiter -pub struct ConcurrencyLimiter { - pub queue_tx: Option>, -} - -impl ConcurrencyLimiter { - /// Create new concurrency limiter with optional queue - pub fn new( - token_bucket: Option>, - queue_size: usize, - queue_timeout: Duration, - ) -> (Self, Option) { - match (token_bucket, queue_size) { - (None, _) => (Self { queue_tx: None }, None), - (Some(bucket), size) if size > 0 => { - let (queue_tx, queue_rx) = mpsc::channel(size); - let processor = QueueProcessor::new(bucket, queue_rx, queue_timeout); - ( - Self { - queue_tx: Some(queue_tx), - }, - Some(processor), - ) - } - (Some(_), _) => (Self { queue_tx: None }, None), - } - } -} - -/// Middleware function for concurrency limiting with optional queuing -pub async fn concurrency_limit_middleware( - State(app_state): State>, - request: Request, - next: Next, -) -> Response { - // Check mesh global rate limit first if mesh is enabled - // If mesh is not enabled, this check is skipped and local rate limiting is used - if let Some(mesh_handler) = &app_state.mesh_handler { - let (is_exceeded, current_count, limit) = - mesh_handler.sync_manager.check_global_rate_limit(); - if is_exceeded { - debug!( - "Global rate limit exceeded: {}/{} req/s", - current_count, limit - ); - return ( - StatusCode::TOO_MANY_REQUESTS, - Json(json!({ - "error": "Rate limit exceeded", - "current_count": current_count, - "limit": limit - })), - ) - .into_response(); - } - } - - let token_bucket = match &app_state.context.rate_limiter { - Some(bucket) => bucket.clone(), - None => { - // Rate limiting disabled, pass through immediately - return next.run(request).await; - } - }; - - // Try to acquire token immediately - if token_bucket.try_acquire(1.0).is_ok() { - debug!("Acquired token immediately"); - Metrics::record_http_rate_limit(metrics_labels::RATE_LIMIT_ALLOWED); - let response = next.run(request).await; - - // Wrap the response body with TokenGuardBody to return token when stream ends - // This ensures that for streaming responses, the token is only returned - // after the entire stream has been sent to the client. - let (parts, body) = response.into_parts(); - let guarded_body = TokenGuardBody::new(body, token_bucket, 1.0); - Response::from_parts(parts, Body::new(guarded_body)) - } else { - // No tokens available, try to queue if enabled - if let Some(queue_tx) = &app_state.concurrency_queue_tx { - debug!("No tokens available, attempting to queue request"); - - // Create a channel for the token response - let (permit_tx, permit_rx) = oneshot::channel(); - - let queued = QueuedRequest { - queued_at: Instant::now(), - permit_tx, - }; - - // Try to send to queue - match queue_tx.try_send(queued) { - Ok(()) => { - // Wait for token from queue processor - match permit_rx.await { - Ok(Ok(())) => { - debug!("Acquired token from queue"); - Metrics::record_http_rate_limit(metrics_labels::RATE_LIMIT_ALLOWED); - let response = next.run(request).await; - - // Wrap the response body with TokenGuardBody to return token when stream ends - let (parts, body) = response.into_parts(); - let guarded_body = TokenGuardBody::new(body, token_bucket, 1.0); - Response::from_parts(parts, Body::new(guarded_body)) - } - Ok(Err(status)) => { - warn!("Queue returned error status: {}", status); - Metrics::record_http_rate_limit(metrics_labels::RATE_LIMIT_REJECTED); - status.into_response() - } - Err(_) => { - error!("Queue response channel closed"); - Metrics::record_http_rate_limit(metrics_labels::RATE_LIMIT_REJECTED); - StatusCode::INTERNAL_SERVER_ERROR.into_response() - } - } - } - Err(_) => { - warn!("Request queue is full, returning 429"); - Metrics::record_http_rate_limit(metrics_labels::RATE_LIMIT_REJECTED); - StatusCode::TOO_MANY_REQUESTS.into_response() - } - } - } else { - warn!("No tokens available and queuing is disabled, returning 429"); - Metrics::record_http_rate_limit(metrics_labels::RATE_LIMIT_REJECTED); - StatusCode::TOO_MANY_REQUESTS.into_response() - } - } -} - -// ============================================================================ -// HTTP Metrics Layer (Layer 1: SMG metrics) -// ============================================================================ - -/// Tower Layer for HTTP metrics collection (SMG Layer 1 metrics) -#[derive(Clone)] -pub struct HttpMetricsLayer { - tracker: Arc, -} - -impl HttpMetricsLayer { - pub fn new(tracker: Arc) -> Self { - Self { tracker } - } -} - -impl Layer for HttpMetricsLayer { - type Service = HttpMetricsMiddleware; - - fn layer(&self, inner: S) -> Self::Service { - HttpMetricsMiddleware { - inner, - in_flight_request_tracker: self.tracker.clone(), - } - } -} - -/// Tower Service for HTTP metrics collection -#[derive(Clone)] -pub struct HttpMetricsMiddleware { - inner: S, - in_flight_request_tracker: Arc, -} - -impl Service for HttpMetricsMiddleware -where - S: Service + Send + Clone + 'static, - S::Future: Send + 'static, -{ - type Response = S::Response; - type Error = S::Error; - type Future = - Pin> + Send>>; - - fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { - self.inner.poll_ready(cx) - } - - fn call(&mut self, req: Request) -> Self::Future { - // Convert method to static string to avoid allocation - let method = method_to_static_str(req.method().as_str()); - let path = normalize_path_for_metrics(req.uri().path()); - let start = Instant::now(); - - let mut inner = self.inner.clone(); - let in_flight_request_tracker = self.in_flight_request_tracker.clone(); - - Box::pin(async move { - let guard = in_flight_request_tracker.track(); - Metrics::set_http_connections_active(in_flight_request_tracker.len()); - - // Capture result before dropping guard to ensure decrement happens on error too - let result = inner.call(req).await; - - drop(guard); - Metrics::set_http_connections_active(in_flight_request_tracker.len()); - - let response = result?; - - let duration = start.elapsed(); - Metrics::record_http_duration(method, &path, duration); - - Ok(response) - }) - } -} - -/// Normalize path for metrics to avoid high cardinality. -/// Replaces dynamic segments (IDs, UUIDs) with `{id}` placeholder. -/// Only allocates when normalization is needed; uses single-pass with byte offsets. -fn normalize_path_for_metrics(path: &str) -> String { - let bytes = path.as_bytes(); - let mut segment_start = 0; - let mut segment_idx = 0; - let mut result: Option = None; - - for (pos, &b) in bytes.iter().enumerate() { - if b == b'/' || pos == bytes.len() - 1 { - // Determine segment end (include last char if not a slash) - let segment_end = if b == b'/' { pos } else { pos + 1 }; - let segment = &path[segment_start..segment_end]; - - // Check segments after index 2 for dynamic IDs - if segment_idx > 2 && !segment.is_empty() && is_dynamic_id(segment) { - // Initialize result with everything before this segment - let result = result.get_or_insert_with(|| { - let mut s = String::with_capacity(path.len()); - s.push_str(&path[..segment_start]); - s - }); - result.push_str("{id}"); - } else if let Some(ref mut r) = result { - // Already normalizing, append this segment as-is - r.push_str(segment); - } - - // Add slash after segment (except at end) - if b == b'/' { - if let Some(ref mut r) = result { - r.push('/'); - } - segment_start = pos + 1; - segment_idx += 1; - } - } - } - - result.unwrap_or_else(|| path.to_owned()) -} - -/// Check if segment looks like a dynamic ID (prefixed ID, UUID, or numeric). -#[inline] -fn is_dynamic_id(s: &str) -> bool { - // Prefixed IDs: resp_xxx, chatcmpl_xxx (len > 10 with underscore) - if s.len() > 10 && s.contains('_') { - return true; - } - // UUIDs: 32+ hex chars with dashes - if s.len() >= 32 && s.bytes().all(|b| b.is_ascii_hexdigit() || b == b'-') { - return true; - } - // Numeric IDs - !s.is_empty() && s.bytes().all(|b| b.is_ascii_digit()) -} - -pub async fn wasm_middleware( - State(app_state): State>, - request: Request, - next: Next, -) -> Response { - // Check if WASM is enabled - if !app_state.context.router_config.enable_wasm { - return next.run(request).await; - } - - // Get WASM manager - let wasm_manager = match &app_state.context.wasm_manager { - Some(manager) => manager, - None => { - return next.run(request).await; - } - }; - - // Get request ID from extensions or generate one - let request_id = request - .extensions() - .get::() - .map(|r| r.0.clone()) - .unwrap_or_else(|| generate_request_id(request.uri().path())); - - // ===== OnRequest Phase ===== - let on_request_attach_point = - WasmModuleAttachPoint::Middleware(MiddlewareAttachPoint::OnRequest); - - let modules_on_request = - match wasm_manager.get_modules_by_attach_point(on_request_attach_point.clone()) { - Ok(modules) => modules, - Err(e) => { - error!("Failed to get WASM modules for OnRequest: {}", e); - return next.run(request).await; - } - }; - - let response = if modules_on_request.is_empty() { - next.run(request).await - } else { - // Decompose request to preserve extensions across reconstruction - let (parts, body) = request.into_parts(); - let method = parts.method; - let uri = parts.uri; - let mut headers = parts.headers; - let extensions = parts.extensions; - - let max_body_size = wasm_manager.get_max_body_size(); - let body_bytes = match axum::body::to_bytes(body, max_body_size).await { - Ok(bytes) => bytes.to_vec(), - Err(e) => { - error!("Failed to read request body for WASM processing: {}", e); - return ( - StatusCode::BAD_REQUEST, - Json(json!({"error": format!("Failed to read request body: {e}")})), - ) - .into_response(); - } - }; - - // Process each OnRequest module - let mut modified_body = body_bytes; - - // Pre-compute strings once before the loop to avoid repeated allocations - let method_str = method.to_string(); - let path_str = uri.path().to_string(); - let query_str = uri.query().unwrap_or("").to_string(); - - for module in modules_on_request { - let wasm_headers = build_wasm_headers_from_axum_headers(&headers); - let wasm_request = WasmRequest { - method: method_str.clone(), - path: path_str.clone(), - query: query_str.clone(), - headers: wasm_headers, - body: modified_body.clone(), - request_id: request_id.clone(), - now_epoch_ms: std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap_or_else(|_| Duration::from_millis(0)) - .as_millis() as u64, - }; - - let action = match wasm_manager - .execute_module_for_attach_point( - &module, - on_request_attach_point.clone(), - WasmComponentInput::MiddlewareRequest(wasm_request), - ) - .await - { - Some(action) => action, - None => continue, - }; - - match action { - Action::Continue => {} - Action::Reject(status) => { - return StatusCode::from_u16(status) - .unwrap_or(StatusCode::BAD_REQUEST) - .into_response(); - } - Action::Modify(modify) => { - apply_modify_action_to_headers(&mut headers, &modify); - if let Some(body_bytes) = modify.body_replace { - modified_body = body_bytes; - } - } - } - } - - // Reconstruct request with modifications, preserving original extensions - let mut final_request = Request::builder() - .method(method) - .uri(uri) - .body(Body::from(modified_body)) - .unwrap_or_else(|_| Request::new(Body::empty())); - *final_request.headers_mut() = headers; - *final_request.extensions_mut() = extensions; - - next.run(final_request).await - }; - - // ===== OnResponse Phase ===== - let on_response_attach_point = - WasmModuleAttachPoint::Middleware(MiddlewareAttachPoint::OnResponse); - - let modules_on_response = - match wasm_manager.get_modules_by_attach_point(on_response_attach_point.clone()) { - Ok(modules) => modules, - Err(e) => { - error!("Failed to get WASM modules for OnResponse: {}", e); - return response; - } - }; - if modules_on_response.is_empty() { - return response; - } - - // Skip WASM OnResponse processing for streaming responses to avoid - // buffering the entire stream into memory (breaks SSE, causes OOM on large streams). - let is_streaming = response - .headers() - .get(header::CONTENT_TYPE) - .and_then(|v| v.to_str().ok()) - .is_some_and(|ct| ct.contains("text/event-stream") || ct.contains("application/x-ndjson")) - || response - .headers() - .get(header::TRANSFER_ENCODING) - .and_then(|v| v.to_str().ok()) - .is_some_and(|te| te.contains("chunked")); - if is_streaming { - warn!("Skipping WASM OnResponse for streaming response; OnResponse modules do not apply to streaming"); - return response; - } - - // Extract response data once before processing modules - let mut status = response.status(); - let mut headers = response.headers().clone(); - let max_body_size = wasm_manager.get_max_body_size(); - let mut body_bytes = match axum::body::to_bytes(response.into_body(), max_body_size).await { - Ok(bytes) => bytes.to_vec(), - Err(e) => { - error!("Failed to read response body for WASM processing: {}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": "Failed to read response body"})), - ) - .into_response(); - } - }; - - // Process each OnResponse module - for module in modules_on_response { - let wasm_headers = build_wasm_headers_from_axum_headers(&headers); - let wasm_response = WasmResponse { - status: status.as_u16(), - headers: wasm_headers, - body: body_bytes.clone(), - }; - - let action = match wasm_manager - .execute_module_for_attach_point( - &module, - on_response_attach_point.clone(), - WasmComponentInput::MiddlewareResponse(wasm_response), - ) - .await - { - Some(action) => action, - None => continue, - }; - - match action { - Action::Continue => {} - Action::Reject(status_code) => { - status = StatusCode::from_u16(status_code).unwrap_or(StatusCode::BAD_REQUEST); - let mut final_response = Response::builder() - .status(status) - .body(Body::from(body_bytes)) - .unwrap_or_else(|_| Response::new(Body::empty())); - *final_response.headers_mut() = headers; - return final_response; - } - Action::Modify(modify) => { - if let Some(new_status) = modify.status { - status = StatusCode::from_u16(new_status).unwrap_or(status); - } - apply_modify_action_to_headers(&mut headers, &modify); - if let Some(new_body) = modify.body_replace { - body_bytes = new_body; - } - } - } - } - - // Reconstruct final response with all modifications - let mut final_response = Response::builder() - .status(status) - .body(Body::from(body_bytes)) - .unwrap_or_else(|_| Response::new(Body::empty())); - *final_response.headers_mut() = headers; - final_response -} - -#[cfg(test)] -mod tests { - use axum::http::{HeaderMap, HeaderValue}; - - use super::*; - use crate::config::RouterConfig; - - #[test] - fn build_storage_request_context_maps_configured_headers() { - let config = RouterConfig { - storage_context_headers: std::collections::HashMap::from([ - ("x-tenant-id".to_string(), "tenant_id".to_string()), - ("x-user-id".to_string(), "user_id".to_string()), - ]), - ..Default::default() - }; - - let mut headers = HeaderMap::new(); - headers.insert("x-tenant-id", HeaderValue::from_static("tenant-abc")); - headers.insert("x-user-id", HeaderValue::from_static("user-123")); - - let ctx = build_storage_request_context(&config, &headers).unwrap(); - - assert_eq!(ctx.get("tenant_id"), Some("tenant-abc")); - assert_eq!(ctx.get("user_id"), Some("user-123")); - } - - #[test] - fn build_storage_request_context_ignores_empty_entries_and_missing_headers() { - let config = RouterConfig { - storage_context_headers: std::collections::HashMap::from([ - (" ".to_string(), "tenant_id".to_string()), - ("x-empty-key".to_string(), " ".to_string()), - ("x-present".to_string(), "present_key".to_string()), - ]), - ..Default::default() - }; - - let mut headers = HeaderMap::new(); - headers.insert("x-present", HeaderValue::from_static(" keep-me ")); - - let ctx = build_storage_request_context(&config, &headers).unwrap(); - - assert_eq!(ctx.get("present_key"), Some("keep-me")); - assert_eq!(ctx.data().len(), 1); - } - - #[test] - fn test_normalize_path_no_ids() { - // Common API paths should pass through unchanged - assert_eq!( - normalize_path_for_metrics("/v1/chat/completions"), - "/v1/chat/completions" - ); - assert_eq!( - normalize_path_for_metrics("/v1/completions"), - "/v1/completions" - ); - assert_eq!(normalize_path_for_metrics("/v1/models"), "/v1/models"); - assert_eq!(normalize_path_for_metrics("/health"), "/health"); - } - - #[test] - fn test_normalize_path_with_prefixed_id() { - // Prefixed IDs (resp_xxx, chatcmpl_xxx) should be normalized - assert_eq!( - normalize_path_for_metrics("/v1/responses/resp_abc123def456"), - "/v1/responses/{id}" - ); - assert_eq!( - normalize_path_for_metrics("/v1/chat/completions/chatcmpl_abc123xyz"), - "/v1/chat/completions/{id}" - ); - } - - #[test] - fn test_normalize_path_with_uuid() { - assert_eq!( - normalize_path_for_metrics("/v1/responses/550e8400-e29b-41d4-a716-446655440000"), - "/v1/responses/{id}" - ); - } - - #[test] - fn test_normalize_path_with_numeric_id() { - assert_eq!( - normalize_path_for_metrics("/v1/workers/12345"), - "/v1/workers/{id}" - ); - } - - #[test] - fn test_is_dynamic_id() { - // Prefixed IDs - assert!(is_dynamic_id("resp_abc123def")); - assert!(is_dynamic_id("chatcmpl_xyz789")); - assert!(!is_dynamic_id("short_id")); // Too short - - // UUIDs - assert!(is_dynamic_id("550e8400-e29b-41d4-a716-446655440000")); - assert!(is_dynamic_id("550e8400e29b41d4a716446655440000")); // No dashes - - // Numeric - assert!(is_dynamic_id("12345")); - assert!(!is_dynamic_id("")); // Empty - - // Regular words - assert!(!is_dynamic_id("completions")); - assert!(!is_dynamic_id("chat")); - } -} diff --git a/model_gateway/src/middleware/auth.rs b/model_gateway/src/middleware/auth.rs new file mode 100644 index 0000000000..70503353c9 --- /dev/null +++ b/model_gateway/src/middleware/auth.rs @@ -0,0 +1,58 @@ +//! Bearer-token auth middleware backed by a precomputed SHA-256 hash. +//! +//! The hash is compared in constant time so the comparison cost does not +//! leak the configured key length. + +use axum::{ + body::Body, + extract::{Request, State}, + http::{header, StatusCode}, + middleware::Next, + response::{IntoResponse, Response}, +}; +use sha2::{Digest, Sha256}; +use subtle::ConstantTimeEq; + +#[derive(Clone)] +pub struct AuthConfig { + /// Precomputed SHA-256 hash of the API key, used for constant-time comparison + /// that doesn't leak key length via timing. + api_key_hash: Option<[u8; 32]>, +} + +impl AuthConfig { + pub fn new(api_key: Option) -> Self { + Self { + api_key_hash: api_key.map(|k| Sha256::digest(k.as_bytes()).into()), + } + } +} + +/// Middleware to validate Bearer token against configured API key. +/// Only active when router has an API key configured. +pub async fn auth_middleware( + State(auth_config): State, + request: Request, + next: Next, +) -> Response { + if let Some(expected_hash) = &auth_config.api_key_hash { + let token = request + .headers() + .get(header::AUTHORIZATION) + .and_then(|h| h.to_str().ok()) + .and_then(|h| h.strip_prefix("Bearer ")); + + let authorized = token.is_some_and(|t| { + Sha256::digest(t.as_bytes()) + .as_slice() + .ct_eq(expected_hash) + .unwrap_u8() + == 1 + }); + if !authorized { + return StatusCode::UNAUTHORIZED.into_response(); + } + } + + next.run(request).await +} diff --git a/model_gateway/src/middleware/concurrency.rs b/model_gateway/src/middleware/concurrency.rs new file mode 100644 index 0000000000..1ca82d0feb --- /dev/null +++ b/model_gateway/src/middleware/concurrency.rs @@ -0,0 +1,301 @@ +//! Per-request concurrency limiting via a token bucket, with optional +//! queuing for backpressure. +//! +//! `ConcurrencyLimiter` wires a bounded `mpsc` channel that +//! `concurrency_limit_middleware` uses to enqueue requests when the +//! bucket is empty; `QueueProcessor` drains that channel and hands tokens +//! back to waiters. `TokenGuardBody` wraps the response body so the token +//! is only released after the entire stream has been delivered. + +use std::{ + pin::Pin, + sync::Arc, + task::{Context, Poll}, + time::{Duration, Instant}, +}; + +use axum::{ + body::Body, + extract::{Request, State}, + http::StatusCode, + middleware::Next, + response::{IntoResponse, Response}, + Json, +}; +use bytes::Bytes; +use http_body::Frame; +use serde_json::json; +use tokio::sync::{mpsc, oneshot}; +use tracing::{debug, error, warn}; + +use super::token_bucket::TokenBucket; +use crate::{ + observability::metrics::{metrics_labels, Metrics}, + server::AppState, +}; + +/// A body wrapper that holds a token and returns it when the body is fully consumed or dropped. +/// This ensures that for streaming responses, the token is only returned after the entire +/// stream has been sent to the client. +pub struct TokenGuardBody { + inner: Body, + /// The token bucket to return tokens to. Uses Option so we can take() on drop. + token_bucket: Option>, + /// Number of tokens to return. + tokens: f64, +} + +impl TokenGuardBody { + /// Create a new TokenGuardBody that will return tokens when dropped. + pub fn new(inner: Body, token_bucket: Arc, tokens: f64) -> Self { + Self { + inner, + token_bucket: Some(token_bucket), + tokens, + } + } +} + +impl Drop for TokenGuardBody { + fn drop(&mut self) { + if let Some(bucket) = self.token_bucket.take() { + debug!( + "TokenGuardBody: stream ended, returning {} tokens to bucket", + self.tokens + ); + // Use lock-free sync return - no runtime needed, guaranteed token return + bucket.return_tokens_sync(self.tokens); + } + } +} + +impl http_body::Body for TokenGuardBody { + type Data = Bytes; + type Error = axum::Error; + + fn poll_frame( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll, Self::Error>>> { + // SAFETY: We never move the inner body, and Body is Unpin + // (it's a type alias for UnsyncBoxBody which is Unpin) + let this = self.get_mut(); + Pin::new(&mut this.inner).poll_frame(cx) + } + + fn is_end_stream(&self) -> bool { + self.inner.is_end_stream() + } + + fn size_hint(&self) -> http_body::SizeHint { + self.inner.size_hint() + } +} + +/// Request queue entry +pub struct QueuedRequest { + /// Time when the request was queued + queued_at: Instant, + /// Channel to send the permit back when acquired + permit_tx: oneshot::Sender>, +} + +/// Queue processor that handles queued requests +pub struct QueueProcessor { + token_bucket: Arc, + queue_rx: mpsc::Receiver, + queue_timeout: Duration, +} + +impl QueueProcessor { + pub fn new( + token_bucket: Arc, + queue_rx: mpsc::Receiver, + queue_timeout: Duration, + ) -> Self { + Self { + token_bucket, + queue_rx, + queue_timeout, + } + } + + pub async fn run(mut self) { + debug!("Starting concurrency queue processor"); + + // Process requests in a single task to reduce overhead + while let Some(queued) = self.queue_rx.recv().await { + // Check timeout immediately + let elapsed = queued.queued_at.elapsed(); + if elapsed >= self.queue_timeout { + warn!("Request already timed out in queue"); + let _ = queued.permit_tx.send(Err(StatusCode::REQUEST_TIMEOUT)); + continue; + } + + let remaining_timeout = self.queue_timeout - elapsed; + + // Try to acquire token for this request + if self.token_bucket.try_acquire(1.0).is_ok() { + // Got token immediately + debug!("Queue: acquired token immediately for queued request"); + let _ = queued.permit_tx.send(Ok(())); + } else { + // Need to wait for token + let token_bucket = self.token_bucket.clone(); + + // Spawn task only when we actually need to wait + #[expect( + clippy::disallowed_methods, + reason = "fire-and-forget permit acquisition: task is bounded by remaining_timeout and communicates via oneshot; dropping the JoinHandle detaches the task but it self-terminates" + )] + tokio::spawn(async move { + if token_bucket + .acquire_timeout(1.0, remaining_timeout) + .await + .is_ok() + { + debug!("Queue: acquired token after waiting"); + let _ = queued.permit_tx.send(Ok(())); + } else { + warn!("Queue: request timed out waiting for token"); + let _ = queued.permit_tx.send(Err(StatusCode::REQUEST_TIMEOUT)); + } + }); + } + } + + warn!("Concurrency queue processor shutting down"); + } +} + +/// State for the concurrency limiter +pub struct ConcurrencyLimiter { + pub queue_tx: Option>, +} + +impl ConcurrencyLimiter { + /// Create new concurrency limiter with optional queue + pub fn new( + token_bucket: Option>, + queue_size: usize, + queue_timeout: Duration, + ) -> (Self, Option) { + match (token_bucket, queue_size) { + (None, _) => (Self { queue_tx: None }, None), + (Some(bucket), size) if size > 0 => { + let (queue_tx, queue_rx) = mpsc::channel(size); + let processor = QueueProcessor::new(bucket, queue_rx, queue_timeout); + ( + Self { + queue_tx: Some(queue_tx), + }, + Some(processor), + ) + } + (Some(_), _) => (Self { queue_tx: None }, None), + } + } +} + +/// Middleware function for concurrency limiting with optional queuing +pub async fn concurrency_limit_middleware( + State(app_state): State>, + request: Request, + next: Next, +) -> Response { + // Check mesh global rate limit first if mesh is enabled + // If mesh is not enabled, this check is skipped and local rate limiting is used + if let Some(mesh_handler) = &app_state.mesh_handler { + let (is_exceeded, current_count, limit) = + mesh_handler.sync_manager.check_global_rate_limit(); + if is_exceeded { + debug!( + "Global rate limit exceeded: {}/{} req/s", + current_count, limit + ); + return ( + StatusCode::TOO_MANY_REQUESTS, + Json(json!({ + "error": "Rate limit exceeded", + "current_count": current_count, + "limit": limit + })), + ) + .into_response(); + } + } + + let token_bucket = match &app_state.context.rate_limiter { + Some(bucket) => bucket.clone(), + None => { + // Rate limiting disabled, pass through immediately + return next.run(request).await; + } + }; + + // Try to acquire token immediately + if token_bucket.try_acquire(1.0).is_ok() { + debug!("Acquired token immediately"); + Metrics::record_http_rate_limit(metrics_labels::RATE_LIMIT_ALLOWED); + let response = next.run(request).await; + + // Wrap the response body with TokenGuardBody to return token when stream ends + // This ensures that for streaming responses, the token is only returned + // after the entire stream has been sent to the client. + let (parts, body) = response.into_parts(); + let guarded_body = TokenGuardBody::new(body, token_bucket, 1.0); + Response::from_parts(parts, Body::new(guarded_body)) + } else { + // No tokens available, try to queue if enabled + if let Some(queue_tx) = &app_state.concurrency_queue_tx { + debug!("No tokens available, attempting to queue request"); + + // Create a channel for the token response + let (permit_tx, permit_rx) = oneshot::channel(); + + let queued = QueuedRequest { + queued_at: Instant::now(), + permit_tx, + }; + + // Try to send to queue + match queue_tx.try_send(queued) { + Ok(()) => { + // Wait for token from queue processor + match permit_rx.await { + Ok(Ok(())) => { + debug!("Acquired token from queue"); + Metrics::record_http_rate_limit(metrics_labels::RATE_LIMIT_ALLOWED); + let response = next.run(request).await; + + // Wrap the response body with TokenGuardBody to return token when stream ends + let (parts, body) = response.into_parts(); + let guarded_body = TokenGuardBody::new(body, token_bucket, 1.0); + Response::from_parts(parts, Body::new(guarded_body)) + } + Ok(Err(status)) => { + warn!("Queue returned error status: {}", status); + Metrics::record_http_rate_limit(metrics_labels::RATE_LIMIT_REJECTED); + status.into_response() + } + Err(_) => { + error!("Queue response channel closed"); + Metrics::record_http_rate_limit(metrics_labels::RATE_LIMIT_REJECTED); + StatusCode::INTERNAL_SERVER_ERROR.into_response() + } + } + } + Err(_) => { + warn!("Request queue is full, returning 429"); + Metrics::record_http_rate_limit(metrics_labels::RATE_LIMIT_REJECTED); + StatusCode::TOO_MANY_REQUESTS.into_response() + } + } + } else { + warn!("No tokens available and queuing is disabled, returning 429"); + Metrics::record_http_rate_limit(metrics_labels::RATE_LIMIT_REJECTED); + StatusCode::TOO_MANY_REQUESTS.into_response() + } + } +} diff --git a/model_gateway/src/middleware/logging.rs b/model_gateway/src/middleware/logging.rs new file mode 100644 index 0000000000..29cda7e1b7 --- /dev/null +++ b/model_gateway/src/middleware/logging.rs @@ -0,0 +1,131 @@ +//! Tracing/logging integration for the HTTP layer. +//! +//! Wires `tower_http::trace::TraceLayer` with custom span/request/response +//! handlers that propagate W3C trace context, attach the request ID into +//! the span, and record HTTP-level metrics via the observability layer. + +use std::time::Duration; + +use axum::{extract::Request, response::Response}; +use tower_http::trace::{MakeSpan, OnRequest, OnResponse, TraceLayer}; +use tracing::{error, field::Empty, info, info_span, warn, Span}; +use tracing_opentelemetry::OpenTelemetrySpanExt; + +use super::{metrics::normalize_path_for_metrics, request_id::RequestId}; +use crate::{ + observability::{ + metrics::{method_to_static_str, Metrics}, + otel_trace::extract_trace_context_http, + }, + routers::error::extract_error_code_from_response, +}; + +/// Custom span maker that includes request ID +#[derive(Clone, Debug)] +pub struct RequestSpan; + +impl MakeSpan for RequestSpan { + fn make_span(&mut self, request: &Request) -> Span { + // Extract incoming W3C trace context (traceparent/tracestate) so that + // server-side spans become children of the caller's distributed trace. + let parent_cx = extract_trace_context_http(request.headers()); + + // Don't try to extract request ID here - it won't be available yet + // The RequestIdLayer runs after TraceLayer creates the span + let span = info_span!( + target: "smg::otel-trace", + "http_request", + method = %request.method(), + uri = %request.uri(), + version = ?request.version(), + request_id = Empty, // Will be set later + status_code = Empty, + latency = Empty, + error = Empty, + module = "smg" + ); + + span.set_parent(parent_cx); + span + } +} + +/// Custom on_request handler +#[derive(Clone, Debug)] +pub struct RequestLogger; + +impl OnRequest for RequestLogger { + fn on_request(&mut self, request: &Request, span: &Span) { + let _enter = span.enter(); + + // Try to get the request ID from extensions + // This will work if RequestIdLayer has already run + if let Some(request_id) = request.extensions().get::() { + span.record("request_id", request_id.0.as_str()); + } + + let method = method_to_static_str(request.method().as_str()); + let path = normalize_path_for_metrics(request.uri().path()); + Metrics::record_http_request(method, &path); + + // Log the request start + info!( + target: "smg::request", + "started processing request" + ); + } +} + +/// Custom on_response handler +#[derive(Clone, Debug, Default)] +pub struct ResponseLogger; + +impl OnResponse for ResponseLogger { + fn on_response(self, response: &Response, latency: Duration, span: &Span) { + let status = response.status(); + let status_code = status.as_u16(); + + let error_code = extract_error_code_from_response(response); + + // Layer 1: HTTP metrics + Metrics::record_http_response(status_code, error_code); + + // Record these in the span for structured logging/observability tools + span.record("status_code", status_code); + // Use microseconds as integer to avoid format! string allocation + span.record("latency", latency.as_micros() as u64); + + // Log the response completion + let _enter = span.enter(); + if status.is_server_error() { + error!( + target: "smg::response", + "request failed with server error" + ); + } else if status.is_client_error() { + warn!( + target: "smg::response", + "request failed with client error" + ); + } else { + info!( + target: "smg::response", + "finished processing request" + ); + } + } +} + +/// Create a configured TraceLayer for HTTP logging +/// Note: Actual request/response logging with request IDs is done in RequestIdService +pub fn create_logging_layer() -> TraceLayer< + tower_http::classify::SharedClassifier, + RequestSpan, + RequestLogger, + ResponseLogger, +> { + TraceLayer::new_for_http() + .make_span_with(RequestSpan) + .on_request(RequestLogger) + .on_response(ResponseLogger) +} diff --git a/model_gateway/src/middleware/metrics.rs b/model_gateway/src/middleware/metrics.rs new file mode 100644 index 0000000000..5e53dfcfcf --- /dev/null +++ b/model_gateway/src/middleware/metrics.rs @@ -0,0 +1,221 @@ +//! HTTP metrics collection (SMG Layer 1 metrics). +//! +//! `HttpMetricsLayer` wraps the inner service to record per-request +//! duration plus the in-flight connection count via +//! `InFlightRequestTracker`. The path is normalized to bound metric +//! cardinality. + +use std::{ + pin::Pin, + sync::Arc, + task::{Context, Poll}, + time::Instant, +}; + +use axum::{extract::Request, response::Response}; +use tower::{Layer, Service}; + +use crate::observability::{ + inflight_tracker::InFlightRequestTracker, + metrics::{method_to_static_str, Metrics}, +}; + +/// Tower Layer for HTTP metrics collection (SMG Layer 1 metrics) +#[derive(Clone)] +pub struct HttpMetricsLayer { + tracker: Arc, +} + +impl HttpMetricsLayer { + pub fn new(tracker: Arc) -> Self { + Self { tracker } + } +} + +impl Layer for HttpMetricsLayer { + type Service = HttpMetricsMiddleware; + + fn layer(&self, inner: S) -> Self::Service { + HttpMetricsMiddleware { + inner, + in_flight_request_tracker: self.tracker.clone(), + } + } +} + +/// Tower Service for HTTP metrics collection +#[derive(Clone)] +pub struct HttpMetricsMiddleware { + inner: S, + in_flight_request_tracker: Arc, +} + +impl Service for HttpMetricsMiddleware +where + S: Service + Send + Clone + 'static, + S::Future: Send + 'static, +{ + type Response = S::Response; + type Error = S::Error; + type Future = + Pin> + Send>>; + + fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { + self.inner.poll_ready(cx) + } + + fn call(&mut self, req: Request) -> Self::Future { + // Convert method to static string to avoid allocation + let method = method_to_static_str(req.method().as_str()); + let path = normalize_path_for_metrics(req.uri().path()); + let start = Instant::now(); + + let mut inner = self.inner.clone(); + let in_flight_request_tracker = self.in_flight_request_tracker.clone(); + + Box::pin(async move { + let guard = in_flight_request_tracker.track(); + Metrics::set_http_connections_active(in_flight_request_tracker.len()); + + // Capture result before dropping guard to ensure decrement happens on error too + let result = inner.call(req).await; + + drop(guard); + Metrics::set_http_connections_active(in_flight_request_tracker.len()); + + let response = result?; + + let duration = start.elapsed(); + Metrics::record_http_duration(method, &path, duration); + + Ok(response) + }) + } +} + +/// Normalize path for metrics to avoid high cardinality. +/// Replaces dynamic segments (IDs, UUIDs) with `{id}` placeholder. +/// Only allocates when normalization is needed; uses single-pass with byte offsets. +pub(super) fn normalize_path_for_metrics(path: &str) -> String { + let bytes = path.as_bytes(); + let mut segment_start = 0; + let mut segment_idx = 0; + let mut result: Option = None; + + for (pos, &b) in bytes.iter().enumerate() { + if b == b'/' || pos == bytes.len() - 1 { + // Determine segment end (include last char if not a slash) + let segment_end = if b == b'/' { pos } else { pos + 1 }; + let segment = &path[segment_start..segment_end]; + + // Check segments after index 2 for dynamic IDs + if segment_idx > 2 && !segment.is_empty() && is_dynamic_id(segment) { + // Initialize result with everything before this segment + let result = result.get_or_insert_with(|| { + let mut s = String::with_capacity(path.len()); + s.push_str(&path[..segment_start]); + s + }); + result.push_str("{id}"); + } else if let Some(ref mut r) = result { + // Already normalizing, append this segment as-is + r.push_str(segment); + } + + // Add slash after segment (except at end) + if b == b'/' { + if let Some(ref mut r) = result { + r.push('/'); + } + segment_start = pos + 1; + segment_idx += 1; + } + } + } + + result.unwrap_or_else(|| path.to_owned()) +} + +/// Check if segment looks like a dynamic ID (prefixed ID, UUID, or numeric). +#[inline] +fn is_dynamic_id(s: &str) -> bool { + // Prefixed IDs: resp_xxx, chatcmpl_xxx (len > 10 with underscore) + if s.len() > 10 && s.contains('_') { + return true; + } + // UUIDs: 32+ hex chars with dashes + if s.len() >= 32 && s.bytes().all(|b| b.is_ascii_hexdigit() || b == b'-') { + return true; + } + // Numeric IDs + !s.is_empty() && s.bytes().all(|b| b.is_ascii_digit()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_normalize_path_no_ids() { + // Common API paths should pass through unchanged + assert_eq!( + normalize_path_for_metrics("/v1/chat/completions"), + "/v1/chat/completions" + ); + assert_eq!( + normalize_path_for_metrics("/v1/completions"), + "/v1/completions" + ); + assert_eq!(normalize_path_for_metrics("/v1/models"), "/v1/models"); + assert_eq!(normalize_path_for_metrics("/health"), "/health"); + } + + #[test] + fn test_normalize_path_with_prefixed_id() { + // Prefixed IDs (resp_xxx, chatcmpl_xxx) should be normalized + assert_eq!( + normalize_path_for_metrics("/v1/responses/resp_abc123def456"), + "/v1/responses/{id}" + ); + assert_eq!( + normalize_path_for_metrics("/v1/chat/completions/chatcmpl_abc123xyz"), + "/v1/chat/completions/{id}" + ); + } + + #[test] + fn test_normalize_path_with_uuid() { + assert_eq!( + normalize_path_for_metrics("/v1/responses/550e8400-e29b-41d4-a716-446655440000"), + "/v1/responses/{id}" + ); + } + + #[test] + fn test_normalize_path_with_numeric_id() { + assert_eq!( + normalize_path_for_metrics("/v1/workers/12345"), + "/v1/workers/{id}" + ); + } + + #[test] + fn test_is_dynamic_id() { + // Prefixed IDs + assert!(is_dynamic_id("resp_abc123def")); + assert!(is_dynamic_id("chatcmpl_xyz789")); + assert!(!is_dynamic_id("short_id")); // Too short + + // UUIDs + assert!(is_dynamic_id("550e8400-e29b-41d4-a716-446655440000")); + assert!(is_dynamic_id("550e8400e29b41d4a716446655440000")); // No dashes + + // Numeric + assert!(is_dynamic_id("12345")); + assert!(!is_dynamic_id("")); // Empty + + // Regular words + assert!(!is_dynamic_id("completions")); + assert!(!is_dynamic_id("chat")); + } +} diff --git a/model_gateway/src/middleware/mod.rs b/model_gateway/src/middleware/mod.rs new file mode 100644 index 0000000000..84cc4395c5 --- /dev/null +++ b/model_gateway/src/middleware/mod.rs @@ -0,0 +1,25 @@ +//! HTTP middleware stack for the gateway server. +//! +//! Each submodule owns a single concern; this `mod.rs` re-exports the +//! types and functions that `server.rs` and other call sites already +//! reference, so this split is invisible to downstream callers. + +pub mod auth; +pub mod concurrency; +pub mod logging; +pub mod metrics; +pub mod request_id; +pub mod storage_context; +pub mod token_bucket; +pub mod wasm; + +pub use auth::{auth_middleware, AuthConfig}; +pub use concurrency::{ + concurrency_limit_middleware, ConcurrencyLimiter, QueueProcessor, QueuedRequest, TokenGuardBody, +}; +pub use logging::{create_logging_layer, RequestLogger, RequestSpan, ResponseLogger}; +pub use metrics::{HttpMetricsLayer, HttpMetricsMiddleware}; +pub use request_id::{RequestId, RequestIdLayer, RequestIdMiddleware}; +pub use storage_context::storage_context_middleware; +pub use token_bucket::TokenBucket; +pub use wasm::wasm_middleware; diff --git a/model_gateway/src/middleware/request_id.rs b/model_gateway/src/middleware/request_id.rs new file mode 100644 index 0000000000..0f3df0420d --- /dev/null +++ b/model_gateway/src/middleware/request_id.rs @@ -0,0 +1,129 @@ +//! Request ID middleware — pulls an existing request ID from configured +//! headers or generates an OpenAI-compatible one, stores it in extensions, +//! and echoes it back via `x-request-id`. + +use std::{ + pin::Pin, + sync::Arc, + task::{Context, Poll}, +}; + +use axum::{extract::Request, http::HeaderValue, response::Response}; +use rand::Rng; +// Re-export RequestId from auth crate for backward compatibility +pub use smg_auth::RequestId; +use tower::{Layer, Service}; + +/// Alphanumeric characters for request ID generation (as bytes for O(1) indexing) +const REQUEST_ID_CHARS: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789"; + +/// Generate OpenAI-compatible request ID based on endpoint. +pub(super) fn generate_request_id(path: &str) -> String { + let prefix = if path.contains("/chat/completions") { + "chatcmpl-" + } else if path.contains("/completions") { + "cmpl-" + } else if path.contains("/generate") { + "gnt-" + } else if path.contains("/responses") { + "resp-" + } else { + "req-" + }; + + // Generate a random string similar to OpenAI's format + // Use byte array indexing (O(1)) instead of chars().nth() (O(n)) + let mut rng = rand::rng(); + let random_part: String = (0..24) + .map(|_| { + let idx = rng.random_range(0..REQUEST_ID_CHARS.len()); + REQUEST_ID_CHARS[idx] as char + }) + .collect(); + + format!("{prefix}{random_part}") +} + +/// Tower Layer for request ID middleware +#[derive(Clone)] +pub struct RequestIdLayer { + headers: Arc>, +} + +impl RequestIdLayer { + pub fn new(headers: Vec) -> Self { + Self { + headers: Arc::new(headers), + } + } +} + +impl Layer for RequestIdLayer { + type Service = RequestIdMiddleware; + + fn layer(&self, inner: S) -> Self::Service { + RequestIdMiddleware { + inner, + headers: self.headers.clone(), + } + } +} + +/// Tower Service for request ID middleware +#[derive(Clone)] +pub struct RequestIdMiddleware { + inner: S, + headers: Arc>, +} + +impl Service for RequestIdMiddleware +where + S: Service + Send + 'static, + S::Future: Send + 'static, +{ + type Response = S::Response; + type Error = S::Error; + type Future = + Pin> + Send>>; + + fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { + self.inner.poll_ready(cx) + } + + fn call(&mut self, mut req: Request) -> Self::Future { + let headers = self.headers.clone(); + + // Extract request ID from headers or generate new one + let mut request_id = None; + + for header_name in headers.iter() { + if let Some(header_value) = req.headers().get(header_name) { + if let Ok(value) = header_value.to_str() { + request_id = Some(value.to_string()); + break; + } + } + } + + let request_id = request_id.unwrap_or_else(|| generate_request_id(req.uri().path())); + + // Insert request ID into request extensions for other middleware/handlers to use + req.extensions_mut().insert(RequestId(request_id.clone())); + + // Call the inner service + let future = self.inner.call(req); + + Box::pin(async move { + let mut response = future.await?; + + // Add request ID to response headers + response.headers_mut().insert( + "x-request-id", + HeaderValue::from_str(&request_id) + .unwrap_or_else(|_| HeaderValue::from_static("invalid-request-id")), + ); + + Ok(response) + }) + } +} diff --git a/model_gateway/src/middleware/storage_context.rs b/model_gateway/src/middleware/storage_context.rs new file mode 100644 index 0000000000..94a748691f --- /dev/null +++ b/model_gateway/src/middleware/storage_context.rs @@ -0,0 +1,114 @@ +//! Storage request-context middleware. +//! +//! Maps configured request headers onto a [`StorageRequestContext`] so the +//! storage layer can read tenant/user/etc. fields without re-parsing headers. + +use std::sync::Arc; + +use axum::{ + body::Body, + extract::{Request, State}, + middleware::Next, + response::Response, +}; +use smg_data_connector::{with_request_context, RequestContext as StorageRequestContext}; + +use crate::{config::RouterConfig, server::AppState}; + +fn extract_header_str(headers: &http::HeaderMap, name: &str) -> Option { + headers + .get(name) + .and_then(|value| value.to_str().ok()) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned) +} + +fn build_storage_request_context( + config: &RouterConfig, + headers: &http::HeaderMap, +) -> Option { + let mut ctx = StorageRequestContext::new(); + + for (header_name, context_key) in &config.storage_context_headers { + let header_name = header_name.trim(); + let context_key = context_key.trim(); + + if header_name.is_empty() || context_key.is_empty() { + continue; + } + + if let Some(value) = extract_header_str(headers, header_name) { + ctx.set(context_key, value); + } + } + + (!ctx.data().is_empty()).then_some(ctx) +} + +pub async fn storage_context_middleware( + State(state): State>, + request: Request, + next: Next, +) -> Response { + if state + .context + .router_config + .storage_context_headers + .is_empty() + { + return next.run(request).await; + } + + match build_storage_request_context(&state.context.router_config, request.headers()) { + Some(ctx) => with_request_context(ctx, next.run(request)).await, + None => next.run(request).await, + } +} + +#[cfg(test)] +mod tests { + use axum::http::{HeaderMap, HeaderValue}; + + use super::*; + + #[test] + fn build_storage_request_context_maps_configured_headers() { + let config = RouterConfig { + storage_context_headers: std::collections::HashMap::from([ + ("x-tenant-id".to_string(), "tenant_id".to_string()), + ("x-user-id".to_string(), "user_id".to_string()), + ]), + ..Default::default() + }; + + let mut headers = HeaderMap::new(); + headers.insert("x-tenant-id", HeaderValue::from_static("tenant-abc")); + headers.insert("x-user-id", HeaderValue::from_static("user-123")); + + let ctx = build_storage_request_context(&config, &headers).unwrap(); + + assert_eq!(ctx.get("tenant_id"), Some("tenant-abc")); + assert_eq!(ctx.get("user_id"), Some("user-123")); + } + + #[test] + fn build_storage_request_context_ignores_empty_entries_and_missing_headers() { + let config = RouterConfig { + storage_context_headers: std::collections::HashMap::from([ + (" ".to_string(), "tenant_id".to_string()), + ("x-empty-key".to_string(), " ".to_string()), + ("x-present".to_string(), "present_key".to_string()), + ]), + ..Default::default() + }; + + let mut headers = HeaderMap::new(); + headers.insert("x-present", HeaderValue::from_static(" keep-me ")); + + let ctx = build_storage_request_context(&config, &headers).unwrap(); + + assert_eq!(ctx.get("present_key"), Some("keep-me")); + assert_eq!(ctx.data().len(), 1); + } +} diff --git a/model_gateway/src/worker/token_bucket.rs b/model_gateway/src/middleware/token_bucket.rs similarity index 100% rename from model_gateway/src/worker/token_bucket.rs rename to model_gateway/src/middleware/token_bucket.rs diff --git a/model_gateway/src/middleware/wasm.rs b/model_gateway/src/middleware/wasm.rs new file mode 100644 index 0000000000..a83c50fa5a --- /dev/null +++ b/model_gateway/src/middleware/wasm.rs @@ -0,0 +1,260 @@ +//! WASM plugin middleware. +//! +//! Dispatches request and response through every WASM module attached at +//! the corresponding `Middleware::OnRequest` / `Middleware::OnResponse` +//! point. Streaming responses skip the OnResponse phase to avoid buffering +//! arbitrary bodies into memory. + +use std::{sync::Arc, time::Duration}; + +use axum::{ + body::Body, + extract::{Request, State}, + http::{header, StatusCode}, + middleware::Next, + response::{IntoResponse, Response}, + Json, +}; +use serde_json::json; +use tracing::{error, warn}; + +use super::request_id::{generate_request_id, RequestId}; +use crate::{ + server::AppState, + wasm::{ + module::{MiddlewareAttachPoint, WasmModuleAttachPoint}, + spec::{ + apply_modify_action_to_headers, build_wasm_headers_from_axum_headers, + smg::gateway::middleware_types::{ + Action, Request as WasmRequest, Response as WasmResponse, + }, + }, + types::WasmComponentInput, + }, +}; + +pub async fn wasm_middleware( + State(app_state): State>, + request: Request, + next: Next, +) -> Response { + // Check if WASM is enabled + if !app_state.context.router_config.enable_wasm { + return next.run(request).await; + } + + // Get WASM manager + let wasm_manager = match &app_state.context.wasm_manager { + Some(manager) => manager, + None => { + return next.run(request).await; + } + }; + + // Get request ID from extensions or generate one + let request_id = request + .extensions() + .get::() + .map(|r| r.0.clone()) + .unwrap_or_else(|| generate_request_id(request.uri().path())); + + // ===== OnRequest Phase ===== + let on_request_attach_point = + WasmModuleAttachPoint::Middleware(MiddlewareAttachPoint::OnRequest); + + let modules_on_request = + match wasm_manager.get_modules_by_attach_point(on_request_attach_point.clone()) { + Ok(modules) => modules, + Err(e) => { + error!("Failed to get WASM modules for OnRequest: {}", e); + return next.run(request).await; + } + }; + + let response = if modules_on_request.is_empty() { + next.run(request).await + } else { + // Decompose request to preserve extensions across reconstruction + let (parts, body) = request.into_parts(); + let method = parts.method; + let uri = parts.uri; + let mut headers = parts.headers; + let extensions = parts.extensions; + + let max_body_size = wasm_manager.get_max_body_size(); + let body_bytes = match axum::body::to_bytes(body, max_body_size).await { + Ok(bytes) => bytes.to_vec(), + Err(e) => { + error!("Failed to read request body for WASM processing: {}", e); + return ( + StatusCode::BAD_REQUEST, + Json(json!({"error": format!("Failed to read request body: {e}")})), + ) + .into_response(); + } + }; + + // Process each OnRequest module + let mut modified_body = body_bytes; + + // Pre-compute strings once before the loop to avoid repeated allocations + let method_str = method.to_string(); + let path_str = uri.path().to_string(); + let query_str = uri.query().unwrap_or("").to_string(); + + for module in modules_on_request { + let wasm_headers = build_wasm_headers_from_axum_headers(&headers); + let wasm_request = WasmRequest { + method: method_str.clone(), + path: path_str.clone(), + query: query_str.clone(), + headers: wasm_headers, + body: modified_body.clone(), + request_id: request_id.clone(), + now_epoch_ms: std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_else(|_| Duration::from_millis(0)) + .as_millis() as u64, + }; + + let action = match wasm_manager + .execute_module_for_attach_point( + &module, + on_request_attach_point.clone(), + WasmComponentInput::MiddlewareRequest(wasm_request), + ) + .await + { + Some(action) => action, + None => continue, + }; + + match action { + Action::Continue => {} + Action::Reject(status) => { + return StatusCode::from_u16(status) + .unwrap_or(StatusCode::BAD_REQUEST) + .into_response(); + } + Action::Modify(modify) => { + apply_modify_action_to_headers(&mut headers, &modify); + if let Some(body_bytes) = modify.body_replace { + modified_body = body_bytes; + } + } + } + } + + // Reconstruct request with modifications, preserving original extensions + let mut final_request = Request::builder() + .method(method) + .uri(uri) + .body(Body::from(modified_body)) + .unwrap_or_else(|_| Request::new(Body::empty())); + *final_request.headers_mut() = headers; + *final_request.extensions_mut() = extensions; + + next.run(final_request).await + }; + + // ===== OnResponse Phase ===== + let on_response_attach_point = + WasmModuleAttachPoint::Middleware(MiddlewareAttachPoint::OnResponse); + + let modules_on_response = + match wasm_manager.get_modules_by_attach_point(on_response_attach_point.clone()) { + Ok(modules) => modules, + Err(e) => { + error!("Failed to get WASM modules for OnResponse: {}", e); + return response; + } + }; + if modules_on_response.is_empty() { + return response; + } + + // Skip WASM OnResponse processing for streaming responses to avoid + // buffering the entire stream into memory (breaks SSE, causes OOM on large streams). + let is_streaming = response + .headers() + .get(header::CONTENT_TYPE) + .and_then(|v| v.to_str().ok()) + .is_some_and(|ct| ct.contains("text/event-stream") || ct.contains("application/x-ndjson")) + || response + .headers() + .get(header::TRANSFER_ENCODING) + .and_then(|v| v.to_str().ok()) + .is_some_and(|te| te.contains("chunked")); + if is_streaming { + warn!("Skipping WASM OnResponse for streaming response; OnResponse modules do not apply to streaming"); + return response; + } + + // Extract response data once before processing modules + let mut status = response.status(); + let mut headers = response.headers().clone(); + let max_body_size = wasm_manager.get_max_body_size(); + let mut body_bytes = match axum::body::to_bytes(response.into_body(), max_body_size).await { + Ok(bytes) => bytes.to_vec(), + Err(e) => { + error!("Failed to read response body for WASM processing: {}", e); + return ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(json!({"error": "Failed to read response body"})), + ) + .into_response(); + } + }; + + // Process each OnResponse module + for module in modules_on_response { + let wasm_headers = build_wasm_headers_from_axum_headers(&headers); + let wasm_response = WasmResponse { + status: status.as_u16(), + headers: wasm_headers, + body: body_bytes.clone(), + }; + + let action = match wasm_manager + .execute_module_for_attach_point( + &module, + on_response_attach_point.clone(), + WasmComponentInput::MiddlewareResponse(wasm_response), + ) + .await + { + Some(action) => action, + None => continue, + }; + + match action { + Action::Continue => {} + Action::Reject(status_code) => { + status = StatusCode::from_u16(status_code).unwrap_or(StatusCode::BAD_REQUEST); + let mut final_response = Response::builder() + .status(status) + .body(Body::from(body_bytes)) + .unwrap_or_else(|_| Response::new(Body::empty())); + *final_response.headers_mut() = headers; + return final_response; + } + Action::Modify(modify) => { + if let Some(new_status) = modify.status { + status = StatusCode::from_u16(new_status).unwrap_or(status); + } + apply_modify_action_to_headers(&mut headers, &modify); + if let Some(new_body) = modify.body_replace { + body_bytes = new_body; + } + } + } + } + + // Reconstruct final response with all modifications + let mut final_response = Response::builder() + .status(status) + .body(Body::from(body_bytes)) + .unwrap_or_else(|_| Response::new(Body::empty())); + *final_response.headers_mut() = headers; + final_response +} diff --git a/model_gateway/src/worker/mod.rs b/model_gateway/src/worker/mod.rs index 024ce6f4dd..30b2b7c931 100644 --- a/model_gateway/src/worker/mod.rs +++ b/model_gateway/src/worker/mod.rs @@ -13,7 +13,6 @@ pub mod monitor; pub mod registry; pub mod resilience; pub mod service; -pub mod token_bucket; // FIXME: worker.rs is a 1800-line monolith containing the Worker trait, // BasicWorker impl, HealthChecker, WorkerType, ConnectionMode, and more. // Break it apart into focused modules (e.g. health_checker.rs, types.rs).