fix(harness): decode a Python file before searching it when its codec can spell a name (TH-9490)

.github/workflows/sdk-smoke.yml

@@ -95,6 +95,7 @@ jobs:
           tests/harness/test_hosted_entrypoint_provision_retry.py
           tests/harness/test_instrumentation_evidence.py
           tests/harness/test_agent_type_refusal.py
+          tests/harness/test_python_source_codecs.py
           tests/harness/test_executor_completed_counts.py
           tests/harness/test_reading_entry.py
           tests/harness/test_agent_reading.py

src/fi/alk/harness/credentials.py

@@ -26,7 +26,7 @@
 
 from fi.simulate.runtime.spec import SecretRef
 
-from .python_source import ParseBudget, parsed, source_encoding, utf8_size
+from .python_source import EscapeBudget, ParseBudget, decoded_first, parsed, utf8_size
 from .secrets import process_secret_environment_names
 from .source_folder import SubfolderNotFound, agent_folder, agent_scope, in_agent_scope, is_within
 
@@ -1434,6 +1434,10 @@ def script_module(path: Path, name: str) -> Path | None:
                 return hit
         return None
 
+    # One total for the whole walk: bundling has no lookup limit and no parse total, so this is
+    # all that bounds what decoding first costs there, a unicode-escape file's second reading
+    # included (at most this much more parsing per walk).
+    escapes = EscapeBudget()
     seen: set[Path] = set()
     pending = deque(sorted(_code_in(folder)))
     own = set(pending)  # Read once already, so importing one doesn't queue it again.
@@ -1460,25 +1464,62 @@ def script_module(path: Path, name: str) -> Path | None:
         if len(raw) > _MAX_FILE_BYTES:
             continue  # Grew since it was looked at.
         bytes_read += len(raw)
-        # As read_text reads it, line ends included; a .py in the encoding its coding line names,
-        # as Python reads it.
-        encoding = (source_encoding(raw) if path.suffix == ".py" else None) or "utf-8-sig"
-        try:
-            # Strictly: replacing runs Python's own error handler for each byte in most codecs,
-            # 60 ms a megabyte, where UTF-8's is in C. Python itself refuses such a file.
-            text = raw.decode(encoding)
-        except (UnicodeError, LookupError):
-            # Read as UTF-8, as is a codec source_encoding refuses: a crafted file Python reads in
-            # one of those can then give fewer imports here than it does to Python.
-            text = raw.decode("utf-8-sig", errors="replace")
-        content = text.replace("\r\n", "\n").replace("\r", "\n")
+        # A .py in a codec that can spell a name in other bytes is read as Python reads it while
+        # the walk's total lasts: decoded after its line ends, so hz's `~` before a CRLF joins.
+        taken, decoded, encoding = (
+            decoded_first(raw, escapes) if path.suffix == ".py" else (False, None, None)
+        )
+        read_before = None
+        if taken:
+            if decoded is None:
+                # Python refuses the file, so it imports nothing.
+                continue
+            content = decoded.replace("\r\n", "\n").replace("\r", "\n")
+            if encoding is None and "\r" in decoded:
+                # unicode-escape is also read as UTF-8, where `\r` is two characters: decoded, it's
+                # a CR the parser ends a comment or a string at where Python reads on, so what that
+                # reading imports is followed too.
+                read_before = raw.decode("utf-8-sig", errors="replace")
+                read_before = read_before.replace("\r\n", "\n").replace("\r", "\n")
+        else:
+            # As read_text reads it, line ends included; a .py in the encoding its coding line
+            # names, as Python reads it.
+            encoding = encoding or "utf-8-sig"
+            try:
+                # Strictly: replacing runs Python's own error handler for each byte in most
+                # codecs, 60 ms a megabyte, where UTF-8's is in C. Python itself refuses such a
+                # file.
+                text = raw.decode(encoding)
+            except (UnicodeError, LookupError):
+                # Read as UTF-8, as is a codec refused by name (and unicode-escape past the
+                # total): a crafted file Python reads in one of those can then give fewer imports
+                # here than it does to Python.
+                text = raw.decode("utf-8-sig", errors="replace")
+            content = text.replace("\r\n", "\n").replace("\r", "\n")
         hits: list[Path] = []
         if path.suffix == ".py":
             # A file with no import in it can't lead anywhere, so it doesn't use up the parse total
             # before the files that do (Pine's TH-9490 review); a file the total cuts says so.
             if "import" not in content:
                 continue
