File size: 4,925 Bytes
f89b1ac
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
from __future__ import annotations

import json
import math
import os
import signal
import subprocess
import sys
import tempfile
from dataclasses import dataclass

DEFAULT_ADDRESS_SPACE_BYTES = 512 * 1024 * 1024
DEFAULT_FILE_BYTES = 2 * 1024 * 1024
DEFAULT_OPEN_FILES = 64
DEFAULT_PROCESSES = 32

_RUNNER_BOOTSTRAP = r"""
import json
import runpy
import sys

try:
    import resource
except ImportError:  # pragma: no cover
    resource = None


def _set_limit(name, value):
    if resource is None or not hasattr(resource, name):
        return
    limit = int(value)
    try:
        _, current_hard = resource.getrlimit(getattr(resource, name))
        soft = min(limit, current_hard) if current_hard >= 0 else limit
        resource.setrlimit(getattr(resource, name), (soft, current_hard))
    except (OSError, ValueError):
        return


config = json.loads(sys.argv[1])
_set_limit("RLIMIT_CORE", 0)
_set_limit("RLIMIT_CPU", config["cpu_seconds"])
_set_limit("RLIMIT_FSIZE", config["file_bytes"])
_set_limit("RLIMIT_NOFILE", config["open_files"])
_set_limit("RLIMIT_AS", config["address_space_bytes"])
_set_limit("RLIMIT_NPROC", config["processes"])

mode = config["mode"]
if mode == "script":
    script = sys.argv[2]
    sys.argv = sys.argv[2:]
    runpy.run_path(script, run_name="__main__")
elif mode == "code":
    sys.argv = ["-c"]
    exec(config["code"], {"__name__": "__main__"})
else:  # pragma: no cover
    raise SystemExit(f"Unsupported execution mode: {mode}")
"""


@dataclass(frozen=True)
class PythonRunResult:
    returncode: int
    stdout: str
    stderr: str
    timed_out: bool = False


def _safe_env(workspace_dir: str) -> dict[str, str]:
    return {
        "HOME": workspace_dir,
        "TMPDIR": workspace_dir,
        "LANG": "C.UTF-8",
        "LC_ALL": "C.UTF-8",
        "PATH": "",
        "PYTHONDONTWRITEBYTECODE": "1",
        "PYTHONHASHSEED": "0",
        "PYTHONIOENCODING": "utf-8",
        "PYTHONNOUSERSITE": "1",
    }


def _limit_config(timeout_s: float) -> dict[str, int]:
    return {
        "cpu_seconds": max(1, int(math.ceil(timeout_s)) + 1),
        "file_bytes": DEFAULT_FILE_BYTES,
        "open_files": DEFAULT_OPEN_FILES,
        "address_space_bytes": DEFAULT_ADDRESS_SPACE_BYTES,
        "processes": DEFAULT_PROCESSES,
    }


def _read_limited_text(handle, limit: int) -> str:
    handle.seek(0)
    data = handle.read(limit + 1)
    if isinstance(data, bytes):
        return data.decode("utf-8", errors="replace")[:limit]
    return str(data)[:limit]


def _terminate_process(proc: subprocess.Popen[bytes]) -> None:
    if proc.poll() is not None:
        return
    if os.name != "nt":
        try:
            os.killpg(proc.pid, signal.SIGKILL)
            return
        except ProcessLookupError:
            return
    proc.kill()


def _run_python_command(
    config: dict[str, object],
    *,
    cwd: str,
    argv: list[str],
    timeout_s: float,
    stdout_limit: int,
    stderr_limit: int,
) -> PythonRunResult:
    command = [
        sys.executable,
        "-I",
        "-B",
        "-c",
        _RUNNER_BOOTSTRAP,
        json.dumps(config, ensure_ascii=True),
        *argv,
    ]
    start_new_session = os.name != "nt"

    with tempfile.TemporaryFile() as stdout_file, tempfile.TemporaryFile() as stderr_file:
        proc = subprocess.Popen(
            command,
            cwd=cwd,
            env=_safe_env(cwd),
            stdin=subprocess.DEVNULL,
            stdout=stdout_file,
            stderr=stderr_file,
            start_new_session=start_new_session,
        )
        timed_out = False
        try:
            proc.wait(timeout=timeout_s)
        except subprocess.TimeoutExpired:
            timed_out = True
            _terminate_process(proc)
            proc.wait()

        return PythonRunResult(
            returncode=proc.returncode if proc.returncode is not None else -1,
            stdout=_read_limited_text(stdout_file, stdout_limit),
            stderr=_read_limited_text(stderr_file, stderr_limit),
            timed_out=timed_out,
        )


def run_python_script(
    script_name: str,
    *,
    cwd: str,
    args: list[str],
    timeout_s: float,
    stdout_limit: int,
    stderr_limit: int,
) -> PythonRunResult:
    config = {"mode": "script", **_limit_config(timeout_s)}
    return _run_python_command(
        config,
        cwd=cwd,
        argv=[script_name, *args],
        timeout_s=timeout_s,
        stdout_limit=stdout_limit,
        stderr_limit=stderr_limit,
    )


def run_python_code(
    code: str,
    *,
    cwd: str,
    timeout_s: float,
    stdout_limit: int,
    stderr_limit: int,
) -> PythonRunResult:
    config = {"mode": "code", "code": code, **_limit_config(timeout_s)}
    return _run_python_command(
        config,
        cwd=cwd,
        argv=[],
        timeout_s=timeout_s,
        stdout_limit=stdout_limit,
        stderr_limit=stderr_limit,
    )