/** @file

;******************************************************************************
;* Copyright (c) 2012 - 2016, 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 <UefiDriveIdentification.h>

/*++

Routine Description:

  Read the EFI variable (VendorGuid/Name) and return a dynamically allocated
  buffer, and the size of the buffer. If failure return NULL.

Arguments:

  Name         - String part of EFI variable name

  VendorGuid   - GUID part of EFI variable name

  VariableSize - Returns the size of the EFI variable that was read

Returns:

  Dynamically allocated memory that contains a copy of the EFI variable.
  Caller is responsible freeing the buffer.

  NULL         - Variable was not read

--*/
VOID *
GetVariableAndSize (
  IN  CHAR16                            *Name,
  IN  EFI_GUID                          *VendorGuid,
  OUT UINTN                             *VariableSize
  )
{
  EFI_STATUS                            Status;
  UINTN                                 BufferSize;
  VOID                                  *Buffer;

  Buffer = NULL;

  //
  // Pass in a zero size buffer to find the required buffer size.
  //
  BufferSize  = 0;
  Status      = gRT->GetVariable (
                       Name,
                       VendorGuid,
                       NULL,
                       &BufferSize,
                       Buffer
                       );

  if (Status == EFI_BUFFER_TOO_SMALL) {
    //
    // Allocate the buffer to return
    //
    Buffer = AllocateZeroPool (BufferSize);

    if (Buffer == NULL) {
      return NULL;
    }

    //
    // Read variable into the allocated buffer.
    //
    Status = gRT->GetVariable (
                    Name,
                    VendorGuid,
                    NULL,
                    &BufferSize,
                    Buffer
                    );

    if (EFI_ERROR (Status)) {
      BufferSize = 0;
      FreePool (Buffer);
      Buffer = NULL;
    }
  }

  *VariableSize = BufferSize;
  return Buffer;
}

EFI_STATUS
UefiDriveIdentificationEntry (
  IN EFI_HANDLE                         ImageHandle,
  IN EFI_SYSTEM_TABLE                   *SystemTable
  )
{
  EFI_STATUS                            Status;
  EFI_EVENT                             ReadyToBootEvent;

  Status = EfiCreateEventReadyToBootEx (
             TPL_CALLBACK,
             ConvertBootToLBoot,
             NULL,
             &ReadyToBootEvent
             );

  return Status;
}

