/** @file

;******************************************************************************
;* Copyright (c) 2021, Insyde Software Corp. All Rights Reserved.
;*
;* You may not reproduce, distribute, publish, display, perform, modify, adapt,
;* transmit, broadcast, present, recite, release, license or otherwise exploit
;* any part of this publication in any form, by any means, without the prior
;* written permission of Insyde Software Corporation.
;*
;******************************************************************************
*/
#include "LimitedToInternalStorage.h"
#include <SetupConfig.h>
#include <Library/UefiLib.h>
#include <Library/BaseMemoryLib.h>
#include <Library/BaseLib.h>
#include <Library/DevicePathLib.h>
#include <Library/MemoryAllocationLib.h>
#include <Library/UefiBootServicesTableLib.h>
#include <Library/UefiRuntimeServicesTableLib.h>
#include <Library/PrintLib.h>
#include <Library/DebugLib.h>
#include <Library/VariableLib.h>
#include <Library/GenericBdsLib.h>
#include <Library/H2OCpLib.h>
#include <Protocol/BlockIo.h>
#include <Protocol/DevicePath.h>

/**
  Is Limited to internal storage.

  @param  None.

  @retval TRUE                          Boot device will be limited to internal storage.
  @retval FALSE                         Boot device will not be limited to internal storage.
**/
BOOLEAN
L05IsLimitedToInternalStorage (
  VOID
  )
{
  EFI_STATUS                            Status;
  UINTN                                 BufferSize;
  SYSTEM_CONFIGURATION                  *SetupNvData;
  BOOLEAN                               IsLimited;

  SetupNvData = NULL;
  IsLimited   = FALSE;

  BufferSize  = sizeof (SYSTEM_CONFIGURATION);
  SetupNvData = AllocateZeroPool (BufferSize);

  if (SetupNvData == NULL) {
    return FALSE;
  }

  Status = CommonGetVariable (
             L"Setup",
             &gSystemConfigurationGuid,
             &BufferSize,
             SetupNvData
             );

  if (EFI_ERROR (Status)) {
    return FALSE;
  }

  if (SetupNvData->L05NaturalFileGuard == 1) {  // 0:Disabled, 1:Enabled
    IsLimited = TRUE;
  }

  if (SetupNvData->L05DeviceGuard == 1) {  // 0:Disabled, 1:Enabled
    IsLimited = TRUE;
  }

  FreePool (SetupNvData);

  return IsLimited;
}