-            names = list(dict.fromkeys(_python_imports(content, parse_limit, parse_budget)))
+            readings = [content] if read_before is None else [content, read_before]
+            larger = max(readings, key=utf8_size)
+            if len(readings) == 1 or utf8_size(larger) > parse_limit:
+                # Each reading is charged as it's parsed; one past the limit is neither.
+                found = [
+                    name
+                    for reading in readings
+                    for name in _python_imports(reading, parse_limit, parse_budget)
+                ]
+            elif parse_budget is None or parse_budget.spend(larger):
+                # One file read two ways is charged once, as the larger reading, so it spends no
+                # more of the total than it did read one way.
+                found = [
+                    name for reading in readings for name in _python_imports(reading, parse_limit)
+                ]
+            else:
+                found = []
+            names = list(dict.fromkeys(found))
             cut |= parse_budget is not None and parse_budget.left < 0
             # One file can hold more names than the walk may try; the rest aren't tried.
             cut |= len(names) > import_limit - imports_tried

src/fi/alk/harness/python_source.py

@@ -32,10 +32,11 @@
 # token is charged its line. This much is at most about 0.03 s of that work on 3.12.3. Over 5,000
 # real files of 4 to 64 KB, 64 KB of code comes to about 0.5 M, and 1.9 M at the most.
 LARGEST_TOKENIZED_WORK = 8 * 1024 * 1024
-# What one scan decodes in unicode-escape, in bytes. Python runs such a file, and its escapes can
-# spell any name (`import \x70laywright`), so it's decoded as Python does; but its decode looks up
-# each \N{...} name, and 3,000 files of 128 KB held the save 6-7 s with no total. This much
-# decodes in about 0.05 s; past it, the file is searched as bytes.
+# What one scan decodes before searching, in bytes. Python runs a file in a codec that can spell a
+# name in other bytes (`import \x70laywright` in unicode-escape, `import +AHA-laywright` in
+# utf-7), so such a file is decoded as Python does; but unicode-escape's decode looks up each
+# \N{...} name, and 3,000 files of 128 KB held the save 6-7 s with no total. This much decodes in
+# about 0.05 s; past it, the file is searched as bytes.
 LARGEST_ESCAPED_TOTAL = 4 * 1024 * 1024
 
 
@@ -45,6 +46,99 @@
 _LONGEST_MARK_RUN = 32
 # Text codecs a coding line may name that aren't read: see _coding.
 _UNREAD_CODECS = frozenset({"idna", "punycode", "undefined", "unicode-escape"})
