diff --git a/studio/backend/core/inference/_html_to_md.py b/studio/backend/core/inference/_html_to_md.py
index d96b8168e2..f999120ffb 100644
--- a/studio/backend/core/inference/_html_to_md.py
+++ b/studio/backend/core/inference/_html_to_md.py
@@ -17,18 +17,26 @@ from html.parser import HTMLParser
__all__ = ["html_to_markdown"]
-_SKIP_TAGS = frozenset({"script", "style", "head", "noscript", "svg", "math"})
+_SKIP_TAGS = frozenset(
+ {
+ "script",
+ "style",
+ "head",
+ "noscript",
+ "svg",
+ "math",
+ "nav",
+ "footer",
+ }
+)
_BLOCK_TAGS = frozenset(
{
"p",
"div",
"section",
"article",
- "header",
- "footer",
"main",
"aside",
- "nav",
"figure",
"figcaption",
"details",
diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py
index 2ac8f76322..d425daa49d 100644
--- a/studio/backend/core/inference/tools.py
+++ b/studio/backend/core/inference/tools.py
@@ -8,14 +8,18 @@ Supports web search (DuckDuckGo), Python code execution, and terminal commands.
"""
import ast
+import http.client
import os
os.environ["UNSLOTH_IS_PRESENT"] = "1"
+import random
+import ssl
import subprocess
import sys
import tempfile
import threading
+import urllib.request
from loggers import get_logger
@@ -154,8 +158,74 @@ def execute_tool(
return f"Unknown tool: {name}"
-_MAX_PAGE_CHARS = 16000 # limit fetched page text
-_MAX_FETCH_BYTES = _MAX_PAGE_CHARS * 4 + 1 # cap raw download size
+_MAX_PAGE_CHARS = 16000 # limit fetched page text (after HTML-to-MD conversion)
+# Raw download cap. Must be larger than _MAX_PAGE_CHARS because SSR pages
+# embed large sections (CSS, JS, SVGs) that are stripped during
+# HTML-to-Markdown conversion. 512 KB is enough to reach article content
+# on GitBook / Next.js / Docusaurus pages whose alone can be 200 KB.
+_MAX_FETCH_BYTES = 512 * 1024
+
+_USER_AGENTS = (
+ "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36",
+ "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36",
+ "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36",
+ "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:133.0) Gecko/20100101 Firefox/133.0",
+ "Mozilla/5.0 (Macintosh; Intel Mac OS X 10.15; rv:133.0) Gecko/20100101 Firefox/133.0",
+ "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/18.2 Safari/605.1.15",
+)
+
+_tls_ctx = ssl.create_default_context()
+
+
+class _NoRedirect(urllib.request.HTTPRedirectHandler):
+ def redirect_request(self, req, fp, code, msg, headers, newurl):
+ return None
+
+
+class _PinnedHTTPSConnection(http.client.HTTPSConnection):
+ """HTTPS connection that connects to a pinned IP but uses a different
+ hostname for SNI and certificate verification.
+
+ The SSRF IP-pinning rewrites URLs to raw IPs. A normal HTTPSConnection
+ would then send no SNI and verify the cert against the IP, both of which
+ fail. This subclass splits the two concerns: TCP connects to the pinned
+ IP (``host`` parameter) while TLS uses ``sni_hostname`` for the
+ ClientHello and cert check.
+ """
+
+ def __init__(self, host: str, *, sni_hostname: str, **kwargs):
+ super().__init__(host, **kwargs)
+ self._sni_hostname = sni_hostname
+
+ def connect(self):
+ # TCP connect to the pinned IP stored in self.host (+ tunnel if
+ # a proxy is configured via set_tunnel, though we do not use one).
+ http.client.HTTPConnection.connect(self)
+ # TLS handshake with the real hostname for SNI + cert verification.
+ self.sock = self._context.wrap_socket(
+ self.sock,
+ server_hostname = self._sni_hostname,
+ )
+
+
+class _SNIHTTPSHandler(urllib.request.HTTPSHandler):
+ """HTTPS handler that sends the correct SNI hostname during TLS handshake.
+
+ The SSRF IP-pinning rewrites URLs to raw IPs, which breaks SNI and cert
+ verification. This handler returns a ``_PinnedHTTPSConnection`` that
+ connects to the pinned IP but verifies TLS against the original hostname.
+ """
+
+ def __init__(self, hostname: str):
+ super().__init__(context = _tls_ctx)
+ self._sni_hostname = hostname
+
+ def https_open(self, req):
+ return self.do_open(self._sni_connection, req)
+
+ def _sni_connection(self, host, **kwargs):
+ kwargs["context"] = _tls_ctx
+ return _PinnedHTTPSConnection(host, sni_hostname = self._sni_hostname, **kwargs)
def _validate_and_resolve_host(hostname: str, port: int) -> tuple[bool, str, str]:
@@ -215,33 +285,32 @@ def _fetch_page_text(
return reason
try:
- import urllib.request
from urllib.error import HTTPError as _HTTPError
from urllib.parse import urljoin, urlunparse
- # Disable auto-redirect so we can validate each hop for SSRF.
- # urllib raises HTTPError for 3xx when the handler returns None,
- # so we catch that and extract the Location header manually.
- class _NoRedirect(urllib.request.HTTPRedirectHandler):
- def redirect_request(self, req, fp, code, msg, headers, newurl):
- return None
-
- opener = urllib.request.build_opener(_NoRedirect)
- max_bytes = max_chars * 4 + 1
+ max_bytes = _MAX_FETCH_BYTES
current_url = url
current_host = parsed.hostname
+ ua = random.choice(_USER_AGENTS)
for _hop in range(5):
# Pin to the validated IP to prevent DNS rebinding.
# Rewrite the URL to use the IP and set the Host header.
cp = urlparse(current_url)
- ip_netloc = f"{pinned_ip}:{cp.port}" if cp.port else pinned_ip
+ # 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))
+ opener = urllib.request.build_opener(
+ _NoRedirect,
+ _SNIHTTPSHandler(current_host),
+ )
+
req = urllib.request.Request(
pinned_url,
headers = {
- "User-Agent": "UnslothStudio/1.0",
+ "User-Agent": ua,
"Host": current_host,
},
)