/**
  Limited to Internal Storage Process.
  Check if platform is Limited to internal storage.
    YES - If boot option is not internal storage, set attributes to inactive.
    NO  - Set attributes to active.

  @param  None.
**/
VOID
LimitedToInternalStorageProcess (
  VOID
  )
{
  EFI_STATUS                            Status;
  UINTN                                 BlockIoHandleCount;
  EFI_HANDLE                            *BlockIoBuffer;
  UINTN                                 BlockIndex;
  EFI_DEVICE_PATH_PROTOCOL              *BlockIoDevicePath;
  EFI_BLOCK_IO_PROTOCOL                 *BlkIo;
  UINT16                                *BootOrder;
  UINTN                                 BootOrderSize;
  UINTN                                 BootDevNum;
  UINTN                                 Index;
  UINT16                                BootOptionName[BOOT_OPTION_MAX_CHAR];
  UINT8                                 *BootOption;
  UINTN                                 BootOptionSize;
  UINT8                                 *TempPtr;
  UINT32                                *Attribute;
  CHAR16                                *Description;
  UINTN                                 DescriptionSize;
  EFI_DEVICE_PATH_PROTOCOL              *DevicePath;
  UINT16                                DevicePathSize;
  BOOLEAN                               InternalStorage;

  BlockIoHandleCount = 0;
  BlockIoBuffer      = NULL;
  BlockIoDevicePath  = NULL;
  BlkIo              = NULL;
  BootOrder          = NULL;
  BootOrderSize      = 0;
  BootDevNum         = 0;
  InternalStorage    = FALSE;

  BlockIoBuffer = NULL;
  Status = gBS->LocateHandleBuffer (
                  ByProtocol,
                  &gEfiBlockIoProtocolGuid,
                  NULL,
                  &BlockIoHandleCount,
                  &BlockIoBuffer
                  );

  if (EFI_ERROR (Status) || BlockIoBuffer == NULL) {
    return;
  }

  Status = CommonGetVariableDataAndSize (
             L"BootOrder",
             &gEfiGlobalVariableGuid,
             &BootOrderSize,
             (VOID **) &BootOrder
             );

  if (EFI_ERROR (Status) || BootOrder == NULL) {
    return;
  }

  BootDevNum = BootOrderSize / sizeof (UINT16);

  for (Index = 0; Index < BootDevNum; Index++) {
    InternalStorage = FALSE;

    UnicodeSPrint (BootOptionName, sizeof (BootOptionName), L"Boot%04x", BootOrder[Index]);
    Status = CommonGetVariableDataAndSize (
               BootOptionName,
               &gEfiGlobalVariableGuid,
               &BootOptionSize,
               (VOID **) &BootOption
               );

    if (EFI_ERROR (Status) || BootOption == NULL) {
      continue;
    }

    TempPtr = BootOption;
    Attribute = (UINT32 *) TempPtr;
    TempPtr += sizeof (UINT32);
    DevicePathSize = *((UINT16 *) TempPtr);
    TempPtr += sizeof (UINT16);
    Description = (CHAR16 *) TempPtr;
    DescriptionSize = StrSize (Description);
    TempPtr += DescriptionSize;
    DevicePath = (EFI_DEVICE_PATH_PROTOCOL *) TempPtr;
    TempPtr += DevicePathSize;

    if ((DevicePathType (DevicePath) == MEDIA_DEVICE_PATH) &&
        (DevicePathSubType (DevicePath) == MEDIA_HARDDRIVE_DP)) {
      //
      // Find hard drive media device path.
      //
      if (BlockIoBuffer != NULL && BlockIoHandleCount != 0) {
        for (BlockIndex = 0; BlockIndex < BlockIoHandleCount; BlockIndex++) {
          BlockIoDevicePath = DevicePathFromHandle (BlockIoBuffer[BlockIndex]);

          if (BlockIoDevicePath == NULL) {
            continue;
          }

          if (MatchPartitionDevicePathNode (BlockIoDevicePath, (HARDDRIVE_DEVICE_PATH *) DevicePath)) {
            //
            // Find the matched partition device path.
            //
            Status = gBS->HandleProtocol (
                            BlockIoBuffer[BlockIndex],
                            &gEfiBlockIoProtocolGuid,
                            (VOID **) &BlkIo
                            );

            if (EFI_ERROR (Status)) {
              continue;
            }

            if (!BlkIo->Media->RemovableMedia) {
              //
              // Not removable media, that is internal storage.
              //
              InternalStorage = TRUE;
              break;
            }
          }
        }
      }
    }

    if (!InternalStorage && L05IsLimitedToInternalStorage ()) {
      //
      // The boot option is not internal storage.
      // So set Attributes to inactive for limited to internal storage.
      //
      *Attribute  &= (~((UINT32) LOAD_OPTION_ACTIVE));

    } else {
      *Attribute  |= (((UINT32) LOAD_OPTION_ACTIVE));
    }

    Status = CommonSetVariable (
               BootOptionName,
               &gEfiGlobalVariableGuid,
               EFI_VARIABLE_NON_VOLATILE | EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_RUNTIME_ACCESS,
               BootOptionSize,
               BootOption
               );

    FreePool (BootOption);
  }

  if (BootOrder != NULL) {
    FreePool (BootOrder);
  }

  return;
}

