Simplify web fetch binary guard

This commit is contained in:
oobabooga 2026-07-14 22:15:48 -03:00
commit 10592853fc
2 changed files with 38 additions and 130 deletions

View file

@ -1381,17 +1381,12 @@ _MAX_PAGE_CHARS = 16000 # cap fetched page text (after HTML-to-MD conversion)
# sections stripped during conversion; 512 KB reaches article content even
# where <head> alone is ~200 KB.
_MAX_FETCH_BYTES = 512 * 1024
# Chars that don't occur in real text: undecodable bytes (U+FFFD) plus control
# chars (C0 minus tab/newline/CR and ESC, DEL, C1). ESC is excluded so a text
# page of ANSI-colored terminal output isn't mistaken for binary. A decoded page
# is treated as binary when more than 1/_BINARY_CHAR_DIVISOR (12.5%) of its chars
# are these, tolerating up to _MIN_BINARY_CHARS so minor glitches don't drop it.
# Undecodable bytes and controls, excluding text whitespace and ESC for ANSI logs.
# More than 12.5% is binary 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
# Leading bytes of common binary formats that servers sometimes mislabel as
# text/*; matched by signature because their replacement-char density alone can
# be low (e.g. a PDF whose first chunk is mostly ASCII object/xref syntax).
# 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
@ -1405,17 +1400,13 @@ _BINARY_MAGIC = (
b"\x28\xb5\x2f\xfd", # zstd
)
# Undeclared legacy Western text should have substantial ASCII structure even
# when its non-ASCII bytes are cp1252. Requiring at least 75% ASCII text bytes
# prevents arbitrary high-byte binary from becoming printable-looking gibberish
# merely because cp1252 defines a character for nearly every byte.
# 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:
"""True when more than 1/_BINARY_CHAR_DIVISOR of ``text`` is binary chars
(control/undecodable, per _BINARY_CHAR_RE), past the _MIN_BINARY_CHARS floor."""
"""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
)
@ -1536,27 +1527,16 @@ _TEXT_APPLICATION_SUBTYPES = frozenset(
)
def _is_texty_content_type(content_type: str | None) -> bool:
"""True for MIME types safe to decode or inspect as possible text.
This includes HTML, plain text, JSON/XML/YAML feeds, and generic
``application/octet-stream`` downloads whose bytes must be sniffed. Known
binary types (PDF, images, archives, Office/ZIP) return False so the fetcher
never decodes them into replacement chars that poison the model context
(unslothai/unsloth#7084).
A missing Content-Type coerces to ``text/plain`` upstream, so unlabeled
bodies pass here and are caught instead by the binary-char fallback.
"""
def _is_text_candidate_content_type(content_type: str | None) -> bool:
"""Whether a MIME type is textual or ambiguous enough for byte sniffing."""
ct = (content_type or "").partition(";")[0].strip().lower()
if not ct:
return True # unlabeled: let the binary-char fallback decide
return True
if ct.startswith("text/"):
return True
if ct.startswith("application/"):
# Match the exact subtype, not a loose substring, so a .docx labeled
# application/vnd.openxmlformats-... isn't read as xml. RFC 6839
# +json/+xml suffixes (ld+json, xhtml+xml, ...) pass too.
# Exact matches avoid treating "openxmlformats" as XML. Structured
# +json/+xml suffixes remain valid text candidates.
subtype = ct[len("application/") :].removeprefix("x-")
return (
subtype == "octet-stream"
@ -1646,33 +1626,24 @@ def _fetch_page_text(
else:
return "Failed to fetch URL: too many redirects."
# Reject binary bodies (PDF, image, archive): decoding them as text
# floods the model context with U+FFFD replacement chars (#7084).
# Reject MIME types known to be binary before decoding.
content_type = resp.headers.get_content_type()
if not _is_texty_content_type(content_type):
# Trim to a clean MIME token: get_content_type() can echo control
# chars from an obs-folded header, and this string is returned to
# the model.
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)"
# A PDF/image/archive mislabeled as text/*: catch by signature, since a
# text-heavy first chunk can slip past the replacement-char density check.
# Catch text-labeled binary whose header and first chunk look textual.
if raw_bytes.startswith(_BINARY_MAGIC):
return f"(binary content, {len(raw_bytes)} bytes; not readable as text)"
declared = resp.headers.get_content_charset()
raw_html = raw_bytes.decode(declared or "utf-8", errors = "replace")
# Fallback for binary mislabeled as text/* or sent with no Content-Type,
# including valid-UTF-8 binary (NUL/control-heavy payloads) that decodes
# without replacement chars: a real text page has few binary chars.
# Catch mislabeled or unlabeled binary, including valid UTF-8 controls.
if _looks_binary(raw_html):
# An undeclared non-UTF-8 page (latin-1/cp1252) also decodes to many
# U+FFFD. Retry as cp1252 only when the bytes have substantial ASCII
# text structure; otherwise high-byte binary would become printable
# looking gibberish because cp1252 maps nearly every byte.
# Rescue undeclared cp1252 only when the bytes have text structure.
alt = (
raw_bytes.decode("cp1252", "replace")
if declared is None and _has_single_byte_text_evidence(raw_bytes)

View file

@ -1,11 +1,7 @@
# 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 for unslothai/unsloth#7084: the web_search fetcher must not decode
a binary body into control characters or a flood of U+FFFD replacement chars
that poison the model context. It rejects known binary Content-Types up front
and sniffs generic, mislabeled, or unlabeled bodies. Real text is unaffected.
"""
"""Regression tests for binary bodies poisoning web_search model context (#7084)."""
from __future__ import annotations
@ -58,9 +54,6 @@ def _fetch_with(monkeypatch, body: bytes, content_type: str | None) -> str:
return tools._fetch_page_text("https://example.com/thing", timeout = 5)
# ── content-type classifier ──
@pytest.mark.parametrize(
"content_type,expected",
[
@ -71,27 +64,22 @@ def _fetch_with(monkeypatch, body: bytes, content_type: str | None) -> str:
("application/xml", True),
("application/xhtml+xml", True),
("application/ld+json", True),
("application/yaml", True), # raw configs / API specs
("application/yaml", True),
("application/x-yaml", True),
("application/x-ndjson", True), # newline-delimited JSON is text
("application/x-ndjson", True),
("application/ndjson", True),
("application/pdf", False),
("image/png", False),
("image/svg+xml", False), # SVG source isn't extracted downstream; reject
# Generic downloads need byte sniffing because they can contain text.
("image/svg+xml", False),
("application/octet-stream", True),
("application/zip", False),
# A .docx ZIP must not pass just because "xml" is inside "openxmlformats".
("application/vnd.openxmlformats-officedocument.wordprocessingml.document", False),
("", True), # unlabeled: defer to the binary-char fallback
("", True),
(None, True),
],
)
def test_is_texty_content_type(content_type, expected):
assert tools._is_texty_content_type(content_type) is expected
# ── fetcher end-to-end (mocked network) ──
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):
@ -100,12 +88,6 @@ def test_pdf_rejected_by_content_type(monkeypatch):
assert "non-text content" in out and "application/pdf" in out
def test_image_rejected_by_content_type(monkeypatch):
out = _fetch_with(monkeypatch, b"\x89PNG\r\n\x1a\n" + bytes(range(256)) * 4, "image/png")
assert "<EFBFBD>" not in out
assert "non-text content" in out and "image/png" 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")
@ -113,54 +95,31 @@ def test_text_octet_stream_kept_after_sniffing(monkeypatch):
assert "non-text content" not in out and "binary content" not in out
def test_binary_octet_stream_rejected_after_sniffing(monkeypatch):
out = _fetch_with(monkeypatch, bytes(range(256)) * 20, "application/octet-stream")
assert "binary content" in out
def test_binary_mislabeled_as_text_caught_by_fallback(monkeypatch):
# A server sends binary but labels it text/plain -> the type check passes,
# so the binary-char fallback must catch it.
out = _fetch_with(monkeypatch, bytes(range(256)) * 20, "text/plain")
@pytest.mark.parametrize("content_type", ["application/octet-stream", "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
def test_valid_utf8_binary_caught_by_control_chars(monkeypatch):
# NUL/control-heavy binary is valid UTF-8, so it decodes with no U+FFFD;
# the fallback must still catch it via control-char density (codex #2).
body = bytes([0, 1, 2, 3, 4, 5, 6, 7]) * 400 # all valid UTF-8, 0 replacement chars
# 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
def test_pdf_mislabeled_as_text_caught_by_magic(monkeypatch):
# A PDF served as text/plain: the content-type gate passes, so the magic-byte
# sniff must catch it (its ASCII-heavy head can stay under the ratio).
pdf = b"%PDF-1.7\n" + b"1 0 obj<</Type/Catalog>>endobj\n" * 300
out = _fetch_with(monkeypatch, pdf, "text/plain")
assert "binary content" in out
def test_zip_mislabeled_as_text_caught_by_magic(monkeypatch):
out = _fetch_with(monkeypatch, b"PK\x03\x04" + b"filename.txt content " * 200, "text/plain")
assert "binary content" in out
@pytest.mark.parametrize(
"magic",
[b"\x1f\x8b", b"BZh", b"\xfd7zXZ\x00", b"\x28\xb5\x2f\xfd"], # gzip, bzip2, xz, zstd
[b"%PDF-", b"PK\x03\x04", b"\x1f\x8b", b"BZh", b"\xfd7zXZ\x00", b"\x28\xb5\x2f\xfd"],
)
def test_compression_mislabeled_as_text_caught_by_magic(monkeypatch, magic):
# A small compressed body mislabeled as text can slip past the density check;
# the signature must reject it.
out = _fetch_with(monkeypatch, magic + b"some short compressed-looking body", "text/plain")
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
def test_latin1_text_without_charset_kept(monkeypatch):
# Accent-dense cp1252 text with no declared charset decodes to many U+FFFD as
# UTF-8; the cp1252 retry must keep it instead of dropping it as binary.
# The cp1252 retry should rescue accent-heavy text with ASCII structure.
body = (
"Muller lauft uber die Strasse: schoene, groesse. MARKERWORD ".replace("ue", "ü")
+ "äöüß éèà "
@ -170,36 +129,20 @@ def test_latin1_text_without_charset_kept(monkeypatch):
assert "MARKERWORD" in out
def test_control_heavy_binary_survives_cp1252_retry(monkeypatch):
# Genuine control-heavy binary (no magic) must stay rejected even after the
# cp1252 retry, which cannot rescue real binary.
out = _fetch_with(monkeypatch, bytes([0, 1, 2, 3, 4, 5, 6, 7]) * 400, "text/plain")
assert "binary content" in out
def test_high_byte_binary_not_rescued_as_cp1252(monkeypatch):
# cp1252 maps almost every high byte to a printable character. Without the
# ASCII-structure requirement, this binary payload would be returned as text.
# 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):
# A text log with per-token ANSI color codes is text; ESC is excluded from
# the binary-char set so it isn't dropped as binary.
# 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_binary_unlabeled_caught_by_fallback(monkeypatch):
# No Content-Type coerces to text/plain upstream; the fallback still catches it.
out = _fetch_with(monkeypatch, bytes(range(256)) * 20, None)
assert "<EFBFBD>" not in out
assert "binary content" 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")
@ -208,8 +151,7 @@ def test_html_page_unaffected(monkeypatch):
def test_content_type_sanitized_in_message(monkeypatch):
# An obs-folded Content-Type can smuggle control chars into get_content_type();
# the returned message must be trimmed to a clean MIME token.
# 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
@ -219,23 +161,18 @@ def test_content_type_sanitized_in_message(monkeypatch):
@pytest.mark.parametrize(
"n_bad,n_total,expect_binary",
[
# Straddle the 12.5% ratio (well above the 16-char floor): just under vs
# just over len//8. Locks the divisor so it can't silently drift.
(120, 1000, False), # 120 <= 1000//8 (125) -> kept
(130, 1000, True), # 130 > 1000//8 (125) -> binary
(120, 1000, False),
(130, 1000, True),
],
)
def test_binary_char_ratio_boundary(monkeypatch, n_bad, n_total, expect_binary):
# Body of n_total chars: n_bad NUL control bytes + ASCII filler. NUL stays
# binary through the cp1252 retry, so only the ratio threshold decides here.
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):
# A mostly-clean page with a handful of bad bytes stays (below the floor),
# so we don't drop legitimate pages over minor encoding glitches.
# 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