/**
 * HiiInfoEmbedded.c - UEFI DXE Driver Source
 * Module: HiiInfoEmbedded.efi (Index 0252)
 * Source: HR650X BIOS
 *
 * This module extracts Human Interface Infrastructure (HII) embedded information
 * from a BIOS image and serializes it into a compact compressed binary format
 * suitable for embedding into the BIOS ROM hole space. It includes a full DEFLATE
 * (LZ77 + Huffman coding) compression implementation for the output.
 *
 * Architecture: x86_64, Image Size: 0xd3c0, Functions: 99
 */

#include <Uefi.h>
#include <Library/DebugLib.h>
#include <Library/UefiBootServicesTableLib.h>
#include <Library/UefiRuntimeServicesTableLib.h>
#include <Library/MemoryAllocationLib.h>
#include <Library/BaseMemoryLib.h>
#include <Library/PrintLib.h>
#include <Library/DevicePathLib.h>
#include <Library/UefiLib.h>
#include <Protocol/Shell.h>
#include <Protocol/HiiPackageList.h>

//
// GUID definitions used by this module
//
#define HII_PACKAGE_LIST_GUID \
  { 0x7475EE58, 0x0E1C, 0x4F47, { 0x8D, 0x1A, 0x94, 0x93, 0x3F, 0x0D, 0x5F, 0x69 } }

#define HII_PLATFORM_SETUP_FORMSET_GUID \
  { 0x93039971, 0x1757, 0x4B18, { 0xAE, 0x4A, 0x86, 0x0B, 0xD4, 0x9D, 0x82, 0x09 } }

// ---------------------------------------------------------------------------
// Global Data - Boot Services / Runtime Services / HII Protocol pointers
// ---------------------------------------------------------------------------

EFI_HANDLE          gImageHandle  = NULL;
EFI_SYSTEM_TABLE   *gST           = NULL;
EFI_BOOT_SERVICES  *gBS           = NULL;
EFI_RUNTIME_SERVICES_TABLE *gRT   = NULL;
VOID               *gDebugConOut  = NULL;

EFI_HANDLE          gHobHandle    = NULL;
EFI_HII_DATABASE_PROTOCOL  *gHiiDatabase  = NULL;
EFI_HII_STRING_PROTOCOL    *gHiiString    = NULL;
EFI_SHELL_PROTOCOL         *gShellProtocol = NULL;
EFI_SHELL_PARAMETERS_PROTOCOL *gShellParams = NULL;
EFI_HANDLE                  gHiiImageHandle = NULL;

// Shell file operation function pointers
VOID *ShellGetFileInfo   = NULL;
VOID *ShellSetFileInfo   = NULL;
VOID *ShellReadFile      = NULL;
VOID *ShellWriteFile     = NULL;
VOID *ShellCloseFile     = NULL;
VOID *ShellDeleteFile    = NULL;
VOID *ShellGetFileSize   = NULL;
VOID *ShellCreateFile    = NULL;
VOID *ShellFlushFileEx   = NULL;
VOID *ShellOpenFile      = NULL;

// Linked list of HII string entries
VOID *gHiiStringList = NULL;

// State flags
UINT8   gBiosImageFromFile   = 0;
UINT8   gHiiDbAvailable      = 0;
UINT64  gCommandLineArgs     = 0;
UINT64  gArgCount            = 0;

// String buffer for output
VOID *gStringBuffer = NULL;

// ---------------------------------------------------------------------------
// Dynamic string buffer structure
// ---------------------------------------------------------------------------
typedef struct {
  CHAR8  *Buffer;
  UINT32  Allocated;
  UINT32  Used;
  UINT32  Pad;
} DYNAMIC_STRING_BUFFER;

// ---------------------------------------------------------------------------
// HII string linked list entry structure
// ---------------------------------------------------------------------------
#pragma pack(push, 1)
typedef struct _HII_STRING_ENTRY {
  UINT16  StringId;
  UINT8   PackageType;
  UINT64  StringPtr;
  struct _HII_STRING_ENTRY *Next;
} HII_STRING_ENTRY;
#pragma pack(pop)