/**
  Limited to Internal Storage update Target HDD Variable.

  @param  None.
**/
VOID
LimitedToInternalStorageUpdateTargetHddVariable (
  VOID
  )
{
  EFI_STATUS                            Status;
  EFI_DEVICE_PATH_PROTOCOL              *TargetHddDevPathVar;
  UINTN                                 BufferSize;
  SYSTEM_CONFIGURATION                  *SetupNvData;
  UINT8                                 BootType;

  TargetHddDevPathVar = NULL;
  SetupNvData         = NULL;

  if (!L05IsLimitedToInternalStorage ()) {
    return;
  }

  //
  // Get TargetHddDevPath variable of fast boot.
  //
  TargetHddDevPathVar = CommonGetVariableData (
                          L"TargetHddDevPath",
                          &gEfiGenericVariableGuid
                          );

  if (TargetHddDevPathVar != NULL) {
    FreePool (TargetHddDevPathVar);
    return;
  }

  BufferSize  = sizeof (SYSTEM_CONFIGURATION);
  SetupNvData = AllocateZeroPool (BufferSize);

  if (SetupNvData == NULL) {
    return;
  }

  Status = CommonGetVariable (
             L"Setup",
             &gSystemConfigurationGuid,
             &BufferSize,
             SetupNvData
             );

  if (EFI_ERROR (Status)) {
    return;
  }

  gBdsServices->GetBootType (gBdsServices, &BootType);

  if (SetupNvData->Win8FastBoot == 0 &&
      BootType == EFI_BOOT_TYPE) {
    UpdateTargetHddVariable ();
  }

  FreePool (SetupNvData);

  return;
}

/**
  Limited to internal storage at Boot Device Enum After check point.
  It will update attribute of boot optoion to limited to internal storage.

  @param  Event                         Event whose notification function is being invoked.
  @param  Context                       Pointer to the notification function's context.
**/
VOID
EFIAPI
LimitedToInternalStorageNotify (
  IN EFI_EVENT                          Event,
  IN VOID                               *Context
  )
{
  BOOLEAN                               IsBootOrderHookEnabled;

  IsBootOrderHookEnabled = BdsLibIsBootOrderHookEnabled ();

  //
  // Update Option Attribute of BootOrder variable.
  //
  if (IsBootOrderHookEnabled) {
    BdsLibDisableBootOrderHook ();
  }

  LimitedToInternalStorageProcess ();

  //
  // Update Option Attribute of PhysicalBootOrder variable.
  //
  BdsLibEnableBootOrderHook ();
  LimitedToInternalStorageProcess ();

  if (!IsBootOrderHookEnabled) {
    BdsLibDisableBootOrderHook ();
  }

  //
  // Update Target HDD Variable
  //
  LimitedToInternalStorageUpdateTargetHddVariable ();

  return;
}

/**
  Limited dummy boot option to internal storage  after gEfiVariableWriteArchProtocolGuid is installed.
  When platform is limited to internal storage, update attribute of dummy boot option to limited to internal storage.

  @param  Event                         Pointer to this event
  @param  Context                       Event hanlder private data
**/
VOID
EFIAPI
LimitedDummyBootOptionToInternalStorageNotify (
  IN  EFI_EVENT                         Event,
  IN  VOID                              *Context
  )
{
  EFI_STATUS                            Status;
  VOID                                  *Interface;
  UINT16                                BootOptionName[BOOT_OPTION_MAX_CHAR];
  UINTN                                 BootOptionSize;
  UINT8                                 *BootOption;
  UINTN                                 Index;
  UINT32                                *Attribute;

  Status = gBS->LocateProtocol (
                  &gEfiVariableWriteArchProtocolGuid,
                  NULL,
                  (VOID **) &Interface
                  );

  if (EFI_ERROR (Status)) {
    return;
  }

  gBS->CloseEvent (Event);

  for (Index = (DummyBootOptionStartNum + 1); Index < DummyBootOptionEndNum; Index++) {
    UnicodeSPrint (BootOptionName, sizeof (BootOptionName), L"Boot%04x", Index);
    Status = CommonGetVariableDataAndSize (
               BootOptionName,
               &gEfiGlobalVariableGuid,
               &BootOptionSize,
               (VOID **) &BootOption
               );

    if (EFI_ERROR (Status) || BootOption == NULL) {
      continue;
    }

    Attribute = (UINT32 *) BootOption;

    if (L05IsLimitedToInternalStorage ()) {
      //
      // Set Attributes to inactive for limited to internal storage.
      //
      *Attribute  &= (~((UINT32) LOAD_OPTION_ACTIVE));

    } else {
      *Attribute  |= (((UINT32) LOAD_OPTION_ACTIVE));
    }

    Status = CommonSetVariable (
               BootOptionName,
               &gEfiGlobalVariableGuid,
               EFI_VARIABLE_NON_VOLATILE | EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_RUNTIME_ACCESS,
               BootOptionSize,
               BootOption
               );

    gBS->FreePool (BootOption);
  }

  return;
}

