fix(harness): the save's import walk says when it leaves code unread, and reads more of it (TH-9543)

.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/agent_type_check.py

@@ -157,6 +157,10 @@
 _WALKED_SCRIPTS = 2 * 1024 * 1024
 # The import names it tries then.
 _WALKED_IMPORTS = 5_000
+# And the Python it parses then, each reading of a file charged: twice the other scans' total, so an
+# agent's own code of a few hundred KB isn't cut, while crafted files that only this walk parses
+# hold the save about 1.3 s on 3.12 and 3.1 s on 3.11.
+_WALKED_PARSED = 512 * 1024
 # And the lookups on disk it makes then, a file found costing one for each part of its path: a name
 # can resolve a package for each of its parts, and 1,000 crafted names took the save 12 s.
 _WALKED_LOOKUPS = 100_000
@@ -642,15 +646,16 @@ def _walked(root: Path, subfolder: str) -> frozenset[Path]:
         root,
         subfolder,
         parse_limit=LARGEST_PARSED_MODULE,
-        parse_budget=ParseBudget(),
+        parse_budget=ParseBudget(_WALKED_PARSED),
         script_budget=ParseBudget(_WALKED_SCRIPTS),
         import_limit=_WALKED_IMPORTS,
         lookup_limit=_WALKED_LOOKUPS,
         stopped=stopped,
     )
     if stopped and stopped[0]:
         logger.warning(
-            "the save's import walk for %s stopped at its limits; %d files reached",
+            "the save's import walk for %s stopped at its limits or left code unread; "
+            "%d files reached",
             subfolder,
             len(found),
         )

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."""