diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py index 0ef6dd46cf..b31db3faf6 100644 --- a/studio/backend/core/inference/tools.py +++ b/studio/backend/core/inference/tools.py @@ -49,6 +49,7 @@ from loggers import get_logger logger = get_logger(__name__) _EXEC_TIMEOUT = 300 # 5 minutes +_DISABLE_DNS_PINNING_ENV = "UNSLOTH_STUDIO_DISABLE_DNS_PINNING" # Splits the UI source-map from the result; loops strip it (like __IMAGES__). RAG_SOURCES_SENTINEL = "\n__RAG_SOURCES__:" @@ -4194,13 +4195,18 @@ def _fetch_url_raw( budget_error = _fetch_budget_exceeded(deadline, cancel_event) if budget_error is not None: return budget_error, "", "" - # Pin to the validated IP (prevents DNS rebinding): rewrite URL to - # the IP, set the Host header. cp = urlparse(current_url) - # Bracket IPv6 addresses so the netloc is valid in a URL. - ip_str = f"[{pinned_ip}]" if ":" in pinned_ip else pinned_ip - ip_netloc = f"{ip_str}:{cp.port}" if cp.port else ip_str - pinned_url = urlunparse(cp._replace(netloc = ip_netloc)) + validated_netloc = f"[{current_host}]" if ":" in current_host else current_host + if cp.port: + validated_netloc = f"{validated_netloc}:{cp.port}" + if os.environ.get(_DISABLE_DNS_PINNING_ENV) == "1": + # Enterprise proxies need the hostname in CONNECT for policy and TLS interception. + request_url = urlunparse(cp._replace(netloc = validated_netloc)) + else: + # Pin to the validated IP to prevent DNS rebinding. + ip_str = f"[{pinned_ip}]" if ":" in pinned_ip else pinned_ip + ip_netloc = f"{ip_str}:{cp.port}" if cp.port else ip_str + request_url = urlunparse(cp._replace(netloc = ip_netloc)) opener = urllib.request.build_opener( _NoRedirect, @@ -4209,11 +4215,11 @@ def _fetch_url_raw( headers = { "User-Agent": ua, - "Host": current_host, + "Host": validated_netloc, } if extra_headers: headers.update(extra_headers) - req = urllib.request.Request(pinned_url, headers = headers) + req = urllib.request.Request(request_url, headers = headers) try: # Cap the socket timeout at the time left on the overall deadline # so a single slow hop cannot outlast the whole fetch budget. diff --git a/studio/backend/run.py b/studio/backend/run.py index 90fd28670c..f1fc8c6062 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -1885,6 +1885,12 @@ def _build_arg_parser(): default = None, help = "Force server-side tools off for every request.", ) + parser.add_argument( + "--disable-dns-pinning", + action = "store_true", + help = "Allow hostname-based web fetches for enterprise proxies. WARNING: weakens " + "DNS-rebinding protection; hostname and redirect validation remain enabled.", + ) parser.add_argument( "--parallel", "--n-parallel", @@ -1924,6 +1930,10 @@ if __name__ == "__main__": parser.error( "--secure requires the Cloudflare tunnel; do not combine it with --no-cloudflare" ) + if args.disable_dns_pinning: + os.environ["UNSLOTH_STUDIO_DISABLE_DNS_PINNING"] = "1" + else: + os.environ.setdefault("UNSLOTH_STUDIO_DISABLE_DNS_PINNING", "0") kwargs = dict( host = args.host, diff --git a/studio/backend/tests/test_secure_tunnel_gate.py b/studio/backend/tests/test_secure_tunnel_gate.py index a8c0c2305f..b491134045 100644 --- a/studio/backend/tests/test_secure_tunnel_gate.py +++ b/studio/backend/tests/test_secure_tunnel_gate.py @@ -85,6 +85,14 @@ def test_arg_parser_secure_polarity_and_not_secure_alias(): assert parser.parse_args(["--not-secure", "--secure"]).secure is True +def test_arg_parser_dns_pinning_opt_out_defaults_off(): + import run + + parser = run._build_arg_parser() + assert parser.parse_args([]).disable_dns_pinning is False + assert parser.parse_args(["--disable-dns-pinning"]).disable_dns_pinning is True + + def test_run_server_accepts_enable_tools_kwarg(): import inspect diff --git a/studio/backend/tests/test_web_fetch_extraction.py b/studio/backend/tests/test_web_fetch_extraction.py index b794ee3e81..d4c3d123c3 100644 --- a/studio/backend/tests/test_web_fetch_extraction.py +++ b/studio/backend/tests/test_web_fetch_extraction.py @@ -15,6 +15,8 @@ from __future__ import annotations import sys from pathlib import Path +import pytest + _BACKEND_DIR = str(Path(__file__).resolve().parent.parent) if _BACKEND_DIR not in sys.path: sys.path.insert(0, _BACKEND_DIR) @@ -715,6 +717,61 @@ def test_fetch_url_raw_missing_content_type_reported_empty(monkeypatch): assert content_type == "" +@pytest.mark.parametrize( + "disable_dns_pinning,expected_url", + [ + (False, "https://203.0.113.7:8443/page?q=1"), + (True, "https://example.com:8443/page?q=1"), + ], +) +def test_fetch_url_raw_dns_pinning_proxy_opt_out(monkeypatch, disable_dns_pinning, expected_url): + import email + import urllib.request + + import core.inference.tools as tools_mod + + class _FakeResp: + headers = email.message_from_string("Content-Type: text/plain\n") + + def __init__(self): + self._body = b"ok" + + def read(self, n = -1): + body, self._body = self._body, b"" + return body + + requested = [] + + class _FakeOpener: + def open( + self, + req, + timeout = None, + ): + requested.append(req) + return _FakeResp() + + resolved = [] + + def resolve(host, port): + resolved.append((host, port)) + return True, "", "203.0.113.7" + + monkeypatch.setenv("UNSLOTH_STUDIO_DISABLE_DNS_PINNING", "1" if disable_dns_pinning else "0") + monkeypatch.setattr(tools_mod, "_validate_and_resolve_host", resolve) + monkeypatch.setattr(urllib.request, "build_opener", lambda *handlers: _FakeOpener()) + + err, body, _content_type = tools_mod._fetch_url_raw( + "https://user:secret@example.com:8443/page?q=1" + ) + + assert err is None + assert body == "ok" + assert resolved == [("example.com", 8443)] + assert [req.full_url for req in requested] == [expected_url] + assert requested[0].get_header("Host") == "example.com:8443" + + def test_fetch_page_text_missing_content_type_html_sniffed(monkeypatch): # A header-less server returning an HTML body must still be converted. def fake_fetch( diff --git a/tests/studio/test_cli_studio_defaults.py b/tests/studio/test_cli_studio_defaults.py index b03fa57603..a39956fab0 100644 --- a/tests/studio/test_cli_studio_defaults.py +++ b/tests/studio/test_cli_studio_defaults.py @@ -68,3 +68,10 @@ def test_studio_run_host_is_loopback(): f"`unsloth studio run` --host default must be '127.0.0.1' (loopback) " f"but got '{host_default}'." ) + + +def test_dns_pinning_opt_out_is_registered_safe_by_default(): + source = _STUDIO_CMD_PY.read_text() + for func_name in ("studio_default", "run"): + default = _find_typer_option_default(source, func_name, "--disable-dns-pinning") + assert default is False, f"{func_name} must keep DNS pinning enabled by default" diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py index 84c22aa4ac..c41963aab9 100644 --- a/unsloth_cli/commands/studio.py +++ b/unsloth_cli/commands/studio.py @@ -1285,6 +1285,12 @@ def studio_default( help = "Force server-side tools (web search, code execution) on or off for " "every request. Default: on for every bind, with the per-chat UI toggle honored.", ), + disable_dns_pinning: bool = typer.Option( + False, + "--disable-dns-pinning", + help = "Allow hostname-based web fetches for enterprise proxies. WARNING: weakens " + "DNS-rebinding protection; hostname and redirect validation remain enabled.", + ), password: str = typer.Option( "", "--password", @@ -1354,6 +1360,15 @@ def studio_default( err = True, ) raise typer.Exit(2) + if disable_dns_pinning: + typer.echo( + "Error: --disable-dns-pinning on `unsloth studio` applies to the " + f"plain-server path only. For `unsloth studio {ctx.invoked_subcommand}`, " + f"put it after the subcommand: `unsloth studio {ctx.invoked_subcommand} " + "--disable-dns-pinning ...`", + err = True, + ) + raise typer.Exit(2) # Same for --api-only: dropping it here would silently serve the UI. if api_only: typer.echo( @@ -1398,6 +1413,10 @@ def studio_default( # default (plain-server path; the `run` subcommand has its own --verbose). if verbose: _enable_verbose_access_logs() + if disable_dns_pinning: + os.environ["UNSLOTH_STUDIO_DISABLE_DNS_PINNING"] = "1" + else: + os.environ.setdefault("UNSLOTH_STUDIO_DISABLE_DNS_PINNING", "0") # Use the studio venv if present and not already in it. Resolve the child # launcher BEFORE the gate: a headless gate strips the seeded @@ -1739,6 +1758,13 @@ def run( "every request. Default: on for every bind." ), ), + disable_dns_pinning: bool = typer.Option( + False, + "--disable-dns-pinning", + rich_help_panel = _RUN_PANEL_TOOLS, + help = "Allow hostname-based web fetches for enterprise proxies. WARNING: weakens " + "DNS-rebinding protection; hostname and redirect validation remain enabled.", + ), tool_call_healing: Optional[bool] = typer.Option( None, "--enable-tool-call-healing/--disable-tool-call-healing", @@ -1944,6 +1970,10 @@ def run( _enable_verbose_access_logs() if not any(a in ("--verbose", "-v", "--log-verbose") for a in extra_llama_args): extra_llama_args.append("--log-verbose") + if disable_dns_pinning: + os.environ["UNSLOTH_STUDIO_DISABLE_DNS_PINNING"] = "1" + else: + os.environ.setdefault("UNSLOTH_STUDIO_DISABLE_DNS_PINNING", "0") # Promote legacy exact `-m`/`-hfr`/`-f` back into typer params; # clusters stay in extras.