truf-server/tests/owned_process_helper.py

import os
import subprocess
import sys
import time

sys.dont_write_bytecode = True


CREATE_NO_WINDOW = 0x08000000


def write_value(path, value):
    with open(path, "w", encoding="ascii") as output:
        output.write(str(value))
        output.flush()


def wait_for_file(path, timeout=5):
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        if os.path.exists(path):
            return True
        time.sleep(0.02)
    return False


def main():
    mode = sys.argv[1]
    if mode == "forked-proxy":
        from pathlib import Path

        sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "app"))
        from owned_process import OwnedProcess

        marker_dir, release_path = Path(sys.argv[2]), sys.argv[3]
        child = OwnedProcess(
            [sys.executable, "-I", "-S", "-B", os.path.abspath(__file__),
             "linger", str(marker_dir / "payload-0.pid")],
            stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
        )
        write_value(marker_dir / "host-1.pid", child.host_pid)
        write_value(marker_dir / "payload-2.pid", os.getpid())
        if not wait_for_file(marker_dir / "payload-0.pid"):
            return 90
        child._control_lock.acquire()
        fork_pid = os.fork()
        if fork_pid == 0:
            del child
            os._exit(0)
        child._control_lock.release()
        write_value(marker_dir / "payload-1.pid", fork_pid)
        _, status = os.waitpid(fork_pid, 0)
        if status != 0 or child.poll() is not None:
            return 92
        write_value(marker_dir / "fork-detached", 1)
        while not os.path.exists(release_path):
            time.sleep(0.02)
        os._exit(23)

    if mode == "adopted-observer":
        from pathlib import Path

        marker_dir, release_path = Path(sys.argv[2]), sys.argv[3]
        owner = subprocess.Popen(
            [sys.executable, "-I", "-S", "-B", os.path.abspath(__file__),
             "nested-owned", "1", str(marker_dir), str(marker_dir / "orphan-release")],
            stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
        )
        write_value(marker_dir / "payload-2.pid", os.getpid())
        owner.wait()
        while not os.path.exists(release_path):
            time.sleep(0.02)
        return 23

    if mode == "nested-owned":
        from pathlib import Path

        sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "app"))
        from owned_process import OwnedProcess

        depth, marker_dir, release_path = int(sys.argv[2]), Path(sys.argv[3]), sys.argv[4]
        if depth:
            child = OwnedProcess(
                [sys.executable, "-I", "-S", "-B", os.path.abspath(__file__),
                 mode, str(depth - 1), str(marker_dir), ""],
                stdin=subprocess.DEVNULL,
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
            )
            write_value(marker_dir / f"host-{depth}.pid", child.host_pid)
        write_value(marker_dir / f"payload-{depth}.pid", os.getpid())
        while not release_path or not os.path.exists(release_path):
            time.sleep(0.02)
        os._exit(23)

    if mode == "linger":
        write_value(sys.argv[2], os.getpid())
        while True:
            time.sleep(1)

    if mode == "escaped-session":
        from pathlib import Path

        marker_dir = Path(sys.argv[2])
        child = subprocess.Popen(
            [sys.executable, "-I", "-S", "-B", os.path.abspath(__file__),
             "linger", str(marker_dir / "escaped.pid")],
            stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL, start_new_session=True,
        )
        write_value(marker_dir / "payload-0.pid", os.getpid())
        if not wait_for_file(marker_dir / "escaped.pid"):
            child.kill()
            return 90
        while True:
            time.sleep(1)

    if mode in ("tree", "spawn-and-exit"):
        ready_path = sys.argv[2]
        child = subprocess.Popen(
            [sys.executable, os.path.abspath(__file__), "linger", ready_path],
            stdin=subprocess.DEVNULL,
            close_fds=True,
            creationflags=CREATE_NO_WINDOW if os.name == "nt" else 0,
        )
        if not wait_for_file(ready_path):
            child.terminate()
            return 90
        if mode == "spawn-and-exit":
            return int(sys.argv[3])
        while True:
            time.sleep(1)

    if mode == "marker":
        write_value(sys.argv[2], os.getpid())
        return 0

    if mode == "stdio":
        sys.stdout.write("owned stdout\n")
        sys.stdout.flush()
        sys.stderr.write("owned stderr\n")
        sys.stderr.flush()
        if len(sys.argv) > 3 and sys.argv[3] == "wait-stdin":
            sys.stdin.buffer.read()
        return int(sys.argv[2])

    if mode == "partial":
        sys.stdout.write("partial stdout\n")
        sys.stdout.flush()
        sys.stderr.write("partial stderr\n")
        sys.stderr.flush()
        while True:
            time.sleep(1)

    return 91


if __name__ == "__main__":
    raise SystemExit(main())