diff --git a/gateway/platforms/yuanbao_media.py b/gateway/platforms/yuanbao_media.py index 85abb30494a8..e876d620773f 100644 --- a/gateway/platforms/yuanbao_media.py +++ b/gateway/platforms/yuanbao_media.py @@ -220,18 +220,19 @@ async def download_url( # SSRF protection: yuanbao downloads model-supplied and inbound URLs # server-side. Reject private/internal targets up front, and re-validate # every redirect hop so a public URL can't 302 to http://169.254.169.254/. - from tools.url_safety import is_safe_url + from tools.url_safety import is_safe_url, redirect_target_from_response if not is_safe_url(url): raise ValueError(f"Blocked unsafe URL (SSRF protection): {url}") async def _redirect_guard(response: httpx.Response) -> None: - if response.is_redirect and response.next_request: - redirect_url = str(response.next_request.url) - if not is_safe_url(redirect_url): - raise ValueError( - f"Blocked redirect to private/internal address: {redirect_url}" - ) + # response.next_request is frequently None inside an httpx response hook + # even for a genuine 302, so resolve the hop from the Location header. + redirect_url = redirect_target_from_response(response) + if redirect_url and not is_safe_url(redirect_url): + raise ValueError( + f"Blocked redirect to private/internal address: {redirect_url}" + ) max_bytes = max_size_mb * 1024 * 1024 async with httpx.AsyncClient( diff --git a/tests/gateway/test_yuanbao_media_ssrf.py b/tests/gateway/test_yuanbao_media_ssrf.py index 329a18e7c920..a8217913cd08 100644 --- a/tests/gateway/test_yuanbao_media_ssrf.py +++ b/tests/gateway/test_yuanbao_media_ssrf.py @@ -90,3 +90,71 @@ def stream(self, method, url, **kw): # The guarded client must register a redirect event hook. assert fetched["hooks"] is not None assert "response" in fetched["hooks"] + + @pytest.mark.asyncio + async def test_redirect_to_metadata_blocked(self, monkeypatch): + """A 302 whose Location points at a metadata address is rejected even + though httpx leaves response.next_request unset inside the hook. + + Regression: the guard used to key on response.next_request, which is + None inside an httpx response event hook, so redirect-based SSRF (a + public URL 302-ing to 169.254.169.254) was never blocked. + """ + import gateway.platforms.yuanbao_media as ym + from tools import url_safety + + # Public pre-flight passes; the redirect target does not. + monkeypatch.setattr( + url_safety, "is_safe_url", lambda u: str(u).startswith("https://example.com") + ) + + captured = {} + + class _FakeResp: + headers = {"content-type": "image/png", "content-length": "3"} + is_redirect = False + next_request = None + + def raise_for_status(self): + pass + + async def aiter_bytes(self, _n): + yield b"png" + + class _FakeStream: + async def __aenter__(self): + return _FakeResp() + + async def __aexit__(self, *a): + return False + + class _FakeClient: + def __init__(self, *a, **kw): + captured["guard"] = kw["event_hooks"]["response"][0] + + async def __aenter__(self): + return self + + async def __aexit__(self, *a): + return False + + async def head(self, url): + return _FakeResp() + + def stream(self, method, url, **kw): + return _FakeStream() + + monkeypatch.setattr(ym.httpx, "AsyncClient", _FakeClient) + + await download_url("https://example.com/image.png") + guard = captured["guard"] + + class _RedirectResp: + # A genuine 302 as seen inside the hook: Location set, next_request None. + is_redirect = True + url = "https://example.com/image.png" + headers = {"location": "http://169.254.169.254/latest/meta-data/"} + next_request = None + + with pytest.raises(ValueError, match="Blocked redirect"): + await guard(_RedirectResp())