Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
from dstack._internal.utils.logging import get_logger

logger = get_logger(__name__)
UVICORN_AUTOMATIC_HEADERS = ("Server", "Date")


async def proxy(
Expand Down Expand Up @@ -73,7 +74,7 @@ async def proxy(
return fastapi.responses.StreamingResponse(
stream_response(upstream_response),
status_code=upstream_response.status_code,
headers=upstream_response.headers,
headers=clean_response_headers(upstream_response.headers),
)


Expand All @@ -87,6 +88,14 @@ def _is_whitelisted_path(path: str, whitelisted_paths: tuple[str, ...]) -> bool:
return False


def clean_response_headers(headers: httpx.Headers) -> httpx.Headers:
headers = httpx.Headers(headers) # copy
for header in UVICORN_AUTOMATIC_HEADERS:
if header in headers:
del headers[header]
return headers


async def stream_response(response: httpx.Response) -> AsyncGenerator[bytes, None]:
try:
async for chunk in response.aiter_raw():
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -88,7 +88,6 @@ async def test_proxy(mock_replica_client_httpbin, method: str) -> None:
content=req_body,
)
assert resp.status_code == 200
assert resp.headers["server"].startswith("Pytest-HTTPBIN")
resp_body = resp.json()
assert resp_body["url"] == f"http://test-host:8888/{method}?a=b&c="
assert resp_body["args"] == {"a": "b", "c": ""}
Expand Down Expand Up @@ -201,6 +200,46 @@ async def test_redirect_to_service_root(mock_replica_client_httpbin) -> None:
assert resp.request.url == url + "/"


@pytest.mark.asyncio
@pytest.mark.parametrize(
"response_headers",
[
pytest.param(
{
"X-Custom-Header": "1",
"Server": "test",
"Date": "Mon, 11 May 2026 00:00:00 GMT",
},
id="mixed-case",
),
pytest.param(
{
"x-custom-header": "1",
"server": "test",
"date": "Mon, 11 May 2026 00:00:00 GMT",
},
id="lower-case",
),
],
)
async def test_drop_uvicorn_headers(
mock_replica_client_httpbin, response_headers: dict[str, str]
) -> None:
repo = ProxyTestRepo()
await repo.set_project(make_project("test-proj"))
await repo.set_service(make_service("test-proj", "httpbin"))
_, client = make_app_client(repo)
resp = await client.post(
"http://test-host/proxy/services/test-proj/httpbin/response-headers",
params=response_headers,
)
assert resp.status_code == 200
assert "X-Custom-Header" in resp.headers
# These should be stripped by the proxy, as they are then set by uvicorn
assert "Server" not in resp.headers
assert "Date" not in resp.headers


@pytest.mark.asyncio
@pytest.mark.parametrize(
("token", "status"), [("correct-token", 200), ("incorrect-token", 403), ("", 403), (None, 403)]
Expand Down
Loading