fix(studio): support hostname-based enterprise proxies (#7416)

* fix(studio): support hostname-based enterprise proxies

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix(studio): strip userinfo from proxy fetch targets

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>

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.

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,

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
 

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(

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"

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.