#define MAX_COMPRESSED_SIZE  0x400000
#define READ_CHUNK_SIZE      0x1000000

// ===========================================================================
// UEFI Library Boot Services Table Constructor
// ===========================================================================
VOID
EFIAPI
UefiBootServicesTableLib_Constructor (
  VOID
  )
{
  gImageHandle = ImageHandle;
  ASSERT (gImageHandle != NULL);
  gST = SystemTable;
  ASSERT (gST != NULL);
  gBS = SystemTable->BootServices;
  ASSERT (gBS != NULL);
  gRT = SystemTable->RuntimeServices;
  ASSERT (gRT != NULL);

  GetHobList();
  gHiiDatabase  = NULL;
  gShellParams  = NULL;
  gShellProtocol = NULL;
  gHiiImageHandle = NULL;
  ModuleEntryInit (ImageHandle);
}

// ===========================================================================
// Module Entry Point (AutoGen)
// ===========================================================================
EFI_STATUS
EFIAPI
_ModuleEntryPoint (
  IN EFI_HANDLE        ImageHandle,
  IN EFI_SYSTEM_TABLE *SystemTable
  )
{
  EFI_STATUS  Status;
  UefiBootServicesTableLib_Constructor ();
  Status = HiiInfoEmbedMain (ImageHandle, SystemTable);
  ModuleDestructor (ImageHandle);
  return Status;
}

// ===========================================================================
// Module Destructor (AutoGen)
// ===========================================================================
EFI_STATUS
ModuleDestructor (
  IN EFI_HANDLE ImageHandle
  )
{
  EFI_STATUS  Status;
  Status = UnloadImageProtocol (ImageHandle);
  if (EFI_ERROR (Status)) {
    DEBUG ((EFI_D_ERROR, "\nASSERT_EFI_ERROR (Status = %r)\n", Status));
    ASSERT (!EFI_ERROR (Status));
  }
  return Status;
}

// ===========================================================================
// Unload Image Protocol
// ===========================================================================
EFI_STATUS
UnloadImageProtocol (
  IN EFI_HANDLE ImageHandle
  )
{
  if (gShellProtocol != NULL) {
    gBS->UninstallMultipleProtocolInterfaces (
           gImageHandle ? gImageHandle : ImageHandle,
           &gEfiShellProtocolGuid,
           ImageHandle, NULL);
    gShellProtocol = NULL;
  }
  if (gShellParams != NULL) {
    gBS->UninstallMultipleProtocolInterfaces (
           ImageHandle, &gEfiShellParametersProtocolGuid,
           ImageHandle, NULL);
    gShellParams = NULL;
  }
  if (gHiiDatabase != NULL) {
    gBS->UninstallMultipleProtocolInterfaces (
           ImageHandle, &gEfiHiiDatabaseProtocolGuid,
           ImageHandle, NULL);
    gHiiDatabase = NULL;
  }
  if (gHiiImageHandle != NULL) {
    gBS->UninstallMultipleProtocolInterfaces (
           ImageHandle, &gEfiHiiImageProtocolGuid,
           ImageHandle, NULL);
    gHiiImageHandle = NULL;
  }
  gImageHandle = NULL;
  return EFI_SUCCESS;
}

