Open_Ultra_SAST/tests/test_engine_task_entry.py

"""Offline AX worker entry controls; HTTP is loopback only."""

import importlib.util
import io
import json
import tarfile
import threading
import time
import urllib.error
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path

import pytest

SCRIPTS = Path(__file__).resolve().parents[1] / "benchmarks/learn"


def load(name):
    with pytest.MonkeyPatch.context() as patch:
        patch.syspath_prepend(str(SCRIPTS))
        spec = importlib.util.spec_from_file_location(name, SCRIPTS / f"{name}.py")
        module = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(module)
        return module


def test_engine_cluster_heap_flags(monkeypatch):
    worker = load("engine_trace_worker")
    monkeypatch.setenv("JAVA_TOOL_OPTIONS", "-Xmx9999m")
    backend = worker.MeasuredBackend(30, 10, lambda: None, heap_profile="cluster")
    flags = (
        "-Xms1700m -Xmx1700m -XX:MaxMetaspaceSize=256m -XX:ReservedCodeCacheSize=128m "
        "-XX:MaxDirectMemorySize=256m -Xss512k -XX:ActiveProcessorCount=2 -XX:+UseG1GC"
    )
    assert backend._jvm_env()["JAVA_TOOL_OPTIONS"] == flags
    assert backend._heap_flag() == "-J-Xmx1700m"
    assert worker.MeasuredBackend(30, 10, lambda: None)._heap_flag() == "-J-Xmx2560m"


@pytest.mark.parametrize("readiness", ["ready", "delayed", "closed", "source_denied"])
@pytest.mark.parametrize("failure", [False, True])
@pytest.mark.parametrize("transport", ["memory", "loopback"])
def test_task_entry_http_roundtrip(tmp_path, monkeypatch, failure, transport, readiness):
    entry = load("engine_task_entry")
    archive = io.BytesIO()
    with tarfile.open(fileobj=archive, mode="w") as tar:
        data = b"print('input proof')\n"
        info = tarfile.TarInfo("app.py")
        info.size = len(data)
        tar.addfile(info, io.BytesIO(data))
    pin = {"pin": "abc", "units": [{"unit": "u", "supported": True}], "questions": {"u": {}}}
    uploads = []
    gets = []
    puts = []
    started = time.monotonic()

    class Handler(BaseHTTPRequestHandler):
        def do_GET(self):
            gets.append((self.path, time.monotonic() - started))
            if readiness == "source_denied" and self.path == "/source":
                self.send_response(403)
                self.end_headers()
                return
            if readiness == "closed" or (readiness == "delayed" and gets[-1][1] < 3):
                self.send_response(403 if len(gets) % 2 else 503)
                self.end_headers()
                return
            data = archive.getvalue() if self.path == "/source" else json.dumps(pin).encode()
            self.send_response(200)
            self.end_headers()
            self.wfile.write(data)

        def do_PUT(self):
            body = self.rfile.read(int(self.headers["Content-Length"]))
            puts.append(time.monotonic())
            if readiness == "closed" and len(puts) == 1:
                self.send_response(403)
                self.end_headers()
                return
            uploads.append(body)
            self.send_response(200)
            self.end_headers()

        def log_message(self, *args):
            pass

    fake_worker = tmp_path / "worker.py"
    fake_worker.write_text(
        "import argparse,json,pathlib,sys\np=argparse.ArgumentParser()\n"
        "for key in ['worker','root','output','deadline','question-deadline','heap-profile']: p.add_argument('--'+key)\n"
        "a=p.parse_args()\nassert a.heap_profile == 'cluster'\n"
        "assert (pathlib.Path(a.root)/'app.py').stat().st_size > 0\n"
        + ("sys.exit(7)\n" if failure else "pathlib.Path(a.output).write_text(json.dumps({'done':True,'units':[]}))\n")
    )
    if transport == "loopback":
        try:
            server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
        except PermissionError:
            pytest.skip("sandbox prohibits loopback sockets; in-memory HTTP roundtrip still runs")
        thread = threading.Thread(target=server.serve_forever, daemon=True)
        thread.start()
        base = f"http://127.0.0.1:{server.server_port}"
    else:
        base = "https://fake-store.invalid"

        def urlopen(request, timeout):
            method = request.get_method() if hasattr(request, "get_method") else "GET"
            url = request.full_url if hasattr(request, "full_url") else request
            body = request.data or b"" if method == "PUT" else b""
            raw = f"{method} {url.removeprefix(base)} HTTP/1.0\r\nContent-Length: {len(body)}\r\n\r\n".encode() + body

            class Socket:
                def __init__(self):
                    self.sent = bytearray()

                def makefile(self, *args):
                    return io.BytesIO(raw)

                def sendall(self, data):
                    self.sent.extend(data)

            connection = Socket()
            Handler(connection, ("127.0.0.1", 0), None)
            headers, payload = bytes(connection.sent).split(b"\r\n\r\n", 1)
            status = int(headers.split()[1])
            if status != 200:
                raise urllib.error.HTTPError(url, status, "denied", {}, io.BytesIO(payload))
            return io.BytesIO(payload)

        monkeypatch.setattr(entry.urllib.request, "urlopen", urlopen)
    try:
        for name, suffix in [("SOURCE_URL", "source"), ("QUESTIONS_URL", "questions"), ("RESULT_URL", "result")]:
            monkeypatch.setenv(name, base + "/" + suffix)
        budget = 0.7 if readiness == "closed" else 60
        assert entry.main(work_root=tmp_path / "work", worker_script=fake_worker, readiness_budget=budget) == 0
        if readiness == "delayed":
            assert len(gets) >= 5
            assert gets[-1][1] >= 3
        if readiness == "closed":
            assert 0.7 <= puts[0] - started < 1.5
            assert len(puts) == 2
            assert puts[1] - puts[0] >= 0.5
            assert all(path == "/questions" for path, _ in gets)
    finally:
        if transport == "loopback":
            server.shutdown()
            thread.join()
            server.server_close()
    assert len(uploads) == 1
    with tarfile.open(fileobj=io.BytesIO(uploads[0])) as tar:
        result = json.load(tar.extractfile("result.json"))
    assert result["done"]
    if readiness == "closed":
        assert result["status"] == "failed"
        assert result["reason"] == "GET questions: egress not open after 60 s"
        assert json.loads((tmp_path / "work/out/result.json").read_text()) == result
    elif readiness == "source_denied":
        assert result["status"] == "failed"
        assert result["reason"] == "GET source: HTTP 403"
        assert result["units"][0]["reason"] == "GET source: HTTP 403"
        assert base not in json.dumps(result)
        assert json.loads((tmp_path / "work/out/result.json").read_text()) == result
    elif failure:
        assert result["units"][0]["status"] == "failed"
        assert "7" in result["units"][0]["reason"]


