Harden web fetch binary sniffing

This commit is contained in:
oobabooga 2026-07-14 22:07:50 -03:00
commit c4404dd2a1
2 changed files with 64 additions and 19 deletions

View file

@ -1405,6 +1405,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.
_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
@ -1413,6 +1420,13 @@ def _looks_binary(text: str) -> bool:
_MIN_BINARY_CHARS, len(text) // _BINARY_CHAR_DIVISOR
)
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",
@ -1522,15 +1536,18 @@ _TEXT_APPLICATION_SUBTYPES = frozenset(
def _is_texty_content_type(content_type: str | None) -> bool:
"""True for MIME types safe to decode as text (HTML pages, plain text,
JSON/XML/YAML feeds). Binary types (PDF, images, archives, Office/ZIP,
octet-stream) return False so the fetcher never decodes them into a flood of
replacement chars that poison the model context (unslothai/unsloth#7084).
"""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.
"""
ct = (content_type or "").lower()
ct = (content_type or "").partition(";")[0].strip().lower()
if not ct:
return True # unlabeled: let the binary-char fallback decide
if ct.startswith("text/"):
@ -1540,7 +1557,11 @@ def _is_texty_content_type(content_type: str | None) -> bool:
# application/vnd.openxmlformats-... isn't read as xml. RFC 6839
# +json/+xml suffixes (ld+json, xhtml+xml, ...) pass too.
subtype = ct[len("application/") :].removeprefix("x-")
return subtype in _TEXT_APPLICATION_SUBTYPES or subtype.endswith(("+json", "+xml"))
return (
subtype == "octet-stream"
or subtype in _TEXT_APPLICATION_SUBTYPES
or subtype.endswith(("+json", "+xml"))
)
return False
@ -1648,9 +1669,14 @@ def _fetch_page_text(
# without replacement chars: a real text page has few binary chars.
if _looks_binary(raw_html):
# An undeclared non-UTF-8 page (latin-1/cp1252) also decodes to many
# U+FFFD. Retry as cp1252 (maps almost every byte to a printable
# char): if that reads as text it was a charset mismatch, not binary.
alt = raw_bytes.decode("cp1252", "replace") if declared is None else None
# 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.
alt = (
raw_bytes.decode("cp1252", "replace")
if declared is None and _has_single_byte_text_evidence(raw_bytes)
else None
)
if alt is not None and not _looks_binary(alt):
raw_html = alt
else:

View file

@ -2,10 +2,9 @@
# 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 (PDF, image, octet-stream) into a flood of U+FFFD replacement
chars that poison the model context. It rejects non-text Content-Types up front
and, for binary mislabeled as text/* or sent unlabeled, falls back to a
replacement-char ratio check. Real HTML pages are unaffected.
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.
"""
from __future__ import annotations
@ -68,6 +67,7 @@ def _fetch_with(monkeypatch, body: bytes, content_type: str | None) -> str:
("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),
@ -78,7 +78,8 @@ def _fetch_with(monkeypatch, body: bytes, content_type: str | None) -> str:
("application/pdf", False),
("image/png", False),
("image/svg+xml", False), # SVG source isn't extracted downstream; reject
("application/octet-stream", False),
# Generic downloads need byte sniffing because they can contain text.
("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),
@ -105,6 +106,18 @@ def test_image_rejected_by_content_type(monkeypatch):
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")
assert "plain text artifact" in out
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.
@ -164,6 +177,14 @@ def test_control_heavy_binary_survives_cp1252_retry(monkeypatch):
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.
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.
@ -189,12 +210,10 @@ 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.
out = _fetch_with(
monkeypatch, b"\x00\x01\x02" * 500, "application/octet-stream\r\n data: injected"
)
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/octet-stream" in out
assert "application/pdf" in out
@pytest.mark.parametrize(
@ -206,7 +225,7 @@ def test_content_type_sanitized_in_message(monkeypatch):
(130, 1000, True), # 130 > 1000//8 (125) -> binary
],
)
def test_replacement_ratio_boundary(monkeypatch, n_bad, n_total, expect_binary):
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)