XFESystem

#!/usr/bin/env python3
"""Maintainer-only bounded capture of immutable NVIDIA RM reference inputs.

This does not install/build/load a driver. Preflight.py consumes the resulting
lock without rediscovering or trusting mutable upstream input.
"""
from __future__ import annotations

import argparse
import base64
import hashlib
import http.client
import json
import pathlib
import re
import ssl
import tarfile
import time
import urllib.error
import urllib.parse

COMMIT = "ca2d03e6fafd4d4b56b383497d931c25b283bfcf"
VERSION = "615.78.08"
REPOSITORY = "https://github.com/NVIDIA/open-gpu-kernel-modules"
RAW = "https://raw.githubusercontent.com/NVIDIA/open-gpu-kernel-modules/" + COMMIT + "/"
MAXIMUM_BYTES = 64 * 1024 * 1024
MAXIMUM_FILE_BYTES = 8 * 1024 * 1024
TASK_ROOT = pathlib.Path(__file__).resolve().parents[3]
SOURCE_ROOT = TASK_ROOT / "build/Gpu/RmSources"
INCLUDE_DIRECTORIES = [
    "src/nvidia/kernel/inc", "src/nvidia/interface", "src/common/sdk/nvidia/inc",
    "src/common/sdk/nvidia/inc/hw", "src/nvidia/arch/nvalloc/common/inc",
    "src/nvidia/arch/nvalloc/common/inc/gsp", "src/nvidia/arch/nvalloc/common/inc/deprecated",
    "src/nvidia/arch/nvalloc/unix/include", "src/nvidia/inc", "src/nvidia/inc/os",
    "src/common/shared/inc", "src/common/inc", "src/common/uproc/os/libos-v2.0.0/include",
    "src/common/uproc/os/common/include", "src/common/inc/swref", "src/common/inc/swref/published",
    "src/nvidia/generated", "src/common/nvswitch/kernel/inc", "src/common/nvswitch/interface",
    "src/common/nvswitch/common/inc", "src/common/inc/displayport", "src/common/nvlink/interface",
    "src/common/nvlink/inband/interface", "src/nvidia/src/mm/uvm/interface",
    "src/nvidia/inc/libraries", "src/nvidia/src/libraries", "src/nvidia/inc/kernel",
    "src/common/mbedtls/3.6.2/include", "src/common/mbedtls/3.6.2/nvidia",
]
BUILD_INPUTS = [
    "COPYING", "README.md", "version.mk", "Makefile", "utils.mk", "nv-compiler.sh",
    "src/nvidia/Makefile", "src/nvidia/srcs.mk", "src/nvidia/exports_link_command.txt",
    "src/nvidia/nv-kernel.ld", "src/nvidia/nv-kernel.patchable.ld",
    "src/nvidia/src/libraries/libspdm/nvidia/openspdm.mk",
    "kernel-open/common/inc/os-interface.h", "kernel-open/nvidia/os-interface.c",
    "src/nvidia/generated/g_kern_fsp_nvoc.h", "src/nvidia/generated/g_kern_fsp_nvoc.c",
    "src/nvidia/generated/g_kernel_gsp_nvoc.h", "src/nvidia/generated/g_kernel_gsp_nvoc.c",
    "src/common/inc/swref/published/blackwell/gb202/dev_therm.h",
    "src/common/inc/swref/published/blackwell/gb202/dev_therm_addendum.h",
]
COMPONENTS = [
    "src/nvidia/src/libraries/msgq/msgq.c",
    "src/nvidia/src/kernel/gpu/fsp/arch/blackwell/kern_fsp_gb202.c",
    "src/nvidia/src/kernel/gpu/gsp/arch/blackwell/kernel_gsp_gb202.c",
    "src/nvidia/src/kernel/gpu/gsp/message_queue_cpu.c",
]
_CONNECTIONS: dict[str, http.client.HTTPSConnection] = {}


def sha256(data: bytes) -> str:
    return hashlib.sha256(data).hexdigest()


def git_blob(data: bytes) -> str:
    return hashlib.sha1(b"blob " + str(len(data)).encode("ascii") + b"\0" + data).hexdigest()


def source_path(relative: str) -> pathlib.Path:
    parsed = pathlib.PurePosixPath(relative)
    if parsed.is_absolute() or ".." in parsed.parts or "\\" in relative or ":" in relative:
        raise ValueError("Unsafe source path")
    path = SOURCE_ROOT.joinpath(*parsed.parts)
    if not path.resolve().is_relative_to(SOURCE_ROOT.resolve()) or path.is_symlink():
        raise ValueError("Source escapes private build directory")
    return path


