XFESystem

"""Verify generated command bytes against hash-pinned NVIDIA header fields.

This is an offline protocol oracle. It neither opens a GPU nor executes CUDA.
"""
from __future__ import annotations

import argparse
import base64
import hashlib
import json
from pathlib import Path
import re
import struct
import urllib.request


def checked_sources(root: Path, download: bool) -> tuple[dict, dict[str, str]]:
    lock = json.loads((root / "Tools/Gpu/NvidiaSources.lock.json").read_text(encoding="utf-8"))
    sources = {}
    for item in lock["Files"]:
        file = root / lock["CacheDirectory"] / item["CacheName"]
        if not file.exists():
            if not download:
                raise ValueError(f"Missing pinned reference {file}; use --download to retrieve official sources.")
            request = urllib.request.Request(
                f'https://api.github.com/repos/NVIDIA/open-gpu-kernel-modules/git/blobs/{item["GitBlob"]}',
                headers={"User-Agent": "XFESystem-NvidiaProtocol"},
            )
            with urllib.request.urlopen(request, timeout=25) as response:
                body = json.load(response)
            data = base64.b64decode(body["content"], validate=False)
            file.parent.mkdir(parents=True, exist_ok=True)
            file.write_bytes(data)
        data = file.read_bytes()
        if hashlib.sha256(data).hexdigest() != item["Sha256"]:
            raise ValueError(f"SHA256 mismatch: {file}")
        blob = b"blob " + str(len(data)).encode("ascii") + b"\0" + data
        if hashlib.sha1(blob).hexdigest() != item["GitBlob"]:
            raise ValueError(f"Git blob mismatch: {file}")
        sources[item["CacheName"].lower()] = data.decode("utf-8")
    return lock, sources


class HeaderFields:
    def __init__(self, sources: dict[str, str]):
        self.defines = {}
        for name, source in sources.items():
            if not name.endswith(".h"):
                continue
            for match in re.finditer(r"^#define[ \t]+(\w+)[ \t]+([^\r\n]+)", source, re.MULTILINE):
                self.defines[match[1]] = match[2].split("//", 1)[0].strip()

    def value(self, name: str) -> int:
        return int(self.defines[name].strip("()"), 0)

    def pack(self, name: str, value: int) -> int:
        high, low = map(int, self.defines[name].split(":"))
        if value < 0 or value >= 1 << (high - low + 1):
            raise ValueError(f"Value does not fit {name}")
        return value << low

    def constant(self, name: str, suffix: str) -> int:
        return self.pack(name, self.value(name + "_" + suffix))


