/** @file
  Cloud Boot DXE Driver.

;******************************************************************************
;* 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 "CloudBootDxe.h"

/**
  Check Network Type.

  @param  BootOptionNum                 Input boot option number.
  @param  *NetworkType                  Output the boot option number type by network type.

  @retval EFI_SUCCESS                   The operation completed successfully.
  @retval Others                        An unexpected error occurred.
**/
EFI_STATUS
CheckNetworkType (
  IN  UINT16                            BootOptionNum,
  OUT UINT8                             *NetworkType
  )
{
  EFI_STATUS                            Status;
  UINT16                                BootOptionName[BOOT_OPTION_MAX_CHAR];
  UINT8                                 *BootOption;
  UINTN                                 BootOptionSize;
  UINT8                                 *Ptr;
  EFI_DEVICE_PATH_PROTOCOL              *DevicePath;
  BOOLEAN                               IsWifiDevice;

  Status         = EFI_SUCCESS;
  BootOption     = NULL;
  BootOptionSize = 0;
  Ptr            = NULL;
  DevicePath     = NULL;
  IsWifiDevice   = FALSE;

  UnicodeSPrint (BootOptionName, sizeof (BootOptionName), L"Boot%04x", BootOptionNum);

  Status = CommonGetVariableDataAndSize (
             BootOptionName,
             &gEfiGlobalVariableGuid,
             &BootOptionSize,
             (VOID **) &BootOption
             );

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

  Ptr = BootOption;
  Ptr += sizeof (UINT32);
  Ptr += sizeof (UINT16);
  Ptr += StrSize ((CHAR16 *) Ptr);

  //
  // Now Ptr point to Device Path
  //
  DevicePath   = (EFI_DEVICE_PATH_PROTOCOL *) Ptr;
  *NetworkType = 0xFF;
  Status       = EFI_NOT_FOUND;

  while (!IsDevicePathEnd (DevicePath)) {

    if (DevicePath->Type == MESSAGING_DEVICE_PATH && DevicePath->SubType == MSG_IPv4_DP) {
      *NetworkType = (IsWifiDevice) ? L05_NETWORK_WIFI_IPV4 : L05_NETWORK_IPV4;
      Status = EFI_SUCCESS;

    } else if (DevicePath->Type == MESSAGING_DEVICE_PATH && DevicePath->SubType == MSG_IPv6_DP) {
      *NetworkType = (IsWifiDevice) ? L05_NETWORK_WIFI_IPV6 : L05_NETWORK_IPV6;;
      Status = EFI_SUCCESS;

    } else if (DevicePath->Type == MESSAGING_DEVICE_PATH && DevicePath->SubType == MSG_URI_DP) {
      if (IsWifiDevice) {
        *NetworkType = (*NetworkType == L05_NETWORK_WIFI_IPV4) ? L05_NETWORK_HTTP_WIFI_IPV4 : L05_NETWORK_HTTP_WIFI_IPV6;
      } else {
        *NetworkType = (*NetworkType == L05_NETWORK_IPV4) ? L05_NETWORK_HTTP_IPV4 : L05_NETWORK_HTTP_IPV6;
      }
      Status = EFI_SUCCESS;
      break;

    } else if (DevicePath->Type == MESSAGING_DEVICE_PATH && DevicePath->SubType == MSG_WIFI_DP) {
      IsWifiDevice = TRUE;

    }

    DevicePath = NextDevicePathNode (DevicePath);
  }

  FreePool (BootOption);

  return Status;
}

