XFESystem

#!/usr/bin/env python3
"""One bounded, verified-HTTPS fetch of the exact official RM source archive.

The archive is only a transport. CaptureReferences checks every extracted
file against the pinned Git tree before it enters the source lock.
"""
from __future__ import annotations

import hashlib
import argparse
import http.client
import json
import pathlib
import ssl
import time
import traceback

import CaptureReferences as capture


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--maximum-seconds", type=int, default=180)
    args = parser.parse_args()
    if not 15 <= args.maximum_seconds <= 600:
        raise ValueError("Use a bounded 15..600 second single archive attempt")
    output = capture.TASK_ROOT / "build/Gpu/RmPort" / ("ArchiveFetch-" + str(time.time_ns()))
    output.mkdir(parents=True)
    name = "Archive-" + capture.COMMIT + ".tar.gz"
    target = capture.source_path(name)
    partial = capture.source_path(name + ".Partial-" + output.name)
    url = "https://codeload.github.com/NVIDIA/open-gpu-kernel-modules/tar.gz/" + capture.COMMIT
    receipt = {"Kind": "XfeNvidiaRmArchiveFetch", "Commit": capture.COMMIT, "Url": url,
               "MaximumCacheBytes": capture.MAXIMUM_BYTES, "Passed": False, "GuestCudaPassed": False,
               "ArchiveIsOnlyTransport": True, "BytesReceived": 0, "MaximumSeconds": args.maximum_seconds,
               "SourceSha256": {str(path): capture.sha256(path.read_bytes()) for path in
                                (pathlib.Path(__file__).resolve(), pathlib.Path(capture.__file__).resolve())}}
    connection = None
    started = time.monotonic()
    try:
        capture.SOURCE_ROOT.mkdir(parents=True, exist_ok=True)
        if any(path.is_symlink() for path in capture.SOURCE_ROOT.rglob("*")):
            raise ValueError("Private cache contains a symlink")
        if target.exists() or partial.exists():
            raise ValueError("Refusing to overwrite an existing archive or partial download")
        connection = http.client.HTTPSConnection("codeload.github.com", timeout=15, context=ssl.create_default_context())
        connection.request("GET", "/NVIDIA/open-gpu-kernel-modules/tar.gz/" + capture.COMMIT,
                           headers={"User-Agent": "XFESystem-RmPort-ArchiveCapture/1"})
        digest = hashlib.sha256()
        with connection.getresponse() as response:
            receipt["HttpStatus"] = response.status
            if response.status != 200:
                raise ValueError("Official archive did not return HTTP 200")
            if response.getheader("Content-Type", "").split(";")[0] not in ("application/x-gzip", "application/gzip", "application/octet-stream"):
                raise ValueError("Official archive has an unexpected content type")
            with partial.open("xb", buffering=0) as stream:
                while True:
                    if time.monotonic() - started > args.maximum_seconds:
                        raise TimeoutError("Single archive download exceeded its configured time bound")
                    actual = sum(path.stat().st_size for path in capture.SOURCE_ROOT.rglob("*") if path.is_file())
                    remaining = capture.MAXIMUM_BYTES - actual
                    if remaining < 0:
                        raise ValueError("Private cache already exceeds its bound")
                    block = response.read(min(256 * 1024, remaining + 1))
                    receipt["BytesReceived"] += len(block)
                    if time.monotonic() - started > args.maximum_seconds:
                        raise TimeoutError("Single archive download exceeded its configured time bound")
                    if not block:
                        break
                    if len(block) > remaining:
                        raise ValueError("Archive plus existing cache would exceed 64 MiB; excess bytes were not written")
                    written = stream.write(block)
                    if written != len(block):
                        raise OSError("Short archive write")
                    digest.update(block)
            if receipt["BytesReceived"] == 0:
                raise ValueError("Empty official archive")
        receipt["ArchiveSha256"] = digest.hexdigest()
        partial.replace(target)
        receipt["Archive"] = str(target)
        receipt["Passed"] = True
    except Exception as error:
        receipt["Error"] = type(error).__name__ + ": " + str(error)
        (output / "Failure.log").write_text(traceback.format_exc(), encoding="utf-8")
        if partial.exists():
            receipt["PartialArchive"] = str(partial)
            receipt["PartialBytes"] = partial.stat().st_size
            receipt["PartialSha256"] = capture.sha256(partial.read_bytes())
    finally:
        if connection:
            connection.close()
        receipt["ElapsedSeconds"] = round(time.monotonic() - started, 3)
        receipt["SourceInputsUnchanged"] = all(capture.sha256(pathlib.Path(path).read_bytes()) == expected
                                               for path, expected in receipt["SourceSha256"].items())
        receipt["CacheBytes"] = sum(path.stat().st_size for path in capture.SOURCE_ROOT.rglob("*") if path.is_file())
        receipt["CacheWithinBound"] = receipt["CacheBytes"] <= capture.MAXIMUM_BYTES
        receipt["Passed"] = receipt["Passed"] and receipt["CacheWithinBound"] and receipt["SourceInputsUnchanged"]
        (output / "Result.json").write_text(json.dumps(receipt, indent=2) + "\n", encoding="utf-8", newline="\n")
        print(output / "Result.json", flush=True)
    return 0 if receipt["Passed"] else 1


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