EFI_STATUS
L05VendorDefinedData (
  IN  EFI_DEVICE_PATH_PROTOCOL          *DevicePathPtr,
  OUT UINT8                             *VendorDefinedData
  )
{
  EFI_DEVICE_PATH_PROTOCOL              *DevicePathNode;
  EFI_DEVICE_PATH_PROTOCOL              *FullDevicePath;
  EFI_DEVICE_PATH_PROTOCOL              *TempDevicePath;
  UINTN                                 SataPort;
  BOOLEAN                               IsUsbDevice;
  BOOLEAN                               IsNvmeDevice;

  *VendorDefinedData = L05_UNKNOWN_DEVICE_TYPE;
  DevicePathNode     = DevicePathPtr;
  FullDevicePath     = NULL;
  TempDevicePath     = NULL;
  SataPort           = 0;
  IsUsbDevice        = FALSE;
  IsNvmeDevice       = FALSE;

  //
  // If device path is USB WWID, expand the USB short device path.
  //
  DevicePathPtr = (IS_USB_SHORT_FORM_DEVICE_PATH (DevicePathPtr)) ? BdsLibExpandUsbShortFormDevPath (DevicePathPtr) : DevicePathPtr;

  if (DevicePathPtr == NULL) {
    DevicePathPtr = DevicePathNode;
  }

  DevicePathNode = DevicePathPtr;

  //
  // By L05 Minimum BIOS Spec 1.37, 13 Lenovo UEFI Drive Identification
  // Add NVME device path to the Lenovo UEFI driver identify FilePathList.
  //
  if ((DevicePathType (DevicePathNode) == MEDIA_DEVICE_PATH) &&
      (DevicePathSubType (DevicePathNode) == MEDIA_HARDDRIVE_DP)) {
    //
    // If it's Device Path that starts with a hard drive path, append it with the front part to compose a
    // full device path to check device type.
    //
    FullDevicePath = BdsExpandPartitionPartialDevicePathToFull ((HARDDRIVE_DEVICE_PATH *) DevicePathNode);
    TempDevicePath = (FullDevicePath == NULL) ? DevicePathNode : FullDevicePath;

    if (TempDevicePath != NULL) {
      while (!IsDevicePathEnd (TempDevicePath)) {

        if ((DevicePathType (TempDevicePath) == MESSAGING_DEVICE_PATH) &&
            (DevicePathSubType (TempDevicePath) == MSG_NVME_NAMESPACE_DP)) {
          //
          //Check NVMe devcie
          //
          IsNvmeDevice = TRUE;
          break;
        }

        TempDevicePath = NextDevicePathNode (TempDevicePath);
      }

    }
  }

  if (FullDevicePath != NULL) {
    gBS->FreePool (FullDevicePath);
  }

  //
  //  Check the device path to find which type it is
  //
  while (!IsDevicePathEnd (DevicePathNode)) {
    if ((DevicePathType (DevicePathNode) == MESSAGING_DEVICE_PATH) &&
        (DevicePathSubType (DevicePathNode) == MSG_USB_DP)) {
      //
      //Check USB devcie
      //
      IsUsbDevice = TRUE;

    } else if ((DevicePathType (DevicePathNode) == MESSAGING_DEVICE_PATH) &&
               (DevicePathSubType (DevicePathNode) == MSG_SATA_DP)) {
      SataPort = ((SATA_DEVICE_PATH *) DevicePathNode)->HBAPortNumber;

    } else if (((DevicePathType (DevicePathNode) == MESSAGING_DEVICE_PATH) &&
                (DevicePathSubType (DevicePathNode) == MSG_IPv4_DP)) ||
               ((DevicePathType (DevicePathNode) == MESSAGING_DEVICE_PATH) &&
                (DevicePathSubType (DevicePathNode) == MSG_IPv6_DP))
               ) {
      //
      //Check network device
      //
      *VendorDefinedData = L05_NETWORK_DEVICE_TYPE;
      break;

    } else if ((DevicePathType (DevicePathNode) == MEDIA_DEVICE_PATH) &&
               (DevicePathSubType (DevicePathNode) == MEDIA_HARDDRIVE_DP)) {
      //
      //Check HardDisk
      //
      if (IsUsbDevice) {
        //
        //  Removable device type
        //
        *VendorDefinedData = L05_REMOVABLE_DEVICE_TYPE;

      } else if (IsNvmeDevice) {
        //
        // NVMe device type
        //
        *VendorDefinedData = L05_NVME_DEVICE_TYPE;

      } else {
        //
        // Fixed device type
        //
        *VendorDefinedData = L05_FIXED_DEVICE_TYPE;
      }

      break;

    } else if ((DevicePathType (DevicePathNode) == MEDIA_DEVICE_PATH) &&
               (DevicePathSubType (DevicePathNode) == MEDIA_CDROM_DP)) {
      if (IsUsbDevice) {
        //
        //USB CDROM
        //
        *VendorDefinedData = L05_REMOVABLE_DEVICE_TYPE;

      } else {
        //
        //CDROM - Optical device
        //
        *VendorDefinedData = L05_OPTICAL_DEVICE_TYPE;
      }

      break;

    } else if ((DevicePathType (DevicePathNode) == ACPI_DEVICE_PATH) &&
               (EISA_ID_TO_NUM (((ACPI_HID_DEVICE_PATH *) DevicePathNode)->HID) == 0x0604)) {
      if (IsUsbDevice) {
        //
        //USB Floppy
        //
        *VendorDefinedData = L05_REMOVABLE_DEVICE_TYPE;

      } else {
        *VendorDefinedData = L05_FIXED_DEVICE_TYPE;
      }

      break;
    }

    DevicePathNode = NextDevicePathNode (DevicePathNode);
  }

  return EFI_SUCCESS;
}

