From 66dfb8c156adf5f4225414c4ccafae3ce0d51789 Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Sun, 4 Jan 2026 01:43:02 +0800 Subject: [PATCH] Tiny fix non-PD router http header missing whitelist (#16339) Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- .../py_test/fixtures/mock_worker.py | 3 ++ .../test_header_forwarding.py | 36 +++++++++++++ sgl-model-gateway/src/routers/header_utils.rs | 51 +++++++++++++++++++ .../src/routers/http/pd_router.rs | 12 +---- sgl-model-gateway/src/routers/http/router.rs | 21 ++------ 5 files changed, 95 insertions(+), 28 deletions(-) create mode 100644 sgl-model-gateway/py_test/integration_mock/test_header_forwarding.py diff --git a/sgl-model-gateway/py_test/fixtures/mock_worker.py b/sgl-model-gateway/py_test/fixtures/mock_worker.py index 2f3022517..ec2b16d5e 100644 --- a/sgl-model-gateway/py_test/fixtures/mock_worker.py +++ b/sgl-model-gateway/py_test/fixtures/mock_worker.py @@ -203,6 +203,8 @@ def create_app(args: argparse.Namespace) -> FastAPI: except (json.JSONDecodeError, ValueError): data = {} + received_headers = {k.lower(): v for k, v in request.headers.items()} + now = time.time() ret = { "id": f"cmpl-{int(now*1000)}", @@ -218,6 +220,7 @@ def create_app(args: argparse.Namespace) -> FastAPI: ], "worker_id": worker_id, "echo": data, + "received_headers": received_headers, } return make_json_response(ret, status_code=200) diff --git a/sgl-model-gateway/py_test/integration_mock/test_header_forwarding.py b/sgl-model-gateway/py_test/integration_mock/test_header_forwarding.py new file mode 100644 index 000000000..11556ad7e --- /dev/null +++ b/sgl-model-gateway/py_test/integration_mock/test_header_forwarding.py @@ -0,0 +1,36 @@ +import pytest +import requests + + +@pytest.mark.integration +def test_header_forwarding_whitelist(mock_workers, router_manager): + _, urls, _ = mock_workers(n=1) + rh = router_manager.start_router(worker_urls=urls) + + with requests.Session() as s: + r = s.post( + f"{rh.url}/v1/completions", + json={"model": "test", "prompt": "hi", "max_tokens": 1, "stream": False}, + headers={ + "Authorization": "Bearer test-token", + "X-SMG-Routing-Key": "routing-123", + "X-Request-Id": "req-456", + "X-Correlation-Id": "corr-789", + "traceparent": "00-trace-span-01", + "tracestate": "vendor=value", + "X-Custom-Header": "should-not-forward", + "Cookie": "session=abc", + }, + ) + assert r.status_code == 200 + h = r.json().get("received_headers", {}) + + assert h.get("authorization") == "Bearer test-token" + assert h.get("x-request-id") == "req-456" + assert h.get("x-correlation-id") == "corr-789" + assert h.get("traceparent") == "00-trace-span-01" + assert h.get("tracestate") == "vendor=value" + + assert "x-smg-routing-key" not in h + assert "x-custom-header" not in h + assert "cookie" not in h diff --git a/sgl-model-gateway/src/routers/header_utils.rs b/sgl-model-gateway/src/routers/header_utils.rs index b7ca861f6..3cabc478a 100644 --- a/sgl-model-gateway/src/routers/header_utils.rs +++ b/sgl-model-gateway/src/routers/header_utils.rs @@ -183,3 +183,54 @@ pub fn extract_auth_header( .and_then(|k| HeaderValue::from_str(&format!("Bearer {}", k)).ok()) }) } + +#[inline] +pub fn should_forward_request_header(name: &str) -> bool { + let lower_name = name.to_ascii_lowercase(); + matches!( + lower_name.as_str(), + "authorization" | "x-request-id" | "x-correlation-id" | "traceparent" | "tracestate" + ) || lower_name.starts_with("x-request-id-") +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_should_forward_request_header_whitelist() { + assert!(should_forward_request_header("authorization")); + assert!(should_forward_request_header("Authorization")); + assert!(should_forward_request_header("AUTHORIZATION")); + assert!(should_forward_request_header("x-request-id")); + assert!(should_forward_request_header("X-Request-Id")); + assert!(should_forward_request_header("x-correlation-id")); + assert!(should_forward_request_header("X-Correlation-ID")); + assert!(should_forward_request_header("traceparent")); + assert!(should_forward_request_header("Traceparent")); + assert!(should_forward_request_header("tracestate")); + assert!(should_forward_request_header("Tracestate")); + assert!(should_forward_request_header("x-request-id-user")); + assert!(should_forward_request_header("X-Request-ID-Span")); + assert!(should_forward_request_header("x-request-id-123")); + } + + #[test] + fn test_should_forward_request_header_blocked() { + assert!(!should_forward_request_header("content-type")); + assert!(!should_forward_request_header("Content-Type")); + assert!(!should_forward_request_header("content-length")); + assert!(!should_forward_request_header("host")); + assert!(!should_forward_request_header("Host")); + assert!(!should_forward_request_header("connection")); + assert!(!should_forward_request_header("transfer-encoding")); + assert!(!should_forward_request_header("accept")); + assert!(!should_forward_request_header("accept-encoding")); + assert!(!should_forward_request_header("user-agent")); + assert!(!should_forward_request_header("cookie")); + assert!(!should_forward_request_header("x-custom-header")); + assert!(!should_forward_request_header("x-api-key")); + assert!(!should_forward_request_header("x-smg-routing-key")); + assert!(!should_forward_request_header("X-SMG-Routing-Key")); + } +} diff --git a/sgl-model-gateway/src/routers/http/pd_router.rs b/sgl-model-gateway/src/routers/http/pd_router.rs index fd3d17252..8fb73b05c 100644 --- a/sgl-model-gateway/src/routers/http/pd_router.rs +++ b/sgl-model-gateway/src/routers/http/pd_router.rs @@ -1045,17 +1045,7 @@ impl PDRouter { } if let Some(headers) = headers { for (name, value) in headers.iter() { - let name_lc = name.as_str().to_ascii_lowercase(); - // Whitelist important end-to-end headers, skip hop-by-hop - let forward = matches!( - name_lc.as_str(), - "authorization" - | "x-request-id" - | "x-correlation-id" - | "traceparent" // W3C Trace Context - | "tracestate" // W3C Trace Context - ) || name_lc.starts_with("x-request-id-"); - if forward { + if header_utils::should_forward_request_header(name.as_str()) { if let Ok(val) = value.to_str() { request = request.header(name, val); } diff --git a/sgl-model-gateway/src/routers/http/router.rs b/sgl-model-gateway/src/routers/http/router.rs index 2d78a05e0..9da22c7f0 100644 --- a/sgl-model-gateway/src/routers/http/router.rs +++ b/sgl-model-gateway/src/routers/http/router.rs @@ -3,10 +3,7 @@ use std::{sync::Arc, time::Instant}; use axum::{ body::{to_bytes, Body}, extract::Request, - http::{ - header::{CONTENT_LENGTH, CONTENT_TYPE}, - HeaderMap, HeaderValue, Method, StatusCode, - }, + http::{header::CONTENT_TYPE, HeaderMap, HeaderValue, Method, StatusCode}, response::{IntoResponse, Response}, Json, }; @@ -91,7 +88,6 @@ impl Router { } } - // Helper method to proxy GET requests to the first available worker async fn proxy_get_request(&self, req: Request, endpoint: &str) -> Response { let headers = header_utils::copy_request_headers(&req); @@ -99,10 +95,7 @@ impl Router { Ok(worker_url) => { let mut request_builder = self.client.get(format!("{}/{}", worker_url, endpoint)); for (name, value) in headers { - // Use eq_ignore_ascii_case to avoid string allocation - if !name.eq_ignore_ascii_case("content-type") - && !name.eq_ignore_ascii_case("content-length") - { + if header_utils::should_forward_request_header(&name) { request_builder = request_builder.header(name, value); } } @@ -361,14 +354,10 @@ impl Router { return error::service_unavailable("no_workers", "No available workers"); } - // Pre-filter headers once before the loop to avoid repeated lowercasing let filtered_headers: Vec<_> = headers .map(|hdrs| { hdrs.iter() - .filter(|(name, _)| { - !name.as_str().eq_ignore_ascii_case("content-type") - && !name.as_str().eq_ignore_ascii_case("content-length") - }) + .filter(|(name, _)| header_utils::should_forward_request_header(name.as_str())) .collect() }) .unwrap_or_default(); @@ -542,11 +531,9 @@ impl Router { request_builder = request_builder.header("Authorization", auth_header); } - // Copy all headers from original request if provided if let Some(headers) = headers { for (name, value) in headers { - // Skip Content-Type and Content-Length as .json() sets them - if *name != CONTENT_TYPE && *name != CONTENT_LENGTH { + if header_utils::should_forward_request_header(name.as_str()) { request_builder = request_builder.header(name, value); } }