From c4404dd2a1778ef76f29ee2ad592d1835632ac12 Mon Sep 17 00:00:00 2001 From: oobabooga <112222186+oobabooga@users.noreply.github.com> Date: Tue, 14 Jul 2026 22:07:50 -0300 Subject: [PATCH] Harden web fetch binary sniffing --- studio/backend/core/inference/tools.py | 44 +++++++++++++++---- .../tests/test_web_fetch_binary_guard.py | 39 +++++++++++----- 2 files changed, 64 insertions(+), 19 deletions(-) diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py index 93826bf2af..937219fbf1 100644 --- a/studio/backend/core/inference/tools.py +++ b/studio/backend/core/inference/tools.py @@ -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: diff --git a/studio/backend/tests/test_web_fetch_binary_guard.py b/studio/backend/tests/test_web_fetch_binary_guard.py index 92dd9411be..d69526e898 100644 --- a/studio/backend/tests/test_web_fetch_binary_guard.py +++ b/studio/backend/tests/test_web_fetch_binary_guard.py @@ -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)