EFI_STATUS
L05SetLbootVaraible (
  IN  UINT16                            BootOptionNo
  )
{
  EFI_STATUS                            Status;
  UINT8                                 *Buffer;
  UINTN                                 BufferSize;
  L05_BOOT_OPTION                       *BootOptionPtr;
  UINT8                                 VendorDefinedData;
  CHAR16                                *DeviceName;
  CHAR16                                BootString[10];
  L05_VENDOR_DEVICE_PATH_WITH_DATA      *L05DevicePath;
  EFI_DEVICE_PATH_PROTOCOL              *OrginalDevicePathPtr;
  EFI_DEVICE_PATH_PROTOCOL              *DevicePathPtr;
  EFI_DEVICE_PATH_PROTOCOL              *NewL05DevicePath;
  EFI_DEVICE_PATH_PROTOCOL              *AddedL05DevicePath;
  L05_BOOT_OPTION                       *L05BootOption;

  Status                = EFI_SUCCESS;
  Buffer                = NULL;
  BufferSize            = 0;
  BootOptionPtr         = NULL;
  OrginalDevicePathPtr  = NULL;
  DevicePathPtr         = NULL;
  NewL05DevicePath      = NULL;
  AddedL05DevicePath    = NULL;
  L05BootOption         = NULL;

  CopyMem (BootString, L"Boot", 4 * sizeof (CHAR16));
  UnicodeValueToStringS (&(BootString[4]), sizeof (BootString) - (4 * sizeof (CHAR16)), PREFIX_ZERO, BootOptionNo, 4);

  Buffer = GetVariableAndSize (
             BootString,
             &gEfiGlobalVariableGuid,
             &BufferSize
             );

  if (BufferSize == 0) {
    return EFI_NOT_FOUND;
  }

  BootOptionPtr = (L05_BOOT_OPTION *) Buffer;

  DeviceName = (VOID *) (&(BootOptionPtr->Description));
  DevicePathPtr = (EFI_DEVICE_PATH_PROTOCOL *) ((UINT8 *) (&(BootOptionPtr->Description)) + StrSize (DeviceName));
  OrginalDevicePathPtr = (EFI_DEVICE_PATH_PROTOCOL *) Buffer;
  L05VendorDefinedData (DevicePathPtr, &VendorDefinedData);

  //
  //Create new l05 device path
  //
  L05DevicePath = (L05_VENDOR_DEVICE_PATH_WITH_DATA *) CreateDeviceNode (
                    MESSAGING_DEVICE_PATH,
                    MSG_VENDOR_DP,
                    (UINT16) sizeof (L05_VENDOR_DEVICE_PATH_WITH_DATA)
                    );
  //
  //set data to l05 device path (FilePathList[0])
  //
  L05DevicePath->Guid                 = gEfiL05LenovoBootOptionGuid;
  L05DevicePath->VendorDefinedData[0] = VendorDefinedData;
  //
  //Set first node to link list
  //
  AddedL05DevicePath                  = AppendDevicePathNode (NULL, (EFI_DEVICE_PATH_PROTOCOL *) L05DevicePath);
  //
  //copy orginal Boot#### device path to l05 devcie path (from FilePathList[1])
  //
  NewL05DevicePath                    = AppendDevicePath (AddedL05DevicePath, DevicePathPtr);

  //
  //Allocate memory for new boot option for LBoot
  //
  L05BootOption = AllocateZeroPool (BufferSize + (GetDevicePathSize (NewL05DevicePath) - GetDevicePathSize (DevicePathPtr)));
  //
  //Set Attributes is same as boot#### data
  //
  L05BootOption->Attributes = BootOptionPtr->Attributes;
  //
  //Get new device path length and set it to L05 boot option
  //
  L05BootOption->FilePathListLength = (UINT16) GetDevicePathSize (NewL05DevicePath);
  //
  //Copy Description to L05 boot option, it same as boot####
  //
  CopyMem ((VOID *) &(L05BootOption->Description), (VOID *) &(BootOptionPtr->Description), StrSize (DeviceName));
  //
  //Copy new device path list to L05 boot option
  //
  CopyMem (((UINT8 *) (&(L05BootOption->Description)) + StrSize (DeviceName)), NewL05DevicePath, L05BootOption->FilePathListLength);
  //
  //Copy option data, it same as boot####
  //
  CopyMem (((UINT8 *) (&(L05BootOption->Description)) + StrSize (DeviceName) + L05BootOption->FilePathListLength),
           ((UINT8 *) (&(BootOptionPtr->Description)) + StrSize (DeviceName) + BootOptionPtr->FilePathListLength),
           BufferSize - (sizeof (L05BootOption->Attributes) + sizeof (L05BootOption->FilePathListLength) + StrSize (DeviceName) + BootOptionPtr->FilePathListLength));


  CopyMem (BootString, L"lBoot", 5 * sizeof (CHAR16));
  UnicodeValueToStringS (&(BootString[5]), sizeof (BootString) - (5 * sizeof (CHAR16)), PREFIX_ZERO, BootOptionNo, 4);
  Status = gRT->SetVariable (
                  BootString,
                  &gEfiL05LenovoBootOptionGuid,
                  EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_RUNTIME_ACCESS,
                  BufferSize + sizeof (L05_VENDOR_DEVICE_PATH_WITH_DATA),
                  L05BootOption
                  );


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

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

  return EFI_SUCCESS;
}

VOID
ConvertBootToLBoot (
  IN EFI_EVENT                          Event,
  IN VOID                               *Context
  )
{
  UINT16                                BootDeviceCount;
  UINT16                                BootDeviceNum;
  UINT16                                *BootOrderList;
  UINTN                                 BufferSize;

  BootDeviceCount = 0;
  BootDeviceNum   = 0;
  BootOrderList   = NULL;
  BufferSize      = 0;

  BootOrderList = GetVariableAndSize (
                    L"BootOrder",
                    &gEfiGlobalVariableGuid,
                    &BufferSize
                    );

  if (BufferSize != 0) {
    BootDeviceCount = (UINT16) (BufferSize / sizeof (UINT16));
  }

  for (BootDeviceNum = 0; BootDeviceNum < BootDeviceCount; BootDeviceNum++) {
    L05SetLbootVaraible (BootDeviceNum);
  }
}