/**
  Switch boot order.

  @param  TempAllBootOrder              Point to temp boot order list.
  @param  Index                         Index number.
  @param  UpShift                       Move up or down.
**/
VOID
SwitchBootOrder (
  IN OUT TEMP_SAVE_NETWORK_INDEX        *TempAllBootOrder,
  IN     UINT8                          Index,
  IN     BOOLEAN                        UpShift
  )
{
  TEMP_SAVE_NETWORK_INDEX               SwapTempBootOrder;

  if (UpShift) {
    SwapTempBootOrder.BootOrderList = TempAllBootOrder[Index].BootOrderList;
    SwapTempBootOrder.NetwortType = TempAllBootOrder[Index].NetwortType;
    TempAllBootOrder[Index].BootOrderList = TempAllBootOrder[Index - 1].BootOrderList;
    TempAllBootOrder[Index].NetwortType = TempAllBootOrder[Index - 1].NetwortType;
    TempAllBootOrder[Index - 1].BootOrderList = SwapTempBootOrder.BootOrderList;
    TempAllBootOrder[Index - 1].NetwortType = SwapTempBootOrder.NetwortType;
  } else {
    SwapTempBootOrder.BootOrderList = TempAllBootOrder[Index].BootOrderList;
    SwapTempBootOrder.NetwortType = TempAllBootOrder[Index].NetwortType;
    TempAllBootOrder[Index].BootOrderList = TempAllBootOrder[Index + 1].BootOrderList;
    TempAllBootOrder[Index].NetwortType = TempAllBootOrder[Index + 1].NetwortType;
    TempAllBootOrder[Index + 1].BootOrderList = SwapTempBootOrder.BootOrderList;
    TempAllBootOrder[Index + 1].NetwortType = SwapTempBootOrder.NetwortType;
  }

  return;
}

/**
  Check whether the Wifi Cloud Service is disabled.

  @param  None.

  @retval TRUE                          Wifi Cloud Service is disabled.
  @retval FALSE                         Wifi Cloud Service is enabled.
**/
BOOLEAN
IsWifiCloudServiceDisable (
  VOID
  )
{
  EFI_STATUS                            Status;
  UINTN                                 BufferSize;
  SYSTEM_CONFIGURATION                  *SetupNvData;
  BOOLEAN                               IsDisable;

  SetupNvData = NULL;
  IsDisable   = FALSE;

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

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

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

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

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

  FreePool (SetupNvData);

  return IsDisable;
}