// ===========================================================================
// Module Entry Init - Locate protocols and set up function pointers
// ===========================================================================
EFI_STATUS
ModuleEntryInit (
  IN EFI_HANDLE ImageHandle
  )
{
  EFI_STATUS  Status;

  // Try locating HII Database protocol
  Status = gBS->OpenProtocol (ImageHandle,
    &gEfiHiiDatabaseProtocolGuid, (VOID**)&gHiiDatabase,
    ImageHandle, 0, EFI_OPEN_PROTOCOL_BY_HANDLE_PROTOCOL);
  if (EFI_ERROR (Status)) {
    Status = gBS->LocateProtocol (&gEfiHiiDatabaseProtocolGuid, NULL,
      (VOID**)&gHiiDatabase);
    if (EFI_ERROR (Status)) gHiiDatabase = NULL;
  }

  // Try locating Shell protocol
  gBS->OpenProtocol (ImageHandle, &gEfiShellProtocolGuid,
    (VOID**)&gShellProtocol, ImageHandle, 0,
    EFI_OPEN_PROTOCOL_BY_HANDLE_PROTOCOL);
  gBS->OpenProtocol (ImageHandle, &gEfiShellParametersProtocolGuid,
    (VOID**)&gShellParams, ImageHandle, 0,
    EFI_OPEN_PROTOCOL_BY_HANDLE_PROTOCOL);

  if (gShellProtocol == NULL || gShellParams == NULL) {
    Status = LocateShellProtocol (ImageHandle);
    if (EFI_ERROR (Status)) {
      DEBUG ((EFI_D_ERROR, "Status: 0x%08x\r\n", Status));
      gShellProtocol = NULL;
    }
    gBS->OpenProtocol (ImageHandle, &gEfiShellParametersProtocolGuid,
      (VOID**)&gShellParams, ImageHandle, 0,
      EFI_OPEN_PROTOCOL_BY_HANDLE_PROTOCOL);
    if (EFI_ERROR (Status)) gShellParams = NULL;
  }

  if ((gShellProtocol == NULL || gShellParams == NULL) &&
      (gHiiDatabase == NULL || gHiiImageHandle == NULL)) {
    return EFI_NOT_FOUND;
  }

  // Set up shell function pointers from protocol vtable (indices 22-36)
  if (gHiiDatabase != NULL) {
    ShellGetFileInfo   = (VOID*)((UINT64*)gHiiDatabase)[22];
    ShellSetFileInfo   = (VOID*)((UINT64*)gHiiDatabase)[23];
    ShellReadFile      = (VOID*)((UINT64*)gHiiDatabase)[27];
    ShellWriteFile     = (VOID*)((UINT64*)gHiiDatabase)[28];
    ShellCloseFile     = (VOID*)((UINT64*)gHiiDatabase)[25];
    ShellDeleteFile    = (VOID*)((UINT64*)gHiiDatabase)[29];
    ShellGetFileSize   = (VOID*)((UINT64*)gHiiDatabase)[31];
    ShellCreateFile    = (VOID*)((UINT64*)gHiiDatabase)[32];
    ShellFlushFileEx   = (VOID*)((UINT64*)gHiiDatabase)[33];
    ShellOpenFile      = (VOID*)((UINT64*)gHiiDatabase)[36];
  } else {
    ShellGetFileInfo   = (VOID*)ShellGetFileInfo_stub;
    ShellSetFileInfo   = (VOID*)ShellSetFileInfo_stub;
    ShellReadFile      = (VOID*)ShellReadFile_stub;
    ShellWriteFile     = (VOID*)ShellWriteFile_stub;
    ShellCloseFile     = (VOID*)ShellCloseFile_stub;
    ShellDeleteFile    = (VOID*)ShellDeleteFile_stub;
    ShellGetFileSize   = (VOID*)ShellGetFileSize_stub;
    ShellCreateFile    = (VOID*)ShellCreateFile_stub;
    ShellFlushFileEx   = (VOID*)ShellFlushFileEx_stub;
    ShellOpenFile      = (VOID*)DefaultShellOpenFile;
  }
  return EFI_SUCCESS;
}

