diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py index 8f9241d51f..9f2e7a0f62 100644 --- a/studio/backend/core/inference/tools.py +++ b/studio/backend/core/inference/tools.py @@ -3602,6 +3602,15 @@ _BINARY_MAGIC = ( 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)) @@ -3616,7 +3625,11 @@ def _looks_binary(text: str) -> bool: def _has_binary_magic(data: bytes) -> bool: """Whether a common binary signature follows optional BOM or whitespace.""" - head = data[:1024].lstrip().removeprefix(b"\xef\xbb\xbf").lstrip() + 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) @@ -3858,7 +3871,11 @@ def _fetch_page_text( declared = resp.headers.get_content_charset() declared_codec = codecs.lookup(declared).name if declared else None - raw_html = raw_bytes.decode(declared or "utf-8", errors = "replace") + 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): diff --git a/studio/backend/tests/test_web_fetch_binary_guard.py b/studio/backend/tests/test_web_fetch_binary_guard.py index 411db2276f..57159c024f 100644 --- a/studio/backend/tests/test_web_fetch_binary_guard.py +++ b/studio/backend/tests/test_web_fetch_binary_guard.py @@ -5,6 +5,7 @@ from __future__ import annotations +import codecs import sys from email.message import Message from pathlib import Path @@ -122,6 +123,23 @@ def test_excel_labeled_csv_kept_after_sniffing(monkeypatch): 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 @@ -146,7 +164,18 @@ def test_text_labeled_binary_caught_by_magic(monkeypatch, magic): assert "binary content" in out -@pytest.mark.parametrize("prefix", [b"\xef\xbb\xbf", b" \r\n", b"\t\xef\xbb\xbf "]) +@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<>endobj\n" * 100 out = _fetch_with(monkeypatch, body, "text/plain")