XFESystem

using System.Runtime.InteropServices;

namespace XFESystem.NativeKernel;

// A read-only GB202 secure-boot observation. It does not boot GSP, create an
// RM session/channel, submit work or imply CUDA is available. Register source:
// NVIDIA 615.78.08 dev_therm.h + dev_therm_addendum.h (MIT), pinned in Tools/Gpu.
public static unsafe class NvidiaBootProbe
{
    public const uint FspBootComplete = 0x00ad00bc;
    static uint _state, _sample, _rejection, _sampled;
    static ulong _physical, _elapsed;
    [DllImport("xfe_native", EntryPoint = "xfe_gpu_discovery_info")] static extern ulong DiscoveryInfo(uint field);
    [DllImport("xfe_native", EntryPoint = "xfe_gpu_discovery_record")]
    static extern int DiscoveryRecord(uint index, NativeGpuDeviceDiscovery* record, uint bytes);

    // State: 0 not run, 1 no RTX5090, 2 unsafe/missing resources, 3 map failed,
    // 4 sampled not complete, 5 FSP boot complete, 6 read failed, 7 unmap failed.
    public static void Initialize(NativeBootInfo* boot)
    {
        if (_state != 0) return;
        _state = 1;
        ulong start = MonotonicClock.Nanoseconds();
        NativeGpuDeviceDiscovery record;
        ulong count = DiscoveryInfo(2);
        if (count > 8) count = 0;
        for (uint i = 0; i < count; i++)
        {
            if (DiscoveryRecord(i, &record, (uint)sizeof(NativeGpuDeviceDiscovery)) != 1 ||
                record.VendorId != 0x10de || record.DeviceId != 0x2b85) continue;
            _state = 2;
            ulong page;
            _rejection = RegisterPageRejection(boot, &record, &page);
            if (_rejection != 0) break;
            _physical = page + (FspBootComplete & 4095);
            ulong flags = Native.IrqSave();
            ulong mapping = VirtualMemory.MapDeviceReadOnly(page, 4096);
            Native.IrqRestore(flags);
            if (mapping == 0) { _state = 3; break; }
            _sample = Native.MmioRead32(mapping + (FspBootComplete & 4095));
            _sampled = 1;
            _state = _sample == 0xff ? 5u : _sample == 0xffffffff ? 6u : 4u;
            flags = Native.IrqSave();
            bool unmapped = VirtualMemory.UnmapDevice(mapping, 4096);
            Native.IrqRestore(flags);
            if (!unmapped) _state = 7;
            break;
        }
        ulong end = MonotonicClock.Nanoseconds();
        _elapsed = end >= start ? end - start : 0;
        Serial.Write("XFE_NVIDIA_BOOT_PROBE|State|"); Serial.Number(_state);
        Serial.Write("|FspScratch|"); Serial.Hex(_sample); Serial.Write("|Physical|"); Serial.Hex(_physical);
        Serial.Write("|ElapsedNs|"); Serial.Number(_elapsed);
        Serial.Write("|Rejection|"); Serial.Number(_rejection);
        Serial.Write("|Sampled|"); Serial.Number(_sampled);
        Serial.Line("|ReadOnly|1|GspStarted|0|ChannelCreated|0|Cuda|0|GpuRasterization|0");
    }

    // Validate the entire page against a matching firmware-assigned aperture,
    // not just its first byte. ReBAR alone cannot establish ordinary BAR0 size.
    public static bool TryRegisterPage(NativeBootInfo* boot, NativeGpuDeviceDiscovery* record, ulong* physical)
        => RegisterPageRejection(boot, record, physical) == 0;

    // Same acceptance rules as TryRegisterPage, with the first failed
    // precondition exposed for physical reports. Zero alone permits mapping.
    public static uint RegisterPageRejection(NativeBootInfo* boot, NativeGpuDeviceDiscovery* record, ulong* physical)
    {
        if (physical == null) return 1;
        *physical = 0;
        if (boot == null || record == null || boot->MemoryEntries == null || boot->MemoryEntryCount == 0 ||
            boot->MemoryEntryCount > 4096) return 2;
        if (record->AbiVersion != 1 || record->RecordBytes != sizeof(NativeGpuDeviceDiscovery)) return 3;
        if (record->VendorId != 0x10de || record->DeviceId != 0x2b85 ||
            record->Segment != 0 || record->Bus > 255 || record->Device > 31 || record->Function > 7 ||
            (record->ClassCode >> 16) != 3 || (record->HeaderType & 127) != 0) return 4;
        if (record->CapabilityStatus != 1) return 5;
        if ((record->Command & 2) == 0) return 6;
        if (record->Bar0.Index != 0 || record->Bar0.Reserved != 0) return 7;
        if (record->Bar0.SizeSource != GpuHardwareBridge.SourceFirmware) return 8;
        if ((record->Bar0.Flags & (GpuHardwareBridge.Memory | GpuHardwareBridge.Io | GpuHardwareBridge.SizeKnown |
                GpuHardwareBridge.Invalid | GpuHardwareBridge.Unassigned | GpuHardwareBridge.UpperHalf)) !=
                (GpuHardwareBridge.Memory | GpuHardwareBridge.SizeKnown)) return 9;
        ulong pageOffset = FspBootComplete & ~4095UL;
        ulong address = record->Bar0.Address, size = record->Bar0.Size;
        if (address == 0 || (address & 4095) != 0 || size < pageOffset + 4096 || (size & (size - 1)) != 0 ||
            (address & (size - 1)) != 0 || address >= (1UL << 52) || size > (1UL << 52) - address ||
            ((record->Bar0.Flags & GpuHardwareBridge.Width64) == 0 && (address >= 0x100000000UL || size > 0x100000000UL - address))) return 10;
        ulong page = address + pageOffset;
        for (ulong i = 0; i < boot->MemoryEntryCount; i++)
        {
            MemoryMapEntry* entry = boot->MemoryEntries[i];
            if (entry == null || entry->Length > ~0UL - entry->Base) return 11;
            // A BAR may be absent from the boot memory map or explicitly
            // reserved. It must never alias RAM, ACPI, kernel or GOP memory.
            if (entry->Length != 0 && address < entry->Base + entry->Length && entry->Base < address + size && entry->Kind != 1)
                return 12;
        }
        *physical = page;
        return 0;
    }

    [UnmanagedCallersOnly(EntryPoint = "xfe_gpu_boot_probe_info")]
    public static ulong Info(uint field)
    {
        if (field == 0) return 1;
        if (field == 1) return _state;
        if (field == 2) return _sample;
        if (field == 3) return _physical;
        if (field == 4) return _elapsed;
        if (field == 5) return _rejection;
        if (field == 6) return _sampled;
        if (field == 7) return 1; // Additive diagnostic fields version.
        // No field can be read as an acceleration/firmware-started claim.
        return 0;
    }
}