// ===========================================================================
// Locate Shell Protocol via Handle Database
// ===========================================================================
EFI_STATUS
LocateShellProtocol (
  IN EFI_HANDLE ImageHandle
  )
{
  EFI_STATUS   Status;
  UINTN        HandleCount;
  EFI_HANDLE  *HandleBuffer;
  UINTN        Index;

  HandleCount = 0;
  HandleBuffer = NULL;
  gShellProtocol = NULL;

  Status = gBS->LocateHandleBuffer (ByProtocol,
    &gEfiShellProtocolGuid, NULL, &HandleCount,
    (EFI_HANDLE**)&HandleBuffer);

  if (Status == EFI_BUFFER_TOO_SMALL && HandleCount > 0) {
    HandleBuffer = AllocatePool (HandleCount * sizeof (EFI_HANDLE));
    if (HandleBuffer == NULL) return EFI_OUT_OF_RESOURCES;
    gBS->LocateHandleBuffer (ByProtocol, &gEfiShellProtocolGuid,
      NULL, &HandleCount, (EFI_HANDLE**)&HandleBuffer);
  }

  if (EFI_ERROR (Status)) goto Done;

  if ((HandleCount & ~7ULL) > 0) {
    for (Index = 0; Index < (HandleCount & ~7ULL); Index++) {
      Status = gBS->OpenProtocol (HandleBuffer[Index],
        &gEfiShellProtocolGuid, (VOID**)&gShellProtocol,
        ImageHandle, 0, EFI_OPEN_PROTOCOL_BY_HANDLE_PROTOCOL);
      if (!EFI_ERROR (Status) &&
          CompareGuidWrapper ((UINT8*)gShellProtocol + 96,
            &gEfiDxeServicesTableGuid)) {
        gImageHandle = HandleBuffer[Index];
        break;
      }
    }
    Status = EFI_SUCCESS;
  } else {
    Status = EFI_NOT_FOUND;
  }

Done:
  if (HandleBuffer != NULL) FreePool (HandleBuffer);
  return Status;
}