/**
  Sort Cloud Service boot order.
  When setup item "L05LenovoCloudServices" enable and "L05UefiWifiNetworkBoot" disable,
  and item is move without default position, than "L05UefiWifiNetworkBoot" enable
  the new boot item need combine position with "L05LenovoCloudServices" boot item,
  because they are combine in "Lenovo Cloud Services".

  @param  BootOrderList                 Point to boot order list.
  @param  BootOrderNumber               Number of boot options.

  @retval EFI_SUCCESS                   The operation completed successfully.
  @retval Others                        An unexpected error occurred.
**/
EFI_STATUS
SortCloudServiceOfBootOrder (
  IN OUT UINT16                         *BootOrderList,
  IN OUT UINTN                          BootOrderNumber
  )
{
  EFI_STATUS                            Status;
  UINT8                                 NetworkType;
  UINT8                                 NetworkType2;
  UINT8                                 CloudeServiceType;
  UINT8                                 CloudeServiceIndex;
  UINT8                                 WifiCloudeServiceType;
  UINT8                                 WifiCloudeServiceIndex;
  SYSTEM_CONFIGURATION                  *SystemConfiguration;
  UINT8                                 Index;
  TEMP_SAVE_NETWORK_INDEX               *TempAllBootOrder;
  UINT8                                 NeedShiftTime;

  SystemConfiguration   = NULL;
  TempAllBootOrder      = NULL;
  NetworkType           = 0;
  NetworkType2          = 0;
  NeedShiftTime         = L05_SHIFT_TIME_INIT;
  CloudeServiceType     = L05_NETWORK_TYPE_NULL;
  WifiCloudeServiceType = L05_NETWORK_TYPE_NULL;

  if (IsWifiCloudServiceDisable ()) {
    return EFI_UNSUPPORTED;
  }

  if ((BootOrderList == NULL) || (BootOrderNumber < EFI_L05_HTTP_WIFI_ITEM_COUNT_MAX)) {
    return EFI_INVALID_PARAMETER;
  }

  TempAllBootOrder = AllocateZeroPool (sizeof (TEMP_SAVE_NETWORK_INDEX) * BootOrderNumber);

  if (TempAllBootOrder == NULL) {
    return EFI_OUT_OF_RESOURCES;
  }

  for (Index = 0; Index < BootOrderNumber; Index++) {
    Status = CheckNetworkType (BootOrderList [Index], &NetworkType);
    TempAllBootOrder[Index].BootOrderList = BootOrderList[Index];
    TempAllBootOrder[Index].NetwortType = NetworkType;
  }

  for (Index = 0; Index < (BootOrderNumber - 1); Index++) {
    NetworkType = (TempAllBootOrder[Index].NetwortType & L05_NETWORK_TYPE_MASK);
    NetworkType2 = (TempAllBootOrder[Index + 1].NetwortType & L05_NETWORK_TYPE_MASK);

    if (((NetworkType == L05_NETWORK_HTTP_WIFI_MASK) && (NetworkType2 == L05_NETWORK_HTTP_MASK)) ||
        ((NetworkType == L05_NETWORK_HTTP_MASK) && (NetworkType2 == L05_NETWORK_HTTP_WIFI_MASK))) {
      FreePool (TempAllBootOrder);
      return EFI_ABORTED;
    }
  }

  for (Index = 0; Index < BootOrderNumber; Index++) {
    NetworkType = (TempAllBootOrder[Index].NetwortType & L05_NETWORK_TYPE_MASK);

    if (NetworkType == L05_NETWORK_HTTP_WIFI_MASK) {
      WifiCloudeServiceType  = NetworkType;
      WifiCloudeServiceIndex = Index;

    } else if (NetworkType == L05_NETWORK_HTTP_MASK) {
      CloudeServiceType  = NetworkType;
      CloudeServiceIndex = Index;

    }
  }

  if ((CloudeServiceType == L05_NETWORK_TYPE_NULL) || (WifiCloudeServiceType == L05_NETWORK_TYPE_NULL)) {
    FreePool (TempAllBootOrder);
    return EFI_ABORTED;
  }

  NeedShiftTime = L05_SHIFT_TIME_INIT;
  if (CloudeServiceIndex > WifiCloudeServiceIndex) {
    //
    // CloudeService need shift up
    //
    NeedShiftTime = CloudeServiceIndex - WifiCloudeServiceIndex - L05_NETWORK_HTTP_WIFI_ITEM_COUNT;
    if (NeedShiftTime > L05_SHIFT_TIME_INIT) {
      for (Index = 0; Index < NeedShiftTime; Index++) {
        SwitchBootOrder (TempAllBootOrder, (CloudeServiceIndex - Index), L05_SHIFT_UP);
      }
    }

  } else {
    //
    // WifiCloudeService need shift up
    //
    NeedShiftTime = WifiCloudeServiceIndex - CloudeServiceIndex - L05_NETWORK_HTTP_WIFI_ITEM_COUNT;
    if (NeedShiftTime > L05_SHIFT_TIME_INIT) {
      for (Index = 0; Index < NeedShiftTime; Index++) {
        SwitchBootOrder (TempAllBootOrder, (WifiCloudeServiceIndex - Index), L05_SHIFT_UP);
      }
    }
  }

  for (Index = 0; Index < BootOrderNumber; Index++) {
    BootOrderList[Index] = TempAllBootOrder[Index].BootOrderList;
  }

  FreePool (TempAllBootOrder);

  return EFI_SUCCESS;
}

/**
  Wifi Boot Variable Delete.

  @param  None.
**/
VOID
WifiBootVariableDelete (
  VOID
  )
{
  EFI_STATUS                            Status;
  UINT16                                *BootOrder;
  UINTN                                 BootOrderSize;
  UINTN                                 BootDevNum;
  UINTN                                 Index;
  UINT8                                 IndexNetwork;
  UINT8                                 NetworkTypeCount;
  CHAR16                                BootOptionName[BOOT_OPTION_MAX_CHAR];
  UINT8                                 *BootOption;
  UINTN                                 BootOptionSize;
  BOOLEAN                               DeleteFlag;
  UINT8                                 NetworkType;
  UINT8                                 NeedDeleteType[] = {
                                                            L05_NETWORK_HTTP_WIFI_IPV4,
                                                            L05_NETWORK_HTTP_WIFI_IPV6,
                                                            L05_NETWORK_WIFI_IPV4,
                                                            L05_NETWORK_WIFI_IPV6
                                                            };

  BootOrder        = NULL;
  BootOrderSize    = 0;
  BootDevNum       = 0;
  NetworkType      = 0;
  IndexNetwork     = 0;
  NetworkTypeCount = 0;
  DeleteFlag       = FALSE;

  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++) {
    DeleteFlag = FALSE;
    UnicodeSPrint (BootOptionName, sizeof (BootOptionName), L"Boot%04x", BootOrder[Index]);
    Status = CommonGetVariableDataAndSize (
               BootOptionName,
               &gEfiGlobalVariableGuid,
               &BootOptionSize,
               (VOID **) &BootOption
               );

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

    CheckNetworkType (BootOrder[Index], &NetworkType);

    NetworkTypeCount = sizeof (NeedDeleteType) / sizeof (UINT8);
    for (IndexNetwork = 0; IndexNetwork < NetworkTypeCount; IndexNetwork++) {
      if (NetworkType == NeedDeleteType[IndexNetwork]) {
        //
        // Delete WiFi Variable
        //
        DeleteFlag = TRUE;
      }
    }

    if (DeleteFlag) {
      Status = CommonSetVariable (
                 BootOptionName,
                 &gEfiGlobalVariableGuid,
                 0,
                 0,
                 NULL
                 );
    }

    FreePool (BootOption);
  }

  FreePool (BootOrder);

  return;
}