def download(url: str, maximum: int, attempts: int = 3) -> bytes:
    parsed = urllib.parse.urlsplit(url)
    if parsed.scheme != "https" or parsed.hostname not in ("raw.githubusercontent.com", "api.github.com") or parsed.port not in (None, 443) or parsed.username or parsed.password or parsed.fragment:
        raise ValueError("Only the pinned official HTTPS origins are allowed")
    host = parsed.hostname
    request_path = parsed.path + ("?" + parsed.query if parsed.query else "")
    for attempt in range(attempts):
        try:
            if host not in _CONNECTIONS:
                # Certificate and hostname verification remain enabled. Reuse
                # TLS connections rather than request hundreds of handshakes.
                _CONNECTIONS[host] = http.client.HTTPSConnection(host, timeout=10, context=ssl.create_default_context())
            connection = _CONNECTIONS[host]
            connection.request("GET", request_path, headers={"User-Agent": "XFESystem-RmPort-ReferenceCapture/1"})
            with connection.getresponse() as response:
                if response.status != 200:
                    raise urllib.error.HTTPError(url, response.status, response.reason, response.headers, None)
                length = response.getheader("Content-Length")
                if length is not None and int(length) > maximum:
                    raise ValueError("Reference exceeds byte bound")
                data = response.read(maximum + 1)
            if len(data) > maximum:
                raise ValueError("Reference exceeds byte bound")
            return data
        except (OSError, http.client.HTTPException, urllib.error.URLError) as error:
            connection = _CONNECTIONS.pop(host, None)
            if connection:
                connection.close()
            if isinstance(error, urllib.error.HTTPError) and error.code not in (408, 429) and error.code < 500:
                raise
            if attempt == attempts - 1:
                raise urllib.error.URLError(str(error)) from error
    raise AssertionError("Unreachable")


def load_tree() -> dict[str, dict]:
    tree_file = source_path("Tree.json")
    if not tree_file.exists():
        remaining = MAXIMUM_BYTES - sum(path.stat().st_size for path in SOURCE_ROOT.rglob("*") if path.is_file())
        if remaining <= 0:
            raise ValueError("No bounded cache space remains for pinned tree metadata")
        data = download("https://api.github.com/repos/NVIDIA/open-gpu-kernel-modules/git/trees/" + COMMIT + "?recursive=1", min(12 * 1024 * 1024, remaining))
        tree_file.parent.mkdir(parents=True, exist_ok=True)
        tree_file.write_bytes(data)
    tree = json.loads(tree_file.read_bytes())
    if tree.get("truncated") or tree.get("sha") != COMMIT:
        raise ValueError("Incomplete or different pinned release tree")
    return {entry["path"]: entry for entry in tree["tree"] if entry["type"] == "blob"}


def read_pinned_source(path: str, entry: dict) -> bytes:
    try:
        return download(RAW + path, MAXIMUM_FILE_BYTES, attempts=3)
    except (TimeoutError, urllib.error.URLError) as raw_error:
        # Same immutable Git object through NVIDIA's official GitHub API;
        # transient raw-host failures must never change source identity.
        try:
            payload = json.loads(download("https://api.github.com/repos/NVIDIA/open-gpu-kernel-modules/git/blobs/" + entry["sha"], 12 * 1024 * 1024))
        except (TimeoutError, urllib.error.URLError) as api_error:
            raise urllib.error.URLError("Raw HTTPS failed: " + str(raw_error) + "; immutable Git blob API failed: " + str(api_error)) from api_error
        if payload.get("sha") != entry["sha"] or payload.get("encoding") != "base64" or payload.get("size") != entry["size"]:
            raise ValueError("GitHub blob metadata differs: " + path)
        return base64.b64decode("".join(payload["content"].split()), validate=True)


def capture_file(path: str, tree: dict[str, dict], archive: tarfile.TarFile | None = None,
                 archive_members: dict[str, tarfile.TarInfo] | None = None) -> tuple[str, bytes]:
    entry = tree[path]
    if entry.get("mode") not in ("100644", "100755"):
        raise ValueError("Unsupported upstream file mode: " + path)
    if entry["size"] > MAXIMUM_FILE_BYTES:
        raise ValueError("Per-file bound: " + path)
    target = source_path(path)
    if target.exists():
        data = target.read_bytes()
    elif archive is not None:
        member = (archive_members or {}).get(path)
        if member is None or not member.isfile() or member.size != entry["size"]:
            raise ValueError("Pinned file missing/different in official archive: " + path)
        stream = archive.extractfile(member)
        if stream is None:
            raise ValueError("Cannot read official archive member: " + path)
        with stream:
            data = stream.read(MAXIMUM_FILE_BYTES + 1)
    else:
        data = read_pinned_source(path, entry)
    if len(data) != entry["size"] or git_blob(data) != entry["sha"]:
        raise ValueError("Pinned Git blob differs: " + path)
    if not target.exists():
        target.parent.mkdir(parents=True, exist_ok=True)
        temporary = target.with_name(target.name + ".Downloading")
        temporary.write_bytes(data)
        temporary.replace(target)
    return path, data