+# The codecs a coding line may name in which no byte form spells a name character, NUL, CR or LF
+# but its own single byte, nor spells nothing: a needle missing from the raw bytes is missing from
+# the text Python reads too. The sweep in test_python_source_codecs proves each one, so a
+# codec missing from here is decoded before it's searched.
+_RAW_SAFE_CODECS = frozenset(
+    {
+        "ascii",
+        "big5",
+        "big5hkscs",
+        "charmap",
+        "cp1006",
+        "cp1125",
+        "cp1250",
+        "cp1251",
+        "cp1252",
+        "cp1253",
+        "cp1254",
+        "cp1255",
+        "cp1256",
+        "cp1257",
+        "cp1258",
+        "cp437",
+        "cp720",
+        "cp737",
+        "cp775",
+        "cp850",
+        "cp852",
+        "cp855",
+        "cp856",
+        "cp857",
+        "cp858",
+        "cp860",
+        "cp861",
+        "cp862",
+        "cp863",
+        "cp864",
+        "cp865",
+        "cp866",
+        "cp869",
+        "cp874",
+        "cp932",
+        "cp949",
+        "cp950",
+        "euc_jis_2004",
+        "euc_jisx0213",
+        "euc_jp",
+        "euc_kr",
+        "gb18030",
+        "gb2312",
+        "gbk",
+        "hp-roman8",
+        "iso8859-1",
+        "iso8859-10",
+        "iso8859-11",
+        "iso8859-13",
+        "iso8859-14",
+        "iso8859-15",
+        "iso8859-16",
+        "iso8859-2",
+        "iso8859-3",
+        "iso8859-4",
+        "iso8859-5",
+        "iso8859-6",
+        "iso8859-7",
+        "iso8859-8",
+        "iso8859-9",
+        "johab",
+        "koi8-r",
+        "koi8-t",
+        "koi8-u",
+        "kz1048",
+        "mac-croatian",
+        "mac-cyrillic",
+        "mac-greek",
+        "mac-iceland",
+        "mac-latin2",
+        "mac-roman",
+        "mac-romanian",
+        "mac-turkish",
+        "palmos",
+        "ptcp154",
+        "shift_jis",
+        "shift_jis_2004",
+        "shift_jisx0213",
+        "tis-620",
+        "utf-8",
+        "utf-8-sig",
+    }
+)
+# Refused by name in _coding: never decoded, whatever the total.
+_NEVER_DECODED_FIRST = frozenset({"idna", "punycode", "undefined"})
+# Python reads a lone CR as a line end before it decodes, so hz's `~` and a line end join.
+_CR_TO_LF = bytes.maketrans(b"\r", b"\n")
 _NON_ASCII_RUN = re.compile(r"[^\x00-\x7f]{2,}")
 
 
@@ -65,7 +159,8 @@ def spend(self, text: str) -> bool:
 
 
 class EscapeBudget:
-    """What is left for one scan to decode in unicode-escape, in bytes."""
+    """What is left for one scan to decode before searching, in bytes: every file read that way (a
+    text codec not proven safe to search as bytes, unicode-escape among them) is charged here."""
 
     def __init__(self, total: int | None = None) -> None:
         self.left = LARGEST_ESCAPED_TOTAL if total is None else total
