/**
 * NetworkStackSetupScreen.c
 *
 * Full decompilation of the NetworkStackSetupScreen DXE driver from Lenovo
 * HR650X BIOS.  This driver manages the "Network Stack" setup screen
 * configuration variable. On first boot, it creates the "NetworkStackVar"
 * UEFI variable with default settings. At each subsequent boot it reads the
 * existing variable and checks CMOS NV storage to determine the enabled /
 * disabled status.
 *
 * Binary:  AmiBoardInfo2.efi (PE32+ X64, ~0x18A0 bytes)
 * Source (PDB): e:\hs\AmiModulePkg\BoardInfo\AmiBoardInfo2.c
 *              https://github.com/ami/AmiModulePkg/BoardInfo/AmiBoardInfo2.c
 *
 * Reconstructed by: hr650x-bios-re
 * Tools: IDA Pro 9.0 + Hex-Rays
 */

#include "AmiBoardInfo2.h"

//
// ---------------------------------------------------------------------------
// Global variable storage (mapped from .data section)
// ---------------------------------------------------------------------------
//

EFI_HANDLE           gImageHandle          = NULL;
EFI_SYSTEM_TABLE    *gSystemTable          = NULL;
EFI_BOOT_SERVICES   *gBootServices         = NULL;
EFI_RUNTIME_SERVICES *gRuntimeServices     = NULL;

//
EFI_AMI_BOARD_INFO_PROTOCOL *gBoardInfoProtocol = NULL;

//
EFI_PHYSICAL_ADDRESS gHobListAddress = 0;

//
UINT64 gBootServicesCopy = 0;
UINT64 gRuntimeServicesCopy = 0;
UINT64 gSystemTableCopy = 0;

//
NETWORK_STACK_VAR gNetworkStackVarData;

//
UINT8  gNetworkStackBootFlag = 0;
UINT16 gNetworkStackReservedWord = 0;

//
// ---------------------------------------------------------------------------
// GUID definitions
// ---------------------------------------------------------------------------
//

//
// GUID for the AMI Board Info Protocol.
// The protocol interface is located via gBS->LocateProtocol.
// Stored at .data 0xCC0.
//
EFI_GUID gBoardInfoProtocolGuid = {
    0xCC0, 0xCC1, 0xCC2,
    { 0xCC3, 0xCC4, 0xCC5, 0xCC6, 0xCC7, 0xCC8, 0xCC9, 0xCCA }
};

//
// PiHobList GUID (standard UEFI PiHobList GUID).
// Used to locate the HOB list from the system table configuration table.
// Stored as two halves at .data 0xCD0 (first 8 bytes) and 0xCD8 (second 8).
//
EFI_GUID gPiHobListGuid = {
    0xCD0, 0xCD1, 0xCD2,
    { 0xCD3, 0xCD4, 0xCD5, 0xCD6, 0xCD7, 0xCD8, 0xCD9, 0xCDA }
};

//
// Vendor GUID for the NetworkStackVar variable.
// Stored at .data 0xCE0.
//
EFI_GUID gNetworkStackVarGuid = {
    0xCE0, 0xCE1, 0xCE2,
    { 0xCE3, 0xCE4, 0xCE5, 0xCE6, 0xCE7, 0xCE8, 0xCE9, 0xCEA }
};


//
// ---------------------------------------------------------------------------
// Forward declarations
// ---------------------------------------------------------------------------
//

UINT64
EFIAPI
ReadUnalignedU64 (
    IN VOID *Buffer
    );

BOOLEAN
EFIAPI
CompareGuidU64 (
    IN VOID *Buffer1,
    IN VOID *Buffer2
    );

UINTN
EFIAPI
ReadNetworkStackNvStorage (
    VOID
    );

EFI_AMI_BOARD_INFO_PROTOCOL *
EFIAPI
LocateBoardInfoProtocol (
    VOID
    );

VOID
EFIAPI
AmiBoardInfoAssertCallback (
    IN CHAR8 *FileName,
    IN UINTN LineNumber,
    IN CHAR8 *Expression
    );