// ===========================================================================
// MAIN HII INFO EMBED FUNCTION
// ===========================================================================
EFI_STATUS
HiiInfoEmbedMain (
  IN EFI_HANDLE  ImageHandle,
  IN EFI_SYSTEM_TABLE *SystemTable
  )
{
  EFI_STATUS              Status;
  VOID                   *BiosFileBuffer;
  UINTN                   FileSize;
  VOID                   *BiosHeader;
  EFI_TIME                BuildTime;
  UINTN                   CompressedSize;
  UINTN                   Index;
  CHAR16                  FileName[200];
  VOID                   *FileHandle;
  VOID                   *OutputBuf;
  UINTN                   SerializedSize;
  VOID                   *StrBuf;

  Status = EFI_SUCCESS;
  FileSize = 0;
  BiosFileBuffer = NULL;
  BiosHeader = NULL;
  OutputBuf = NULL;
  gBiosImageFromFile = 0;

  // Validate command line args (count 0-3)
  if (gArgCount == 0 || gArgCount > 3) {
    PrintToConOut (L"Please input BIOS image file.\n");
    return EFI_INVALID_PARAMETER;
  }
  if (gArgCount > 3) {
    PrintToConOut (L"Invalid Parameter.\n");
    return EFI_INVALID_PARAMETER;
  }

  // Parse command: "info" or "info <filename>"
  if (gArgCount == 2) {
    if (!StrCmpWrapper (*(CHAR16**)(gCommandLineArgs + 8), L"info"))
      goto ProcessHii;
  } else if (gArgCount == 3) {
    if (StrCmpWrapper (*(CHAR16**)(gCommandLineArgs + 16), L"info")) {
      PrintToConOut (L"Invalid Parameter.\n");
      return EFI_INVALID_PARAMETER;
    }
    gBiosImageFromFile = 1;
  }

  // Open BIOS image file
  {
    CHAR16 *FileNamePtr = *(CHAR16**)(gCommandLineArgs + 8);
    if (FileNamePtr != NULL) {
      Status = ShellOpenFileByName (FileNamePtr, &FileHandle,
                 EFI_FILE_MODE_READ);
      if (EFI_ERROR (Status)) {
        PrintToConOut (L"File Not Found.\n");
        return EFI_NOT_FOUND;
      }
      Status = gBS->AllocatePool (EfiBootServicesData,
                 MAX_COMPRESSED_SIZE, &BiosFileBuffer);
      if (Status == EFI_BUFFER_TOO_SMALL) {
        PrintToConOut (L"Out of Resources.\n");
        return EFI_OUT_OF_RESOURCES;
      }
      // Read file in chunks
      {
        UINTN ChunkSize = READ_CHUNK_SIZE;
        while (!EFI_ERROR (Status) && ChunkSize == READ_CHUNK_SIZE &&
               FileSize < MAX_COMPRESSED_SIZE) {
          Status = ShellReadFile (FileHandle, &ChunkSize,
                     (UINT8*)BiosFileBuffer + FileSize);
          if (!EFI_ERROR (Status)) FileSize += ChunkSize;
        }
      }
      ShellCloseFile (FileHandle);

      // Scan for BIOS flash image header GUID
      if (FileSize > 0) {
        for (Index = 0; Index < FileSize; Index += 8) {
          if (!CompareMemWrapper ((UINT8*)BiosFileBuffer + Index,
                                   &gFlashBiosIdGuid, 16) &&
              !CompareMemWrapper ((UINT8*)BiosFileBuffer + Index + 144,
                                   &gBiosImageHeaderGuid, 16)) {
            BiosHeader = (UINT8*)BiosFileBuffer + Index;
            break;
          }
        }
        if (BiosHeader == NULL) {
          PrintToConOut (L"Invalid BIOS Image\n");
          gBS->FreePool (BiosFileBuffer);
          return EFI_INVALID_PARAMETER;
        }
      }
    }
  }

ProcessHii:
  // Initialize HII protocols
  Status = HiiInitGlobalProtocols ();
  if (EFI_ERROR (Status)) {
    DEBUG ((EFI_D_ERROR, "InitializeGlobalProtocols FAILED\n"));
    goto Cleanup;
  }

  // Allocate global string info buffer
  gStringBuffer = HiiGetHiiHandleList ();
  if (gStringBuffer != NULL)
    ((DYNAMIC_STRING_BUFFER*)gStringBuffer)->Used = 6;

  if (gHiiDbAvailable) {
    // Read HII database from UEFI variable "HiiDB"
    UINTN     DataSize = 8;
    EFI_GUID  HiiDbGuid = HII_PACKAGE_LIST_GUID;
    VOID     *HiiDbBuffer = NULL;
    Status = gRT->GetVariable (L"HiiDB", &HiiDbGuid, NULL, &DataSize, NULL);
    if (Status == EFI_BUFFER_TOO_SMALL) {
      gBS->AllocatePool (EfiBootServicesData, DataSize, &HiiDbBuffer);
      Status = gRT->GetVariable (L"HiiDB", &HiiDbGuid, NULL,
                 &DataSize, HiiDbBuffer);
      if (!EFI_ERROR (Status)) {
        for (Index = 0; Index < *(UINT32*)HiiDbBuffer;
             Index += *(UINT32*)((UINT8*)HiiDbBuffer + 16)) {
          HiiExtractEmbeddedInfo (0, (UINT8*)HiiDbBuffer + Index);
        }
      } else {
        DEBUG ((EFI_D_ERROR, "Can not find HiiDB. Status = %r\n", Status));
        gHiiDbAvailable = 0;
      }
    } else {
      DEBUG ((EFI_D_ERROR, "Can not find HiiDB. Status = %r\n", Status));
      gHiiDbAvailable = 0;
    }
  }

  if (!gHiiDbAvailable) {
    // Walk all HII package lists
    UINT16 NumPackageLists;
    UINT64 *PackageListArray;
    UINTN  Index2;
    StrBuf = HiiGetAllPackageLists (&NumPackageLists);
    if (StrBuf != NULL && NumPackageLists > 0) {
      PackageListArray = AllocatePool (sizeof(UINT64) * NumPackageLists);
      if (PackageListArray == NULL) goto Cleanup;
      for (Index2 = 0; Index2 < NumPackageLists; Index2++) {
        HiiExtractEmbeddedInfo (
          *(UINT64*)((UINT8*)StrBuf + Index2),
          &PackageListArray[Index2]);
      }
      gBS->FreePool (PackageListArray);
    } else {
      DEBUG ((EFI_D_ERROR, "Can not get any HiiPkgHandle\n"));
      goto Cleanup;
    }
  }

  // Serialize and compress extracted info
  StrBuf = HiiSerializeStringsToBuf (gStringBuffer, 0, 1);
  SerializedSize = StrSizeWrapper (StrBuf);
  CompressedSize = SerializedSize;

  {
    UINTN OutSize = CompressedSize;
    DeflateCompressEx (StrBuf, SerializedSize, NULL, &OutSize);
    CompressedSize = OutSize;
    if (OutSize < 0x1000000) {
      Status = gBS->AllocatePool (EfiBootServicesData, CompressedSize,
                 &OutputBuf);
      if (!EFI_ERROR (Status))
        DeflateCompressEx (StrBuf, SerializedSize, OutputBuf,
          &CompressedSize);
    }
  }

  if (!EFI_ERROR (Status)) {
    gRT->GetTime (&BuildTime, NULL);
    if (BiosHeader != NULL) {
      // Check ROM hole space availability
      if ((CompressedSize & 0xFF000000) != 0 ||
          CompressedSize >= (*(UINT32*)((UINT8*)BiosHeader + 168) - 172)) {
        PrintToConOut (L"Insufficient ROM Hole space to embedded it.\n");
      } else {
        // Embed compressed data into ROM hole
        gBS->SetMem ((UINT8*)BiosHeader + 172, CompressedSize + 28, -1);
        gBS->CopyMem ((UINT8*)BiosHeader + 172 + 4, OutputBuf,
          CompressedSize);
        // Generate filename with timestamp: YYYY.MM.DD.HH.MM.SS.<filename>
        SPrintWrapper (FileName, 200,
          L"%04d.%02d.%02d.%02d.%02d.%02d.%s",
          BuildTime.Year, BuildTime.Month, BuildTime.Day,
          BuildTime.Hour, BuildTime.Minute, BuildTime.Second,
          *(UINT64*)(gCommandLineArgs + 8));
        if (ShellOpenFileByName (FileName, &FileHandle,
              ShellOpenFileEx) >= 0) {
          ShellWriteFile (FileHandle, &SerializedSize, BiosFileBuffer);
          ShellCloseFile (FileHandle);
        }
        PrintToConOut (L"%s is generated.\n", FileName);
      }
    }

    gBS->SetMem (FileName, 400, 0);

    if (gBiosImageFromFile != 0) {
      // Save uncompressed HiiInfoResult file
      SPrintWrapper (FileName, 200,
        L"%04d.%02d.%02d.%02d.%02d.%02d.HiiInfoResult",
        BuildTime.Year, BuildTime.Month, BuildTime.Day,
        BuildTime.Hour, BuildTime.Minute, BuildTime.Second);
      if (ShellOpenFileByName (FileName, &FileHandle,
            ShellOpenFileEx) >= 0) {
        ShellWriteFile (FileHandle, &SerializedSize, StrBuf);
        ShellCloseFile (FileHandle);
      }

      // Save compressed HiiInfoResultCompressed file
      SPrintWrapper (FileName, 200,
        L"%04d.%02d.%02d.%02d.%02d.%02d.HiiInfoResultCompressed",
        BuildTime.Year, BuildTime.Month, BuildTime.Day,
        BuildTime.Hour, BuildTime.Minute, BuildTime.Second);
      if (ShellOpenFileByName (FileName, &FileHandle,
            ShellOpenFileEx) >= 0) {
        if (CompressedSize != 0)
          ShellWriteFile (FileHandle, &CompressedSize, OutputBuf);
        ShellCloseFile (FileHandle);
      }
    }

    if (StrBuf != NULL) FreePool (StrBuf);
  }

Cleanup:
  if (gStringBuffer != NULL) {
    HiiFreeHiiHandleList ();
    gStringBuffer = NULL;
  }
  if (BiosFileBuffer != NULL) gBS->FreePool (BiosFileBuffer);
  return Status;
}