@@ -114,11 +209,12 @@ def read_module(
     in the encoding its coding line or byte-order mark names (UTF-8 without one), or, given
     ``needles``, naming none of them.
 
-    Given ``escapes``, a file in unicode-escape is decoded as Python does while the scan's total
-    lasts, and the needles are looked for in what it decodes to. Given ``undecoded``, a file that
-    names a needle but isn't read because its coding line names a codec refused by name (or
-    unicode-escape past the total) is put there as bytes: Python still runs it (plain ASCII under
-    idna as it would without the line), so a scan looking for names can still find them.
+    Given ``escapes``, a file in a codec not proven safe to search as bytes (unicode-escape,
+    utf-7, hz and others) is decoded as Python does while the scan's total lasts, and a needle
+    counts in what it decodes to, or in its bytes when that isn't ASCII. Given ``undecoded``, a file that names a needle
+    but isn't read because its coding line names a codec refused by name (or unicode-escape past
+    the total) is put there as bytes: Python still runs it (plain ASCII under idna as it would
+    without the line), so a scan looking for names can still find them.
     """
     try:
         # A link can name a device or a pipe, which reads without end.
@@ -131,18 +227,17 @@ def read_module(
     if len(raw) > limit:
         return None
     coding = _coding(raw)
-    if (
-        coding is not None
-        and coding[1].name == "unicode-escape"
-        and escapes is not None
-        and escapes.spend(raw)
-    ):
-        try:
-            # Python ends lines first, so a backslash before a CRLF joins the lines it ends.
-            text = raw.replace(b"\r\n", b"\n").replace(b"\r", b"\n").decode("unicode-escape")
-        except UnicodeError:
+    taken, text = _decoded_first(raw, coding, escapes)
+    if taken:
+        if text is None:
             return None
-        if needles and not any(needle.decode("ascii") in text for needle in needles):
+        # The bytes too, as for any other file: the parser folds a fullwidth name the text spells
+        # (`ｐlaywright`) to the needle, and a needle in the bytes is what lets such a file in. On
+        # ASCII text nothing folds, so whatever a scan can find is in the text itself.
+        if needles and not any(
+            needle.decode("ascii") in text or not text.isascii() and needle in raw
+            for needle in needles
+        ):
             return None
         return text
     # A literal match: a name spelled in other characters that Python folds to the same one,
@@ -177,6 +272,58 @@ def source_encoding(raw: bytes) -> str | None:
     return encoding
 
 
+def decoded_first(
+    raw: bytes, escapes: EscapeBudget | None
+) -> tuple[bool, str | None, str | None]:
+    """``raw`` read as ``read_module`` reads a file before searching it, for a reader that
+    searches no needles: whether it was decoded first, its text (None when Python refuses it),
+    and the encoding ``source_encoding`` gives, to read a file that wasn't. The coding line is
+    read once for both."""
+    coding = _coding(raw)
+    taken, text = _decoded_first(raw, coding, escapes)
+    encoding = None if coding is None or coding[1].name in _UNREAD_CODECS else coding[0]
+    return taken, text, encoding
+
+
+def _decoded_first(
+    raw: bytes, coding: tuple[str, codecs.CodecInfo] | None, escapes: EscapeBudget | None
+) -> tuple[bool, str | None]:
+    """Whether ``raw``, in the codec ``coding`` names, is decoded before it's searched, and its
+    text as Python reads it: None when Python refuses it.
+
+    A codec not proven safe to search as bytes can spell a name in other bytes, so its file is
+    decoded while ``escapes`` lasts. Past it, or with no total, it isn't taken, and the caller
+    reads it as any other file.
+    """
+    if coding is None or escapes is None:
+        return False, None
+    encoding, info = coding
+    if (
+        info.name in _RAW_SAFE_CODECS
+        or info.name in _NEVER_DECODED_FIRST
+        # Not a text encoding, such as rot13 or base64: bytes.decode refuses it before reading.
+        or not getattr(info, "_is_text_encoding", True)
+    ):
+        return False, None
+    if not escapes.spend(raw):
+        return False, None
+    if b"\r" in raw:
+        # Python ends lines before it decodes, so a backslash or hz's `~` before a CRLF joins
+        # the lines it ends.
+        raw = raw.replace(b"\r\n", b"\n").translate(_CR_TO_LF)
+    if not raw.endswith(b"\n"):
+        # Python ends the last line too: a file ending in an open escape still runs.
+        raw += b"\n"
+    try:
+        text = raw.decode(encoding)
+    except (UnicodeError, LookupError):
+        # LookupError only from a codec registered at run time: Python's text codecs don't.
+        return True, None
+    # Python reads nothing past a NUL its decode makes at a line's start, and refuses the file
+    # for one elsewhere: so a name after it is never run.
+    return True, text.partition("\0")[0]
+
+
 def _first_lines(raw: bytes) -> bytes:
     """The first two lines of ``raw``, each ended as Python ends a line (LF, CRLF or a lone CR),
     joined by LF: found where they end, so the rest of the file is never rewritten."""

tests/harness/test_agent_type_refusal.py

@@ -701,9 +701,11 @@ def test_a_file_in_a_refused_codec_still_names_its_tools(
     assert (decoded, labels, parsed) == ([], [], [])
 
 
-def test_an_escaped_import_in_a_unicode_escape_file_is_found(tmp_path):
+@pytest.mark.parametrize("ending", [b"\n", b""], ids=["final-lf", "no-final-lf"])
+def test_an_escaped_import_in_a_unicode_escape_file_is_found(tmp_path, ending):
     # Python decodes the escapes before it parses, so the import is real and so is the finding.
-    (tmp_path / "helper.py").write_bytes(b"# coding: unicode-escape\nimport \\x70laywright\n")
+    # It ends the last line itself, so a file without a final line end runs the same.
+    (tmp_path / "helper.py").write_bytes(b"# coding: unicode-escape\nimport \\x70laywright" + ending)
     assert browser_tools_found(tmp_path) == ["playwright"]