/**
  Contorl Http wired & wireless item active and combine some boot item in "Lenovo Cloud Services".

  @param  Event                         A pointer to the Event that triggered the callback.
  @param  Handle                        Checkpoint handle.
**/
VOID
EFIAPI
L05CloudServiceNotify (
  IN  EFI_EVENT                         Event,
  IN  H2O_CP_HANDLE                     Handle
  )
{
  EFI_STATUS                            Status;
  UINT16                                *BootOrder;
  UINTN                                 BootOrderSize;

  UINT16                                *BootList;
  UINTN                                 BootSize;

  Status   = EFI_NOT_FOUND;
  BootList = NULL;
  BootSize = 0;

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

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

  //
  // Combine some boot item in "Lenovo Cloud Services"
  //
  Status = SortCloudServiceOfBootOrder (BootOrder, BootOrderSize / sizeof (UINT16));

  if (!EFI_ERROR (Status)) {
    Status = CommonSetVariable (
               L"BootOrder",
               &gEfiGlobalVariableGuid,
               EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_RUNTIME_ACCESS | EFI_VARIABLE_NON_VOLATILE,
               BootOrderSize,
               BootOrder
               );
  }

  FreePool (BootOrder);

  return;
}

/**
  Delete wireless item in Variable when "L05UefiWifiNetworkBoot" set 0.

  @param  Event                         A pointer to the Event that triggered the callback.
  @param  Handle                        Checkpoint handle.
**/
VOID
EFIAPI
L05WifiBootVariableDeleteNotify (
  IN  EFI_EVENT                         Event,
  IN  H2O_CP_HANDLE                     Handle
  )
{
  H2OCpUnregisterHandler (Handle);

  if (IsWifiCloudServiceDisable ()) {
    WifiBootVariableDelete ();
  }

  return;
}

/**
  Wait Wifi connecting for Wifi Cloud Service.

  @retval EFI_SUCESS                    Connecting has been set successfully.
  @retval Others                        An unexpected error occurred.
**/
EFI_STATUS
EFIAPI
L05WaitWifiConnecting (
  VOID
  )
{
  EFI_STATUS                            Status;
  EFI_STATUS                            TimerStatus;
  EFI_EVENT                             Timer;
  UINT64                                TimeOutMs;

  TimerStatus = EFI_NOT_READY;
  Timer       = NULL;
  TimeOutMs   = 0;

  Status = gBS->CreateEvent (
                  EVT_TIMER,
                  TPL_CALLBACK,
                  NULL,
                  NULL,
                  &Timer
                  );

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

  Print (L"Wifi is connecting...");

  //
  // Calculate 30 sec
  //
  TimeOutMs = MultU64x32 ((UINT64) WIFI_WAITING_CONNECT_TIMEOUT_SEC, MILLISEC);
  Status = gBS->SetTimer (
                  Timer,
                  TimerRelative,
                  MultU64x32 (TimeOutMs, TIMER_EVENT_ONE_MILLISEC)
                  );

  if (EFI_ERROR (Status)) {
    gBS->CloseEvent(Timer);
    return EFI_DEVICE_ERROR;
  }

  //
  // Wait 30 sec , until connect sucess or time out
  //
  while (TimerStatus == EFI_NOT_READY) {
    TimerStatus = gBS->CheckEvent (Timer);
    if (PcdGetBool (PcdL05WifiAutoConnect)) {
      break;
    }
  }

  if (TimerStatus != EFI_NOT_READY) {
    DEBUG ((DEBUG_INFO, "Time Out\n"));
  }

  gBS->CloseEvent (Timer);
  gST->ConOut->ClearScreen (gST->ConOut);

  return EFI_SUCCESS;
}