EFI_STATUS
EFIAPI
SendBoardInfoIfNetworkStackEnabled (
    IN UINTN Mask,
    IN UINTN Arg2,
    ...
    );

EFI_PHYSICAL_ADDRESS
EFIAPI
GetHobListAddress (
    IN EFI_SYSTEM_TABLE *SystemTable
    );

EFI_STATUS
EFIAPI
InitializeNetworkStackVar (
    IN UINTN             N10Param,
    IN EFI_SYSTEM_TABLE *SystemTable
    );

EFI_STATUS
EFIAPI
NetworkStackSetupScreenEntryPoint (
    IN EFI_HANDLE        ImageHandle,
    IN EFI_SYSTEM_TABLE *SystemTable
    );


//
// ---------------------------------------------------------------------------
// ReadUnalignedU64
// ---------------------------------------------------------------------------
//

/**
  Read an unaligned 64-bit value from a buffer.

  Dereferences the pointer as UINT64*.

  @param[in] Buffer  Pointer to read from (must not be NULL).

  @return UINT64 value at the address.
**/
UINT64
EFIAPI
ReadUnalignedU64 (
    IN VOID *Buffer
    )
{
    ASSERT (Buffer != NULL);
    return *(UINT64 *)Buffer;
}


//
// ---------------------------------------------------------------------------
// CompareGuidU64
// ---------------------------------------------------------------------------
//

/**
  Compare two 8-byte (QWORD) values by unaligned read.

  Reads 8 bytes from each pointer and returns TRUE if equal.

  @param[in] Buffer1  First buffer.
  @param[in] Buffer2  Second buffer.

  @return TRUE if first 8 bytes are equal, FALSE otherwise.
**/
BOOLEAN
EFIAPI
CompareGuidU64 (
    IN VOID *Buffer1,
    IN VOID *Buffer2
    )
{
    return (BOOLEAN)(
        ReadUnalignedU64 (Buffer1) == ReadUnalignedU64 (Buffer2)
        );
}


//
// ---------------------------------------------------------------------------
// ReadNetworkStackNvStorage
// ---------------------------------------------------------------------------
//

/**
  Read the CMOS NV storage diagnostic status register (0x4B) to determine
  whether the network stack has been enabled or disabled by the platform
  setup utility.

  CMOS register 0x4B layout (diagnostic status byte):
    Bit 7  : RTC power-good flag (1 = valid)
    Bits 6:2 : other diagnostic flags (unused)
    Bit 1  : when Bit 0 is 0, set = network stack disabled (value 2)
    Bit 0  : set = network stack enabled (value 1)

  Fallback: if the CMOS register reads 0 but appears valid (>3), it
  reads a fixed byte at physical address 0xFDAF0490 as override.

  @return UINTN
    EFI_NETWORK_STACK_ENABLED  (0x80000004) if NvStorage == 1,
    EFI_NETWORK_STACK_DISABLED (0x80000006) if NvStorage == 2,
    0 if NvStorage is invalid/unavailable.
**/
UINTN
EFIAPI
ReadNetworkStackNvStorage (
    VOID
    )
{
    UINT8 DiagnosticValue;
    UINT8 NvStorage;

    //
    // Save current index, select CMOS diagnostic register (0x4B).
    // Mask bit 7 (power good flag) to access the register number.
    //
    IoWrite8 (
        CMOS_REG_INDEX_PORT,
        (IoRead8 (CMOS_REG_INDEX_PORT) & CMOS_DIAG_STATUS_MASK) |
        CMOS_DIAGNOSTIC_STATUS_REG
        );

    //
    // Read the diagnostic status value.
    //
    DiagnosticValue = IoRead8 (CMOS_REG_DATA_PORT);
    NvStorage       = DiagnosticValue;

    if (DiagnosticValue > 3) {
        //
        // Value appears accessible.
        //
        if (NvStorage == 0) {
            //
            // Fallback: read a fixed byte in SPI/flash memory (0xFDAF0490)
            // and use bit 1 combined with bit 0 as override indicator.
            //
            NvStorage = (*(volatile UINT8 *)0xFDAF0490 & 2) | 1;
        }
    }

    //
    // Validate: only 1 and 2 are acceptable.
    // (NvStorage - 1) > 0xFD catches value == 1 (0), 2 (1), or invalid.
    //
    if ((UINT8)(NvStorage - 1) > 0xFD) {
        //
        // NvStorage is not 1 or 2 (or is 0xFF) -> undetermined.
        //
        return 0;
    }

    if (NvStorage == 1) {
        return EFI_NETWORK_STACK_ENABLED;     // 0x80000004
    }

    return EFI_NETWORK_STACK_DISABLED;        // 0x80000006
}


//
// ---------------------------------------------------------------------------
// LocateBoardInfoProtocol
// ---------------------------------------------------------------------------
//

/**
  Locate the AMI Board Info Protocol via gBS->LocateProtocol.

  This protocol is installed by the board/platform drivers.
  It provides:
    - Offset 0x00: Debug assertion callback (VOID*)(CHAR8*, UINTN, CHAR8*)
    - Offset 0x08: Status/event send callback (EFI_STATUS*)(...)

  @return Pointer to EFI_AMI_BOARD_INFO_PROTOCOL, or NULL if
          not available or pool allocation fails.
**/
EFI_AMI_BOARD_INFO_PROTOCOL *
EFIAPI
LocateBoardInfoProtocol (
    VOID
    )
{
    EFI_STATUS Status;
    VOID       *Pool;

    if (gBoardInfoProtocol != NULL) {
        return gBoardInfoProtocol;
    }

    //
    // Allocate a small pool (EfiBootServicesData) as scratch.
    //
    Pool = AllocatePool (31);    // 31 = EfiBootServicesData
    ZeroMem (Pool, 31);

    if ((UINTN)Pool > 0x10) {
        //
        // Unlikely: pool allocation was larger than 16 bytes, fail.
        //
        FreePool (Pool);
        return NULL;
    }

    //
    // Locate the protocol.
    //
    Status = gBootServices->LocateProtocol (
                     &gBoardInfoProtocolGuid,
                     NULL,                          // Registration
                     (VOID **)&gBoardInfoProtocol
                     );
    if (EFI_ERROR (Status)) {
        gBoardInfoProtocol = NULL;
    }

    return gBoardInfoProtocol;
}


//
// ---------------------------------------------------------------------------
// AmiBoardInfoAssertCallback
// ---------------------------------------------------------------------------
//

/**
  Debug assertion output via the AMI Board Info Protocol.

  @param[in] FileName     Source file name string.
  @param[in] LineNumber   Line number of the assertion.
  @param[in] Expression   Assertion expression string.
**/
VOID
EFIAPI
AmiBoardInfoAssertCallback (
    IN CHAR8 *FileName,
    IN UINTN LineNumber,
    IN CHAR8 *Expression
    )
{
    EFI_AMI_BOARD_INFO_PROTOCOL *Protocol;

    Protocol = LocateBoardInfoProtocol ();
    if (Protocol != NULL) {
        //
        // Offset 0x00: AssertCallback.
        // Signature: VOID (*)(CHAR8 *File, UINTN Line, CHAR8 *Expr)
        //
        Protocol->AssertCallback (FileName, LineNumber, Expression);
    }
}


//
// ---------------------------------------------------------------------------
// SendBoardInfoIfNetworkStackEnabled
// ---------------------------------------------------------------------------
//