/**
  Limited boot next to internal storage  after gEfiVariableWriteArchProtocolGuid is installed.
  When platform is limited to internal storage, delete BootNext variable if BootNext option is dummy boot option.

  @param  Event                         Pointer to this event
  @param  Context                       Event hanlder private data
**/
VOID
EFIAPI
LimitedBootNextToInternalStorageNotify (
  IN  EFI_EVENT                         Event,
  IN  VOID                              *Context
  )
{
  EFI_STATUS                            Status;
  VOID                                  *Interface;
  UINT16                                *BootNext;

  BootNext = NULL;

  Status = gBS->LocateProtocol (
                  &gEfiVariableWriteArchProtocolGuid,
                  NULL,
                  (VOID **) &Interface
                  );

  if (EFI_ERROR (Status)) {
    return;
  }

  gBS->CloseEvent (Event);

  if (!L05IsLimitedToInternalStorage ()) {
    return;
  }

  BootNext = CommonGetVariableData (
               L"BootNext",
               &gEfiGlobalVariableGuid
               );

  if (BootNext == NULL) {
    return;
  }

  if (BdsLibIsDummyBootOption (*BootNext)) {
    CommonSetVariable (L"BootNext", &gEfiGlobalVariableGuid, 0, 0, NULL);
  }

  if (BootNext != NULL) {
    FreePool (BootNext);
  }

  return;
}

/**
  Limited to internal storage.

  @retval EFI_SUCCESS                   The operation completed successfully.
  @retval Others                        An unexpected error occurred.
**/
EFI_STATUS
LimitedToInternalStorage (
  VOID
  )
{
  EFI_STATUS                            Status;
  H2O_CP_HANDLE                         CpHandle;
  VOID                                  *Registration;

  CpHandle     = NULL;

  if (FeaturePcdGet (PcdH2OBdsCpBootDeviceEnumAfterSupported)) {
    Status = H2OCpRegisterHandler (
               &gH2OBdsCpBootDeviceEnumAfterGuid,
               LimitedToInternalStorageNotify,
               H2O_CP_MEDIUM,
               &CpHandle
               );

    if (EFI_ERROR (Status)) {
      DEBUG_CP ((DEBUG_ERROR, "Checkpoint Register Fail: %g (%r)\n", &gH2OBdsCpBootDeviceEnumAfterGuid, Status));

    } else {
      DEBUG_CP ((DEBUG_INFO, "Checkpoint Registered: %g (%r)\n", &gH2OBdsCpBootDeviceEnumAfterGuid, Status));
    }
  }

  Registration = NULL;
  EfiCreateProtocolNotifyEvent (
    &gEfiVariableWriteArchProtocolGuid,
    TPL_CALLBACK,
    LimitedDummyBootOptionToInternalStorageNotify,
    NULL,
    &Registration
    );

  Registration = NULL;
  EfiCreateProtocolNotifyEvent (
    &gEfiVariableWriteArchProtocolGuid,
    TPL_CALLBACK,
    LimitedBootNextToInternalStorageNotify,
    NULL,
    &Registration
    );

  return Status;
}