Simplify web fetch binary guard
This commit is contained in:
parent
c5902a6266
commit
10592853fc
2 changed files with 38 additions and 130 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue