Studio: reject binary web_search fetches instead of decoding them into replacement chars (#7130)
* Studio: reject binary web_search fetches instead of decoding them into replacement chars
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Match web-fetch MIME subtypes exactly and detect control-char binary
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Sniff binary magic bytes and retry undeclared non-UTF-8 pages as text
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Harden web fetch binary sniffing
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Simplify web fetch binary guard
* Sniff unknown MIME types and handle Latin-1
* Sniff ambiguous Office MIME and prefixed magic
* Decode BOM-marked Unicode web content
* Studio: require ASCII evidence for declared Latin-1/cp1252 web fetches
Latin-1 and cp1252 decode every byte to a printable character, so a high-byte
binary body declared as iso-8859-1/windows-1252 decoded cleanly and slipped
past the control-character binary check. Apply the existing ASCII-structure gate
to those declared decodes as well. Scoped to the Latin family so legitimate
non-Latin single-byte pages (Cyrillic, Greek) are not rejected.
* Revert "Studio: require ASCII evidence for declared Latin-1/cp1252 web fetches"
This reverts commit c7fbec216c.
* Studio: tighten web-fetch binary guard comments
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: danielhanchen <michaelhan2050@gmail.com>
Co-authored-by: danielhanchen <danielhanchen@gmail.com>
This commit is contained in:
oobabooga
GitHub
parent
c5ae208abb
commit
770f92e250
2 changed files with 413 additions and 2 deletions
|
|
@ -5,6 +5,7 @@
|
|||
(DuckDuckGo), Python code execution, and terminal commands."""
|
||||
|
||||
import ast
|
||||
import codecs
|
||||
import fnmatch
|
||||
import http.client
|
||||
import os
|
||||
|
|
@ -3599,6 +3600,64 @@ _MAX_PAGE_CHARS = 16000 # cap fetched page text (after HTML-to-MD conversion)
|
|||
# Raw download cap > _MAX_PAGE_CHARS since SSR pages embed large <head> sections
|
||||
# stripped during conversion; 512 KB still reaches article content.
|
||||
_MAX_FETCH_BYTES = 512 * 1024
|
||||
# Control/undecodable chars, excluding text whitespace and ESC (for ANSI logs).
|
||||
# Binary when they exceed 12.5%, after allowing 16 minor encoding glitches.
|
||||
_BINARY_CHAR_RE = re.compile("[\\x00-\\x08\\x0b\\x0c\\x0e-\\x1a\\x1c-\\x1f\\x7f-\\x9f\\ufffd]")
|
||||
_MIN_BINARY_CHARS = 16
|
||||
_BINARY_CHAR_DIVISOR = 8
|
||||
# Common binary signatures that can otherwise look text-heavy when mislabeled.
|
||||
_BINARY_MAGIC = (
|
||||
b"%PDF-", # PDF
|
||||
b"PK\x03\x04", # zip / docx / xlsx / pptx / epub / jar
|
||||
b"\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1", # OLE / legacy Office
|
||||
b"\x89PNG\r\n\x1a\n", # PNG
|
||||
b"\xff\xd8\xff", # JPEG
|
||||
b"GIF87a",
|
||||
b"GIF89a",
|
||||
b"\x1f\x8b", # gzip
|
||||
b"BZh", # bzip2
|
||||
b"\xfd7zXZ\x00", # xz
|
||||
b"\x28\xb5\x2f\xfd", # zstd
|
||||
)
|
||||
|
||||
# Check UTF-32 first because its little-endian BOM starts with the UTF-16 BOM.
|
||||
_UNICODE_BOM_CODECS = (
|
||||
(codecs.BOM_UTF32_LE, "utf-32"),
|
||||
(codecs.BOM_UTF32_BE, "utf-32"),
|
||||
(codecs.BOM_UTF16_LE, "utf-16"),
|
||||
(codecs.BOM_UTF16_BE, "utf-16"),
|
||||
(codecs.BOM_UTF8, "utf-8-sig"),
|
||||
)
|
||||
|
||||
# A cp1252 retry needs 75% ASCII structure so it cannot rescue high-byte binary.
|
||||
_MIN_SINGLE_BYTE_ASCII_RATIO = 3 / 4
|
||||
_ASCII_TEXT_BYTES = frozenset((*range(0x20, 0x7F), 0x09, 0x0A, 0x0D, 0x1B))
|
||||
|
||||
|
||||
def _looks_binary(text: str) -> bool:
|
||||
"""Whether control or undecodable characters exceed the binary threshold."""
|
||||
return len(_BINARY_CHAR_RE.findall(text)) > max(
|
||||
_MIN_BINARY_CHARS, len(text) // _BINARY_CHAR_DIVISOR
|
||||
)
|
||||
|
||||
|
||||
def _has_binary_magic(data: bytes) -> bool:
|
||||
"""Whether a common binary signature follows optional BOM or whitespace."""
|
||||
head = data[:1024].lstrip()
|
||||
for bom, _codec in _UNICODE_BOM_CODECS:
|
||||
if head.startswith(bom):
|
||||
head = head.removeprefix(bom).lstrip()
|
||||
break
|
||||
return head.startswith(_BINARY_MAGIC)
|
||||
|
||||
|
||||
def _has_single_byte_text_evidence(data: bytes) -> bool:
|
||||
"""True when *data* has enough ASCII structure for a cp1252 text retry."""
|
||||
if not data:
|
||||
return True
|
||||
ascii_text_bytes = sum(byte in _ASCII_TEXT_BYTES for byte in data)
|
||||
return ascii_text_bytes / len(data) >= _MIN_SINGLE_BYTE_ASCII_RATIO
|
||||
|
||||
|
||||
_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",
|
||||
|
|
@ -3698,6 +3757,42 @@ def _validate_and_resolve_host(hostname: str, port: int) -> tuple[bool, str, str
|
|||
return True, "", first_ip
|
||||
|
||||
|
||||
# Binary application subtypes rejected by MIME; other application types are
|
||||
# sniffed so textual artifacts such as SQL stay usable.
|
||||
_BINARY_APPLICATION_SUBTYPES = frozenset(
|
||||
{
|
||||
"epub+zip",
|
||||
"gzip",
|
||||
"java-archive",
|
||||
"pdf",
|
||||
"vnd.apple.installer+xml",
|
||||
"wasm",
|
||||
"x-7z-compressed",
|
||||
"x-bzip2",
|
||||
"x-gzip",
|
||||
"x-rar-compressed",
|
||||
"x-tar",
|
||||
"x-xz",
|
||||
"zip",
|
||||
"zstd",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _is_text_candidate_content_type(content_type: str | None) -> bool:
|
||||
"""Whether a MIME type is textual or ambiguous enough for byte sniffing."""
|
||||
match = re.match(r"[\w.+-]+/[\w.+-]+", content_type or "")
|
||||
if not match:
|
||||
return True
|
||||
ct = match.group(0).lower()
|
||||
if ct.startswith("text/"):
|
||||
return True
|
||||
if ct.startswith("application/"):
|
||||
subtype = ct[len("application/") :]
|
||||
return subtype not in _BINARY_APPLICATION_SUBTYPES
|
||||
return False
|
||||
|
||||
|
||||
# First path segments on github.com that are site pages, not repo owners.
|
||||
_GITHUB_NON_OWNER_SEGMENTS = frozenset(
|
||||
{
|
||||
|
|
@ -3975,7 +4070,6 @@ def _fetch_url_raw(
|
|||
else:
|
||||
return "Failed to fetch URL: too many redirects.", "", ""
|
||||
|
||||
charset = resp.headers.get_content_charset() or "utf-8"
|
||||
# get_content_type() defaults to "text/plain" when the header is
|
||||
# absent (RFC 2045); report "" instead so callers can tell a missing
|
||||
# header apart from a server that really declared text/plain.
|
||||
|
|
@ -3983,7 +4077,54 @@ def _fetch_url_raw(
|
|||
content_type = ""
|
||||
else:
|
||||
content_type = (resp.headers.get_content_type() or "").lower()
|
||||
return None, raw_bytes.decode(charset, errors = "replace"), content_type
|
||||
|
||||
# Reject known-binary MIME types before decoding. Binary is returned as the
|
||||
# error string so the caller surfaces the placeholder, not replacement chars.
|
||||
if not _is_text_candidate_content_type(content_type):
|
||||
# Only echo a clean MIME token back to the model.
|
||||
m = re.match(r"[\w.+-]+/[\w.+-]+", content_type or "")
|
||||
safe_type = m.group(0) if m else "unknown type"
|
||||
return (
|
||||
f"(non-text content: {safe_type}, {len(raw_bytes)} bytes; not readable as text)",
|
||||
"",
|
||||
content_type,
|
||||
)
|
||||
|
||||
# Catch text-labeled binary via its magic signature.
|
||||
if _has_binary_magic(raw_bytes):
|
||||
return (
|
||||
f"(binary content, {len(raw_bytes)} bytes; not readable as text)",
|
||||
"",
|
||||
content_type,
|
||||
)
|
||||
|
||||
declared = resp.headers.get_content_charset()
|
||||
declared_codec = codecs.lookup(declared).name if declared else None
|
||||
bom_codec = next(
|
||||
(codec for bom, codec in _UNICODE_BOM_CODECS if raw_bytes.startswith(bom)),
|
||||
None,
|
||||
)
|
||||
raw_html = raw_bytes.decode(declared or bom_codec or "utf-8", errors = "replace")
|
||||
|
||||
# Catch mislabeled or unlabeled binary, including valid UTF-8 controls.
|
||||
if _looks_binary(raw_html):
|
||||
# Rescue undeclared cp1252 only when the bytes have text structure.
|
||||
alt = (
|
||||
raw_bytes.decode("cp1252", "replace")
|
||||
if declared_codec in (None, "iso8859-1")
|
||||
and _has_single_byte_text_evidence(raw_bytes)
|
||||
else None
|
||||
)
|
||||
if alt is not None and not _looks_binary(alt):
|
||||
raw_html = alt
|
||||
else:
|
||||
return (
|
||||
f"(binary content, {len(raw_bytes)} bytes; not readable as text)",
|
||||
"",
|
||||
content_type,
|
||||
)
|
||||
|
||||
return None, raw_html, content_type
|
||||
except _HTTPError as e:
|
||||
return f"Failed to fetch URL: HTTP {e.code} {getattr(e, 'reason', '')}", "", ""
|
||||
except Exception as e:
|
||||
|
|
|
|||
270
studio/backend/tests/test_web_fetch_binary_guard.py
Normal file
270
studio/backend/tests/test_web_fetch_binary_guard.py
Normal file
|
|
@ -0,0 +1,270 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Regression tests for binary bodies poisoning web_search model context (#7084)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import codecs
|
||||
import sys
|
||||
from email.message import Message
|
||||
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)
|
||||
|
||||
from core.inference import tools
|
||||
|
||||
|
||||
class _FakeResp:
|
||||
def __init__(self, body: bytes, content_type: str | None):
|
||||
self._body = body
|
||||
self._pos = 0
|
||||
self.headers = Message()
|
||||
if content_type is not None:
|
||||
self.headers["Content-Type"] = content_type
|
||||
|
||||
def read(self, n: int | None = None) -> bytes:
|
||||
# Advance a cursor like a real stream so the chunked reader reaches EOF.
|
||||
chunk = self._body[self._pos :] if n is None else self._body[self._pos : self._pos + n]
|
||||
self._pos += len(chunk)
|
||||
return chunk
|
||||
|
||||
|
||||
class _FakeOpener:
|
||||
def __init__(self, resp):
|
||||
self._resp = resp
|
||||
|
||||
def open(
|
||||
self,
|
||||
req,
|
||||
timeout = None,
|
||||
):
|
||||
return self._resp
|
||||
|
||||
|
||||
def _fetch_with(monkeypatch, body: bytes, content_type: str | None) -> str:
|
||||
# Pass SSRF validation and skip real DNS/network.
|
||||
monkeypatch.setattr(
|
||||
tools, "_validate_and_resolve_host", lambda host, port: (True, "", "93.184.216.34")
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
tools.urllib.request,
|
||||
"build_opener",
|
||||
lambda *a, **k: _FakeOpener(_FakeResp(body, content_type)),
|
||||
)
|
||||
return tools._fetch_page_text("https://example.com/thing", timeout = 5)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"content_type,expected",
|
||||
[
|
||||
("text/html", True),
|
||||
("text/plain; charset=utf-8", True),
|
||||
("application/json", True),
|
||||
("application/json; charset=utf-8", True),
|
||||
("application/xml", True),
|
||||
("application/xhtml+xml", True),
|
||||
("application/ld+json", True),
|
||||
("application/yaml", True),
|
||||
("application/x-yaml", True),
|
||||
("application/x-ndjson", True),
|
||||
("application/ndjson", True),
|
||||
("application/sql", True),
|
||||
("application/x-www-form-urlencoded", True),
|
||||
("application/pdf", False),
|
||||
("image/png", False),
|
||||
("image/svg+xml", False),
|
||||
("application/octet-stream", True),
|
||||
("application/zip", False),
|
||||
("application/vnd.ms-excel", True),
|
||||
("application/vnd.openxmlformats-officedocument.wordprocessingml.document", True),
|
||||
("", True),
|
||||
(None, True),
|
||||
],
|
||||
)
|
||||
def test_is_text_candidate_content_type(content_type, expected):
|
||||
assert tools._is_text_candidate_content_type(content_type) is expected
|
||||
|
||||
|
||||
def test_pdf_rejected_by_content_type(monkeypatch):
|
||||
out = _fetch_with(monkeypatch, b"%PDF-1.7\n\xff\xd8\xff\x00\x89PNG" * 200, "application/pdf")
|
||||
assert "<EFBFBD>" not in out
|
||||
assert "non-text content" in out and "application/pdf" in out
|
||||
|
||||
|
||||
def test_text_octet_stream_kept_after_sniffing(monkeypatch):
|
||||
body = b"level=info\nmessage=plain text artifact\n" * 100
|
||||
out = _fetch_with(monkeypatch, body, "application/octet-stream")
|
||||
assert "plain text artifact" in out
|
||||
assert "non-text content" not in out and "binary content" not in out
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"content_type",
|
||||
["application/octet-stream", "application/x-custom-binary", "text/plain", None],
|
||||
)
|
||||
def test_binary_candidates_rejected_after_sniffing(monkeypatch, content_type):
|
||||
out = _fetch_with(monkeypatch, bytes(range(256)) * 20, content_type)
|
||||
assert "<EFBFBD>" not in out
|
||||
assert "binary content" in out
|
||||
|
||||
|
||||
@pytest.mark.parametrize("content_type", ["application/sql", "application/x-www-form-urlencoded"])
|
||||
def test_unknown_application_text_kept_after_sniffing(monkeypatch, content_type):
|
||||
out = _fetch_with(monkeypatch, b"select readable_text from artifacts;\n" * 100, content_type)
|
||||
assert "readable_text" in out
|
||||
assert "non-text content" not in out and "binary content" not in out
|
||||
|
||||
|
||||
def test_excel_labeled_csv_kept_after_sniffing(monkeypatch):
|
||||
body = b"name,value\nreadable,42\n" * 100
|
||||
out = _fetch_with(monkeypatch, body, "application/vnd.ms-excel")
|
||||
assert "readable" in out
|
||||
assert "binary content" not in out
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"bom,encoding",
|
||||
[
|
||||
(codecs.BOM_UTF16_LE, "utf-16-le"),
|
||||
(codecs.BOM_UTF16_BE, "utf-16-be"),
|
||||
(codecs.BOM_UTF32_LE, "utf-32-le"),
|
||||
(codecs.BOM_UTF32_BE, "utf-32-be"),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("content_type", ["text/plain", "application/vnd.ms-excel"])
|
||||
def test_bom_unicode_text_without_charset_kept(monkeypatch, bom, encoding, content_type):
|
||||
body = bom + ("name,value\nreadable,42\n" * 100).encode(encoding)
|
||||
out = _fetch_with(monkeypatch, body, content_type)
|
||||
assert "readable" in out
|
||||
assert "binary content" not in out
|
||||
|
||||
|
||||
def test_valid_utf8_binary_caught_by_control_chars(monkeypatch):
|
||||
# These controls are valid UTF-8 and therefore produce no replacement chars.
|
||||
body = bytes([0, 1, 2, 3, 4, 5, 6, 7]) * 400
|
||||
out = _fetch_with(monkeypatch, body, "text/plain")
|
||||
assert "binary content" in out
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"magic",
|
||||
[
|
||||
b"%PDF-",
|
||||
b"PK\x03\x04",
|
||||
b"\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1",
|
||||
b"\x1f\x8b",
|
||||
b"BZh",
|
||||
b"\xfd7zXZ\x00",
|
||||
b"\x28\xb5\x2f\xfd",
|
||||
],
|
||||
)
|
||||
def test_text_labeled_binary_caught_by_magic(monkeypatch, magic):
|
||||
out = _fetch_with(monkeypatch, magic + b" printable text-heavy body" * 100, "text/plain")
|
||||
assert "binary content" in out
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"prefix",
|
||||
[
|
||||
codecs.BOM_UTF8,
|
||||
codecs.BOM_UTF16_LE,
|
||||
codecs.BOM_UTF16_BE,
|
||||
codecs.BOM_UTF32_LE,
|
||||
codecs.BOM_UTF32_BE,
|
||||
b" \r\n",
|
||||
b"\t\xef\xbb\xbf ",
|
||||
],
|
||||
)
|
||||
def test_pdf_magic_after_harmless_prefix(monkeypatch, prefix):
|
||||
body = prefix + b"%PDF-1.7\n" + b"1 0 obj<</Type/Catalog>>endobj\n" * 100
|
||||
out = _fetch_with(monkeypatch, body, "text/plain")
|
||||
assert "binary content" in out
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"content_type,magic",
|
||||
[
|
||||
("application/vnd.ms-excel", b"\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1"),
|
||||
(
|
||||
"application/vnd.openxmlformats-officedocument.wordprocessingml.document",
|
||||
b"PK\x03\x04",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_office_labeled_binary_caught_by_magic(monkeypatch, content_type, magic):
|
||||
out = _fetch_with(monkeypatch, magic + b" printable text-heavy body" * 100, content_type)
|
||||
assert "binary content" in out
|
||||
|
||||
|
||||
def test_latin1_text_without_charset_kept(monkeypatch):
|
||||
# The cp1252 retry should rescue accent-heavy text with ASCII structure.
|
||||
body = (
|
||||
"Muller lauft uber die Strasse: schoene, groesse. MARKERWORD ".replace("ue", "ü")
|
||||
+ "äöüß éèà "
|
||||
) * 30
|
||||
out = _fetch_with(monkeypatch, body.encode("cp1252"), "text/plain")
|
||||
assert "binary content" not in out
|
||||
assert "MARKERWORD" in out
|
||||
|
||||
|
||||
@pytest.mark.parametrize("charset", ["iso-8859-1", "latin-1", "latin1"])
|
||||
def test_declared_latin1_cp1252_punctuation_kept(monkeypatch, charset):
|
||||
body = ("“quoted” " * 100).encode("cp1252")
|
||||
out = _fetch_with(monkeypatch, body, f"text/plain; charset={charset}")
|
||||
assert "quoted" in out
|
||||
assert "binary content" not in out
|
||||
|
||||
|
||||
def test_high_byte_binary_not_rescued_as_cp1252(monkeypatch):
|
||||
# cp1252 maps these bytes to printable characters, but they lack ASCII structure.
|
||||
body = bytes(range(0xA0, 0x100)) * 40
|
||||
out = _fetch_with(monkeypatch, body, "text/plain")
|
||||
assert "binary content" in out
|
||||
|
||||
|
||||
def test_ansi_colored_text_log_kept(monkeypatch):
|
||||
# ESC is excluded from the binary set so ANSI logs remain readable.
|
||||
line = "".join(f"\x1b[32m+{i}\x1b[0m\n" for i in range(300)).encode()
|
||||
out = _fetch_with(monkeypatch, line, "text/plain")
|
||||
assert "binary content" not in out
|
||||
|
||||
|
||||
def test_html_page_unaffected(monkeypatch):
|
||||
html = b"<html><body><h1>Hello</h1><p>Real text content here.</p></body></html>"
|
||||
out = _fetch_with(monkeypatch, html, "text/html; charset=utf-8")
|
||||
assert "Hello" in out
|
||||
assert "non-text content" not in out and "binary content" not in out
|
||||
|
||||
|
||||
def test_content_type_sanitized_in_message(monkeypatch):
|
||||
# Do not echo obs-folded header content into the model response.
|
||||
out = _fetch_with(monkeypatch, b"\x00\x01\x02" * 500, "application/pdf\r\n data: injected")
|
||||
assert "\n" not in out and "\r" not in out
|
||||
assert "injected" not in out
|
||||
assert "application/pdf" in out
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"n_bad,n_total,expect_binary",
|
||||
[
|
||||
(120, 1000, False),
|
||||
(130, 1000, True),
|
||||
],
|
||||
)
|
||||
def test_binary_char_ratio_boundary(monkeypatch, n_bad, n_total, expect_binary):
|
||||
body = b"\x00" * n_bad + b"a" * (n_total - n_bad)
|
||||
out = _fetch_with(monkeypatch, body, "text/plain")
|
||||
assert ("binary content" in out) is expect_binary
|
||||
|
||||
|
||||
def test_text_with_a_few_stray_replacement_chars_kept(monkeypatch):
|
||||
# Minor encoding glitches below the floor should not drop a real page.
|
||||
body = ("Real article text. " * 200).encode() + b"\xff\xfe\xff"
|
||||
out = _fetch_with(monkeypatch, body, "text/html")
|
||||
assert "Real article text." in out
|
||||
assert "binary content" not in out
|
||||
Loading…
Add table
Add a link
Reference in a new issue