def verify(root: Path, directory: Path, download: bool) -> dict:
    lock, sources = checked_sources(root, download)
    fields = HeaderFields(sources)
    metadata = json.loads((directory / "Vector.json").read_text(encoding="utf-8"))
    subchannel = fields.value("NVA06F_SUBCHANNEL_COPY_ENGINE")

    def header(method: int, count: int) -> int:
        return (fields.constant("NVC46F_DMA_SEC_OP", "INC_METHOD") |
                fields.pack("NVC46F_DMA_INCR_ADDRESS", method >> 2) |
                fields.pack("NVC46F_DMA_INCR_SUBCHANNEL", subchannel) |
                fields.pack("NVC46F_DMA_INCR_COUNT", count))

    words = []

    def method(name: str, *values: int) -> None:
        words.append(header(fields.value(name), len(values)))
        words.extend(values)

    target, fence, pushbuffer = metadata["Target"], metadata["Fence"], metadata["Pushbuffer"]
    method("NVCA6F_SET_OBJECT", fields.value("BLACKWELL_DMA_COPY_B"))
    method("NVC8B5_SET_REMAP_CONST_B", metadata["Color"],
           fields.constant("NVC8B5_SET_REMAP_COMPONENTS_DST_X", "CONST_B") |
           fields.constant("NVC8B5_SET_REMAP_COMPONENTS_COMPONENT_SIZE", "FOUR") |
           fields.constant("NVC8B5_SET_REMAP_COMPONENTS_NUM_DST_COMPONENTS", "ONE"))
    method("NVC8B5_OFFSET_OUT_UPPER", fields.pack("NVC8B5_OFFSET_OUT_UPPER_UPPER", target >> 32), target & 0xFFFFFFFF)
    method("NVC8B5_LINE_LENGTH_IN", metadata["Pixels"])
    method("NVC8B5_LAUNCH_DMA",
           fields.constant("NVC8B5_LAUNCH_DMA_DATA_TRANSFER_TYPE", "NON_PIPELINED") |
           fields.constant("NVC8B5_LAUNCH_DMA_FLUSH_ENABLE", "TRUE") |
           fields.constant("NVC8B5_LAUNCH_DMA_FLUSH_TYPE", "SYS") |
           fields.constant("NVC8B5_LAUNCH_DMA_SRC_MEMORY_LAYOUT", "PITCH") |
           fields.constant("NVC8B5_LAUNCH_DMA_DST_MEMORY_LAYOUT", "PITCH") |
           fields.constant("NVC8B5_LAUNCH_DMA_MULTI_LINE_ENABLE", "FALSE") |
           fields.constant("NVC8B5_LAUNCH_DMA_REMAP_ENABLE", "TRUE") |
           fields.constant("NVC8B5_LAUNCH_DMA_DST_TYPE", "VIRTUAL") |
           fields.constant("NVC8B5_LAUNCH_DMA_COPY_TYPE", "DEFAULT") |
           fields.constant("NVC8B5_LAUNCH_DMA_MEMORY_SCRUB_ENABLE", "FALSE") |
           fields.constant("NVC8B5_LAUNCH_DMA_DISABLE_PLC", "TRUE"))
    method("NVC8B5_SET_SEMAPHORE_A", fields.pack("NVC8B5_SET_SEMAPHORE_A_UPPER", fence >> 32),
           fence & 0xFFFFFFFF, metadata["Payload"])
    method("NVC8B5_LAUNCH_DMA",
           fields.constant("NVC8B5_LAUNCH_DMA_DATA_TRANSFER_TYPE", "NONE") |
           fields.constant("NVC8B5_LAUNCH_DMA_FLUSH_ENABLE", "TRUE") |
           fields.constant("NVC8B5_LAUNCH_DMA_FLUSH_TYPE", "SYS") |
           fields.constant("NVC8B5_LAUNCH_DMA_SEMAPHORE_TYPE", "RELEASE_SEMAPHORE_NO_TIMESTAMP") |
           fields.constant("NVC8B5_LAUNCH_DMA_SEMAPHORE_PAYLOAD_SIZE", "ONE_WORD") |
           fields.constant("NVC8B5_LAUNCH_DMA_DISABLE_PLC", "TRUE"))
    commands = struct.pack("<" + "I" * len(words), *words)
    entries = struct.pack("<IIII",
                          fields.pack("NVCA6F_GP_ENTRY0_PB_EXTENDED_BASE_OPERAND", pushbuffer >> 40),
                          fields.constant("NVCA6F_GP_ENTRY1_OPCODE", "SET_PB_SEGMENT_EXTENDED_BASE"),
                          fields.pack("NVCA6F_GP_ENTRY0_GET", (pushbuffer & 0xFFFFFFFF) >> 2),
                          fields.pack("NVCA6F_GP_ENTRY1_GET_HI", (pushbuffer >> 32) & 0xFF) |
                          fields.pack("NVCA6F_GP_ENTRY1_LENGTH", len(commands) >> 2) |
                          fields.constant("NVCA6F_GP_ENTRY1_SYNC", "WAIT"))
    for name, expected in (("Fill32.bin", commands), ("Gpfifo.bin", entries)):
        actual = (directory / name).read_bytes()
        if actual != expected:
            raise ValueError(f"{name} differs from NVIDIA-header field oracle: {actual.hex()} != {expected.hex()}")
    result = {
        "Schema": "XfeNvidiaHeaderOracle1", "Passed": True,
        "Scope": "Offline C# binary encoding versus independent Python NVIDIA-header oracle; no GPU/CUDA execution",
        "NvidiaCommit": lock["Commit"], "License": lock["License"],
        "SourceFilesVerified": len(lock["Files"]), "CommandWords": len(words), "GpfifoEntries": len(entries) // 8,
        "Fill32Sha256": hashlib.sha256(commands).hexdigest(), "GpfifoSha256": hashlib.sha256(entries).hexdigest(),
        "CSharpSourceHashes": {
            str(file.relative_to(root)).replace("\\", "/"): hashlib.sha256(file.read_bytes()).hexdigest()
            for file in sorted((root / "Source/XFESystem.Gpu.Nvidia").glob("*.cs"))
        },
    }
    (directory / "HeaderOracle.json").write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8")
    return result


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--repo-root", type=Path, default=Path(__file__).resolve().parents[2])
    parser.add_argument("--probe-directory", type=Path, required=True)
    parser.add_argument("--download", action="store_true")
    options = parser.parse_args()
    print(json.dumps(verify(options.repo_root.resolve(), options.probe_directory.resolve(), options.download), indent=2))