def resolve_include(owner: str, include: str, quoted: bool, tree: dict[str, dict]) -> str | None:
    candidates = []
    if quoted:
        candidates.append(str(pathlib.PurePosixPath(owner).parent / include))
    candidates += [directory + "/" + include for directory in INCLUDE_DIRECTORIES]
    for candidate in candidates:
        # Normalize official relative includes without permitting cache escape.
        parts: list[str] = []
        for part in pathlib.PurePosixPath(candidate).parts:
            if part == "..":
                if not parts:
                    break
                parts.pop()
            elif part != ".":
                parts.append(part)
        normalized = "/".join(parts)
        if normalized in tree:
            return normalized
    return None


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--closure", choices=["msgq", "all"], default="msgq")
    parser.add_argument("--archive", action="store_true", help="Use only the fixed official archive in the private cache, verifying each Git blob")
    parser.add_argument("--lock", default=str(pathlib.Path(__file__).with_name("RmSources.lock.json")))
    args = parser.parse_args()
    output = pathlib.Path(args.lock).resolve()
    if not output.is_relative_to(pathlib.Path(__file__).parent.resolve()):
        raise ValueError("Lock output must remain in the owned RmPort directory")
    SOURCE_ROOT.mkdir(parents=True, exist_ok=True)
    if any(path.is_symlink() for path in SOURCE_ROOT.rglob("*")):
        raise ValueError("Private cache contains a symlink")
    if sum(path.stat().st_size for path in SOURCE_ROOT.rglob("*") if path.is_file()) > MAXIMUM_BYTES:
        raise ValueError("Private reference cache already exceeds 64 MiB")
    tree = load_tree()
    archive = None
    archive_members = {}
    archive_input = None
    if args.archive:
        archive_name = "Archive-" + COMMIT + ".tar.gz"
        archive_path = source_path(archive_name)
        archive_input = {"Path": archive_name, "Bytes": archive_path.stat().st_size, "Sha256": sha256(archive_path.read_bytes()),
                         "Url": "https://codeload.github.com/NVIDIA/open-gpu-kernel-modules/tar.gz/" + COMMIT,
                         "OnlyTransport": True, "ExtractedFilesRequirePinnedGitBlob": True}
        archive = tarfile.open(archive_path, "r:gz")
        prefix = "open-gpu-kernel-modules-" + COMMIT
        member_count = 0
        total_declared = 0
        for member in archive:
            member_count += 1
            total_declared += member.size
            parts = pathlib.PurePosixPath(member.name).parts
            if member_count > 20000 or total_declared > 512 * 1024 * 1024 or len(member.name) > 2048:
                raise ValueError("Official archive index exceeds its bounded source shape")
            if not parts or parts[0] != prefix or ".." in parts or "\\" in member.name or ":" in member.name:
                raise ValueError("Official archive contains an unexpected member path")
            relative = "/".join(parts[1:])
            if relative in tree and member.isfile():
                if relative in archive_members:
                    raise ValueError("Duplicate official archive member: " + relative)
                archive_members[relative] = member
        archive_input["MemberCount"] = member_count
        archive_input["DeclaredUncompressedBytes"] = total_declared
    selected = {path for path in BUILD_INPUTS + COMPONENTS if path in tree}
    missing_seeds = sorted(set(BUILD_INPUTS + COMPONENTS) - selected)
    include_roots = COMPONENTS if args.closure == "all" else COMPONENTS[:1]
    pending = set(selected) | {"src/common/sdk/nvidia/inc/cpuopsys.h"}
    recursive = set(include_roots) | {"src/common/sdk/nvidia/inc/cpuopsys.h"}
    captured: dict[str, bytes] = {}
    scanned: set[str] = set()
    unresolved: list[dict[str, str]] = []
    cached_total = sum(path.stat().st_size for path in SOURCE_ROOT.rglob("*") if path.is_file())
    failure = None
    try:
        while pending or recursive - scanned:
            batch = sorted(pending - captured.keys())
            for path in batch:
                expected_new = tree[path]["size"] if not source_path(path).exists() else 0
                if cached_total + expected_new > MAXIMUM_BYTES:
                    raise ValueError("Capture would exceed private cache 64 MiB bound: " + path)
                path, data = capture_file(path, tree, archive, archive_members)
                captured[path] = data
                pending.discard(path)
                cached_total += expected_new
                print("Captured", path, len(data), flush=True)
            pending.difference_update(captured)
            for path in sorted(recursive & captured.keys() - scanned):
                scanned.add(path)
                text = captured[path].decode("utf-8", errors="replace")
                for match in re.finditer(r'^\s*#\s*include\s*([<"])([^>"\r\n]+)[>"]', text, re.MULTILINE):
                    name = match.group(2)
                    included = resolve_include(path, name, match.group(1) == '"', tree)
                    if included:
                        if included not in captured:
                            pending.add(included)
                        recursive.add(included)
                    else:
                        unresolved.append({"Owner": path, "Include": name})
    except (TimeoutError, urllib.error.URLError, ValueError) as error:
        # Keep a verifiable partial lock for resumable network/capacity failures;
        # a corrupt cached file still fails again below and cannot enter a lock.
        failure = {"Type": type(error).__name__, "Message": str(error),
                   "PendingPaths": sorted(pending), "UnscannedPaths": sorted(recursive - scanned)}
        print("Capture incomplete:", failure["Type"], failure["Message"], flush=True)
    if archive is not None:
        archive.close()
        if sha256(source_path(archive_input["Path"]).read_bytes()) != archive_input["Sha256"]:
            raise ValueError("Official archive changed during dependency capture")
    # Earlier official probes may have cached additional headers. Clang may
    # encounter these even while recovering from a missing-include error, so
    # bind their actual pinned bytes too; never silently trust cache extras.
    extra_cached = []
    for file in sorted(SOURCE_ROOT.rglob("*")):
        if not file.is_file():
            continue
        relative = str(file.relative_to(SOURCE_ROOT)).replace("\\", "/")
        if relative in tree and relative not in captured:
            path, data = capture_file(relative, tree)
            captured[path] = data
            extra_cached.append(path)
    files = []
    for path, data in sorted(captured.items()):
        spdx = re.search(rb"SPDX-License-Identifier:\s*([^\r\n*]+)", data[:8192])
        row = {"Path": path, "Bytes": len(data), "GitBlob": tree[path]["sha"], "Sha256": sha256(data), "Mode": tree[path]["mode"],
               "License": spdx.group(1).decode("ascii").strip() if spdx else "Per-file notice and COPYING"}
        files.append(row)
    source_list = captured.get("src/nvidia/srcs.mk", b"").decode("utf-8")
    listed = re.findall(r"^SRCS(?:_CXX)?\s*\+=\s*(\S+)", source_list, re.MULTILINE)
    generated = [path for path in listed if path.startswith("generated/")]
    lock = {"Schema": "XfeNvidiaRmReferenceSources1", "Repository": REPOSITORY, "Commit": COMMIT, "Version": VERSION,
            "MaximumCacheBytes": MAXIMUM_BYTES, "IncludeDirectories": INCLUDE_DIRECTORIES, "Components": COMPONENTS,
            "Closure": args.closure, "Files": files, "ExtraVerifiedCachedFiles": extra_cached,
            "CaptureComplete": failure is None, "CaptureFailure": failure,
            "ArchiveInput": archive_input,
            "UnresolvedLiteralIncludes": unresolved, "MissingSeedPaths": missing_seeds,
            "BuildInventory": {"ListedTranslationUnits": len(listed), "ListedGeneratedTranslationUnits": len(generated),
                "ListedNvocTranslationUnits": sum("_nvoc." in path for path in generated),
                "ListedSourceBytes": sum(tree.get("src/nvidia/" + path, {}).get("size", 0) for path in listed),
                "NvocGeneratorInvokedByPublishedMakefile": False},
            "GuestCudaPassed": False, "GpuExecutionPassed": False}
    temporary = output.with_name(output.name + ".Writing")
    temporary.write_text(json.dumps(lock, indent=2) + "\n", encoding="utf-8", newline="\n")
    temporary.replace(output)
    cached_total = sum(path.stat().st_size for path in SOURCE_ROOT.rglob("*") if path.is_file())
    print("Lock", output, "Files", len(files), "CacheBytes", cached_total, flush=True)
    if failure:
        report = TASK_ROOT / "build/Gpu/RmPort" / ("CaptureFailure-" + str(time.time_ns()) + ".json")
        report.parent.mkdir(parents=True, exist_ok=True)
        report.write_text(json.dumps({"Kind": "XfeNvidiaRmCaptureFailure", "Commit": COMMIT,
            "LockSha256": sha256(output.read_bytes()), "Failure": failure,
            "CacheBytes": cached_total, "GuestCudaPassed": False}, indent=2) + "\n", encoding="utf-8")
        print(report, flush=True)
    return 0 if failure is None else 1


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