/**
  Conditional status send.

  Locates the AMI Board Info Protocol, reads the CMOS NV storage
  status via ReadNetworkStackNvStorage(), and if (NvResult & Mask) is
  non-zero, invokes the protocol's StatusSend callback with
  the given arguments.

  @param[in] Mask     Status mask to compare against ReadNetworkStackNvStorage.
  @param[in] Arg2     Second argument to StatusSend callback.
  @param[in] ...      Third argument (varargs, passed as VOID* pointer).

  @return EFI_STATUS  Result from StatusSend, or EFI_UNSUPPORTED
                      if protocol/NV check fails.
**/
EFI_STATUS
EFIAPI
SendBoardInfoIfNetworkStackEnabled (
    IN UINTN Mask,
    IN UINTN Arg2,
    ...
    )
{
    EFI_AMI_BOARD_INFO_PROTOCOL *Protocol;
    EFI_STATUS                   Status;
    UINTN                        NvResult;
    VA_LIST                      Args;
    VOID                         *VaPtr;

    Protocol = LocateBoardInfoProtocol ();
    if (Protocol == NULL) {
        return EFI_UNSUPPORTED;
    }

    NvResult = ReadNetworkStackNvStorage ();
    if ((NvResult & Mask) == 0) {
        //
        // CMOS NV storage status does not match Mask.
        //
        return EFI_UNSUPPORTED;
    }

    //
    // Status matches; invoke StatusSend callback at offset 0x08.
    //
    VA_START (Args, Arg2);
    VaPtr = VA_ARG (Args, VOID *);
    VA_END (Args);

    Status = Protocol->StatusSend (Mask, Arg2, (UINTN)VaPtr);
    return Status;
}


//
// ---------------------------------------------------------------------------
// GetHobListAddress
// ---------------------------------------------------------------------------
//

/**
  Get the HOB (Hand-Off Block) list pointer.

  Walks the system table's configuration table array searching for
  the entry whose VendorGuid matches the PiHobList GUID.  The HOB
  list pointer is cached in the global gHobListAddress.

  @param[in] SystemTable  Pointer to the UEFI system table.

  @return EFI_PHYSICAL_ADDRESS  Pointer to the HOB list, or 0 on failure.
**/
EFI_PHYSICAL_ADDRESS
EFIAPI
GetHobListAddress (
    IN EFI_SYSTEM_TABLE *SystemTable
    )
{
    UINTN  Index;
    UINTN  TableCount;

    if (gHobListAddress != 0) {
        return gHobListAddress;
    }

    //
    // Initialize to NULL.
    //
    gHobListAddress = 0;

    TableCount = SystemTable->NumberOfTableEntries;
    if (TableCount == 0) {
        goto FAIL;
    }

    //
    // Walk the configuration table array.
    // Each entry is sizeof(EFI_CONFIGURATION_TABLE) bytes = 24 bytes:
    //   16 bytes VendorGuid + 8 bytes VendorTable pointer.
    //
    for (Index = 0; Index < TableCount; Index++) {
        EFI_CONFIGURATION_TABLE *ConfigEntry;

        ConfigEntry = (EFI_CONFIGURATION_TABLE *)
            ((UINT8 *)SystemTable->ConfigurationTable +
             (Index * sizeof (EFI_CONFIGURATION_TABLE)));

        if (CompareGuidU64 (
                &gPiHobListGuid,
                &ConfigEntry->VendorGuid
                )) {
            //
            // Found the HOB list entry.
            //
            gHobListAddress = (EFI_PHYSICAL_ADDRESS)ConfigEntry->VendorTable;
            break;
        }
    }

    if (gHobListAddress == 0) {

FAIL:
        //
        // ASSERT failure on debug builds. The original binary called
        // send the assertion details with:
        //   mask  = 0x80000000LL
        //   arg2  = "\nASSERT_EFI_ERROR (Status = %r)\n" string ref
        //   arg3  = 0x800000000000000EuLL (EFI_NOT_FOUND)
        //
        SendBoardInfoIfNetworkStackEnabled (
            0x80000000LL,
            (UINTN)"\nASSERT_EFI_ERROR (Status = %r)\n",
            0x800000000000000EuLL
            );

        AmiBoardInfoAssertCallback (
            "e:\\hs\\MdePkg\\Library\\DxeHobLib\\HobLib.c",
            54,
            "!EFI_ERROR (Status)"
            );
    }

    if (gHobListAddress == 0) {
        AmiBoardInfoAssertCallback (
            "e:\\hs\\MdePkg\\Library\\DxeHobLib\\HobLib.c",
            55,
            "mHobList != ((void *) 0)"
            );
    }

    return gHobListAddress;
}