def test_engine_cluster_profile_recorded_on_input_failure(tmp_path):
    worker = load("engine_trace_worker")
    pin = {
        "units": [{"unit": "u", "file": "missing.py", "language": "python", "supported": True, "reason": ""}],
        "questions": {"u": {}},
    }
    result = worker.worker(pin, tmp_path, tmp_path / "result.json", 30, 10, heap_profile="cluster")
    assert result["units"][0]["status"] == "failed"
    assert result["instrument"]["heap_profile"] == "cluster"
    assert result["instrument"]["heap_mb"] == 1700
    assert result["instrument"]["java_tool_options"] == worker.CLUSTER_JVM_FLAGS


def test_task_entry_rejects_tar_traversal(tmp_path):
    entry = load("engine_task_entry")
    source = tmp_path / "source.tar"
    with tarfile.open(source, "w") as archive:
        member = tarfile.TarInfo("../escape")
        member.size = 1
        archive.addfile(member, io.BytesIO(b"x"))
    with pytest.raises(ValueError, match="unsafe"):
        entry.unpack(source, tmp_path / "case")
    assert not (tmp_path / "escape").exists()


@pytest.mark.parametrize("error", [urllib.error.URLError("refused"), ConnectionResetError("reset")])
def test_readiness_connection_errors_back_off_to_cap(tmp_path, monkeypatch, error):
    entry = load("engine_task_entry")
    now = [0.0]
    pauses = []
    timeouts = []

    def pause(seconds):
        pauses.append(seconds)
        now[0] += seconds

    def refused(request, timeout):
        timeouts.append(timeout)
        raise error

    monkeypatch.setattr(entry.time, "monotonic", lambda: now[0])
    monkeypatch.setattr(entry.time, "sleep", pause)
    monkeypatch.setattr(entry.urllib.request, "urlopen", refused)
    with pytest.raises(RuntimeError, match="^egress not open after 60 s$"):
        entry.download("https://fake.invalid", tmp_path / "input", readiness_budget=60)
    assert now[0] == 60
    assert pauses[:5] == [0.5, 1, 2, 4, 8]
    assert max(pauses) == 8
    assert timeouts[-1] == pauses[-1]


@pytest.mark.parametrize("status, attempts", [(403, 3), (503, 3), (404, 1)])
def test_later_transfers_keep_short_retry(monkeypatch, status, attempts):
    entry = load("engine_task_entry")
    calls = []
    pauses = []

    def refused(request, timeout):
        calls.append(timeout)
        raise urllib.error.HTTPError(request, status, "denied", {}, io.BytesIO())

    monkeypatch.setattr(entry.time, "sleep", pauses.append)
    monkeypatch.setattr(entry.urllib.request, "urlopen", refused)
    with pytest.raises(urllib.error.HTTPError):
        entry.transfer("https://fake.invalid")
    assert len(calls) == attempts
    assert pauses == list(range(1, attempts))