/**
  Is WIFI search need execution when enter WIFI config and wifi support enable.

  @param  PackageListGuid               The current package list guid.

  @retval TRUE                          Wifi search need execution.
  @retval FALSE                         Wifi search don't execution.
**/
BOOLEAN
L05IsWifiSupportEnableAndEnterWifiConfig (
  IN CONST EFI_GUID                     *PackageListGuid
  )
{
  BOOLEAN                               IsWifiSupportEnable;
  EFI_GUID                              FormSetGuid = L05_WIFI_CONFIG_FORMSET_GUID;
  STATIC BOOLEAN                        ConnectDeviceOnce = FALSE;

  IsWifiSupportEnable = FALSE;

  if (ConnectDeviceOnce) {
    return FALSE;
  }

  if ((!(IsWifiCloudServiceDisable ())) && (CompareGuid (PackageListGuid, &FormSetGuid))) {
    IsWifiSupportEnable = TRUE;
    ConnectDeviceOnce = TRUE;
  }
  return IsWifiSupportEnable;
}

/**
  Cloud Boot DXE Entry.

  @param  ImageHandle                   The firmware allocated handle for the UEFI image.
  @param  SystemTable                   A pointer to the EFI System Table.

  @retval EFI_SUCCESS                   The operation completed successfully.
  @retval Others                        An unexpected error occurred.
**/
EFI_STATUS
L05CloudBootDxeEntryPoint (
  IN  EFI_HANDLE                        ImageHandle,
  IN  EFI_SYSTEM_TABLE                  *SystemTable
  )
{
  EFI_STATUS                            Status;
  EFI_L05_CLOUD_BOOT_PROTOCOL           *L05CloudBootPtr;
  H2O_CP_HANDLE                         CpHandle;

  Status          = EFI_SUCCESS;
  L05CloudBootPtr = NULL;
  CpHandle        = NULL;

  //
  // Install Protocol
  //
  L05CloudBootPtr = AllocateZeroPool (sizeof (EFI_L05_CLOUD_BOOT_PROTOCOL));

  if (L05CloudBootPtr == NULL) {
    return EFI_OUT_OF_RESOURCES;
  }

  L05CloudBootPtr->WaitWifiConnecting                    = L05WaitWifiConnecting;
  L05CloudBootPtr->IsWifiSupportEnableAndEnterWifiConfig = L05IsWifiSupportEnableAndEnterWifiConfig;

  Status = gBS->InstallProtocolInterface (
                  &ImageHandle,
                  &gEfiL05CloudBootProtocolGuid,
                  EFI_NATIVE_INTERFACE,
                  L05CloudBootPtr
                  );

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

  //
  // Need to be triggered after LimitedToInternalStorageNotify()
  //
  if (FeaturePcdGet (PcdH2OBdsCpBootDeviceEnumAfterSupported)) {
    Status = H2OCpRegisterHandler (
               &gH2OBdsCpBootDeviceEnumAfterGuid,
               L05CloudServiceNotify,
               H2O_CP_LOW,
               &CpHandle
               );

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

  if (FeaturePcdGet (PcdH2OBdsCpBootDeviceEnumBeforeSupported)) {
    Status = H2OCpRegisterHandler (
               &gH2OBdsCpBootDeviceEnumBeforeGuid,
               L05WifiBootVariableDeleteNotify,
               H2O_CP_LOW,
               &CpHandle
               );
    if (EFI_ERROR (Status)) {
      DEBUG_CP ((DEBUG_ERROR, "Checkpoint Register Fail: %g (%r)\n", &gH2OBdsCpBootDeviceEnumBeforeGuid, Status));
      return Status;
    }
    DEBUG_CP ((DEBUG_INFO, "Checkpoint Registered: %g (%r)\n", &gH2OBdsCpBootDeviceEnumBeforeGuid, Status));
  }

  return Status;
}