//
// ---------------------------------------------------------------------------
// InitializeNetworkStackVar
// ---------------------------------------------------------------------------
//

/**
  Initialize the AmiBoardInfo variable.

  Reads the "NetworkStackVar" (AmiBoardInfoVar) UEFI variable.  If not
  found (EFI_NOT_FOUND = 0x800000000000000EuLL), allocates a buffer,
  sets default values, and writes the variable.

  Default values:
    Flags      = 0x1000000
    BootFlag   = 1
    word_D46   = 0

  @param[in] N10Param      Unused parameter (always 10).
  @param[in] SystemTable    Pointer to the UEFI system table.

  @return EFI_SUCCESS  Always returns success.
**/
EFI_STATUS
EFIAPI
InitializeNetworkStackVar (
    IN UINTN             N10Param,
    IN EFI_SYSTEM_TABLE *SystemTable
    )
{
    UINTN       VarSize;
    EFI_STATUS  Status;

    //
    // Cache services if not already set (by entry point).
    //
    if (gSystemTable == NULL) {
        gSystemTable = SystemTable;
        gBootServices = SystemTable->BootServices;
        gRuntimeServices = SystemTable->RuntimeServices;
    }

    //
    // Try to read the existing variable.
    //
    VarSize = sizeof (gNetworkStackVarData);
    Status  = gRuntimeServices->GetVariable (
                      NETWORK_STACK_VAR_NAME,
                      &gNetworkStackVarGuid,
                      NULL,             // Attributes (returned)
                      &VarSize,
                      &gNetworkStackVarData
                      );

    if (Status == EFI_NOT_FOUND) {
        //
        // Variable does not exist: zero the buffer and set defaults.
        //
        gBootServices->SetMem (&gNetworkStackVarData, sizeof (gNetworkStackVarData), 0);

        gNetworkStackReservedWord = 0;
        gNetworkStackBootFlag = 1;
        gNetworkStackVarData.Flags = 0x1000000;

        //
        // Write the variable as non-volatile + boot + runtime.
        //
        Status = gRuntimeServices->SetVariable (
                        NETWORK_STACK_VAR_NAME,
                        &gNetworkStackVarGuid,
                        (EFI_VARIABLE_NON_VOLATILE |
                         EFI_VARIABLE_BOOTSERVICE_ACCESS |
                         EFI_VARIABLE_RUNTIME_ACCESS),
                        sizeof (gNetworkStackVarData),
                        &gNetworkStackVarData
                        );

        //
        // Original driver does not check SetVariable status.
        //
    }

    return EFI_SUCCESS;
}


//
// ---------------------------------------------------------------------------
// NetworkStackSetupScreenEntryPoint
// ---------------------------------------------------------------------------
//

/**
  Entry point for the AMI Board Info (NetworkStackSetupScreen) DXE driver.

  Saves ImageHandle, SystemTable, BootServices, and RuntimeServices
  into global variables, then initializes the HOB list, and finally
  calls InitializeNetworkStackVar().

  @param[in] ImageHandle  Handle of this driver image.
  @param[in] SystemTable  Pointer to the UEFI system table.

  @return EFI_STATUS      Return value from InitializeNetworkStackVar().
**/
EFI_STATUS
EFIAPI
NetworkStackSetupScreenEntryPoint (
    IN EFI_HANDLE        ImageHandle,
    IN EFI_SYSTEM_TABLE *SystemTable
    )
{
    //
    // Save global handles (standard UEFI boot services library pattern).
    //
    gImageHandle = ImageHandle;
    ASSERT (gImageHandle != NULL);

    gSystemTable = SystemTable;
    ASSERT (gSystemTable != NULL);

    gBootServices = SystemTable->BootServices;
    ASSERT (gBootServices != NULL);

    gRuntimeServices = SystemTable->RuntimeServices;
    ASSERT (gRuntimeServices != NULL);

    //
    // Locate and cache the HOB list pointer.
    //
    GetHobListAddress (SystemTable);

    //
    // Initialize the board info variable ("NetworkStackVar").
    //
    return InitializeNetworkStackVar (10, SystemTable);
}