/** @file
  Boot functions implementation for UefiPxeBc Driver.

  Copyright (c) 2009 - 2018, Intel Corporation. All rights reserved.<BR>
  (C) Copyright 2016 Hewlett Packard Enterprise Development LP<BR>

  SPDX-License-Identifier: BSD-2-Clause-Patent

**/

#include "PxeBcImpl.h"


/**
  Display the string of the boot item.

  If the length of the boot item string beyond 70 Char, just display 70 Char.

  @param[in]  Str           The pointer to the string.
  @param[in]  Len           The length of the string.

**/
VOID
PxeBcDisplayBootItem (
  IN UINT8                 *Str,
  IN UINT8                 Len
  )
{
  UINT8                    Tmp;

  //
  // Cut off the chars behind 70th.
  //
  Len       = (UINT8) MIN (PXEBC_DISPLAY_MAX_LINE, Len);
  Tmp       = Str[Len];
  Str[Len]  = 0;
  AsciiPrint ("%a \n", Str);

  //
  // Restore the original 70th char.
  //
  Str[Len]  = Tmp;
}


/**
  Select and maintain the boot prompt if needed.

  @param[in]  Private          Pointer to PxeBc private data.

  @retval EFI_SUCCESS          Selected boot prompt done.
  @retval EFI_TIMEOUT          Selected boot prompt timed out.
  @retval EFI_NOT_FOUND        The proxy offer is not Pxe10.
  @retval EFI_ABORTED          User cancelled the operation.
  @retval EFI_NOT_READY        Reading the input key from the keyboard has not finish.

**/
EFI_STATUS
PxeBcSelectBootPrompt (
  IN PXEBC_PRIVATE_DATA      *Private
  )
{
  PXEBC_DHCP_PACKET_CACHE    *Cache;
  PXEBC_VENDOR_OPTION        *VendorOpt;
  EFI_PXE_BASE_CODE_MODE     *Mode;
  EFI_EVENT                  TimeoutEvent;
  EFI_EVENT                  DescendEvent;
  EFI_INPUT_KEY              InputKey;
  EFI_STATUS                 Status;
  UINT32                     OfferType;
  UINT8                      Timeout;
  UINT8                      *Prompt;
  UINT8                      PromptLen;
  INT32                      SecCol;
  INT32                      SecRow;

  TimeoutEvent = NULL;
  DescendEvent = NULL;
  Mode         = Private->PxeBc.Mode;
  Cache        = Mode->ProxyOfferReceived ? &Private->ProxyOffer : &Private->DhcpAck;
  OfferType    = Mode->UsingIpv6 ? Cache->Dhcp6.OfferType : Cache->Dhcp4.OfferType;

  //
  // Only DhcpPxe10 and ProxyPxe10 offer needs boot prompt.
  //
  if (OfferType != PxeOfferTypeProxyPxe10 && OfferType != PxeOfferTypeDhcpPxe10) {
    return EFI_NOT_FOUND;
  }

  //
  // There is no specified ProxyPxe10 for IPv6 in PXE and UEFI spec.
  //
  ASSERT (!Mode->UsingIpv6);

  VendorOpt = &Cache->Dhcp4.VendorOpt;
  //
  // According to the PXE specification 2.1, Table 2-1 PXE DHCP Options,
  // we must not consider a boot prompt or boot menu if all of the following hold:
  //   - the PXE_DISCOVERY_CONTROL tag(6) is present inside the Vendor Options(43), and has bit 3 set
  //   - a boot file name has been presented in the initial DHCP or ProxyDHCP offer packet.
  //
  if (IS_DISABLE_PROMPT_MENU (VendorOpt->DiscoverCtrl) &&
      Cache->Dhcp4.OptList[PXEBC_DHCP4_TAG_INDEX_BOOTFILE] != NULL) {
    return EFI_ABORTED;
  }

  if (!IS_VALID_BOOT_PROMPT (VendorOpt->BitMap)) {
    return EFI_TIMEOUT;
  }

  Timeout   = VendorOpt->MenuPrompt->Timeout;
  Prompt    = VendorOpt->MenuPrompt->Prompt;
  PromptLen = (UINT8) (VendorOpt->MenuPromptLen - 1);

  //
  // The valid scope of Timeout refers to PXE2.1 spec.
  //
  if (Timeout == 0) {
    return EFI_TIMEOUT;
  }
  if (Timeout == 255) {
    return EFI_SUCCESS;
  }

  //
  // Create and start a timer as timeout event.
  //
  Status = gBS->CreateEvent (
                  EVT_TIMER,
                  TPL_CALLBACK,
                  NULL,
                  NULL,
                  &TimeoutEvent
                  );
  if (EFI_ERROR (Status)) {
    return Status;
  }

  Status = gBS->SetTimer (
                  TimeoutEvent,
                  TimerRelative,
                  MultU64x32 (Timeout, TICKS_PER_SECOND)
                  );
  if (EFI_ERROR (Status)) {
    goto ON_EXIT;
  }

  //
  // Create and start a periodic timer as descend event by second.
  //
  Status = gBS->CreateEvent (
                  EVT_TIMER,
                  TPL_CALLBACK,
                  NULL,
                  NULL,
                  &DescendEvent
                  );
  if (EFI_ERROR (Status)) {
    goto ON_EXIT;
  }

  Status = gBS->SetTimer (
                  DescendEvent,
                  TimerPeriodic,
                  TICKS_PER_SECOND
                  );
  if (EFI_ERROR (Status)) {
    goto ON_EXIT;
  }

  //
  // Display the boot item and cursor on the screen.
  //
  SecCol = gST->ConOut->Mode->CursorColumn;
  SecRow = gST->ConOut->Mode->CursorRow;

  PxeBcDisplayBootItem (Prompt, PromptLen);

  gST->ConOut->SetCursorPosition (gST->ConOut, SecCol + PromptLen, SecRow);
  AsciiPrint ("(%d) ", Timeout--);

  Status = EFI_TIMEOUT;
  while (EFI_ERROR (gBS->CheckEvent (TimeoutEvent))) {
    if (!EFI_ERROR (gBS->CheckEvent (DescendEvent))) {
      gST->ConOut->SetCursorPosition (gST->ConOut, SecCol + PromptLen, SecRow);
      AsciiPrint ("(%d) ", Timeout--);
    }
    if (gST->ConIn->ReadKeyStroke (gST->ConIn, &InputKey) == EFI_NOT_READY) {
      gBS->Stall (10 * TICKS_PER_MS);
      continue;
    }
    //
    // Parse the input key by user.
    // If <F8> or <Ctrl> + <M> is pressed, return success to display the boot menu.
    //
    if (InputKey.ScanCode == 0) {

      switch (InputKey.UnicodeChar) {

      case CTRL ('c'):
        Status = EFI_ABORTED;
        break;

      case CTRL ('m'):
      case 'm':
      case 'M':
        Status = EFI_SUCCESS;
        break;

      default:
        continue;
      }

    } else {

      switch (InputKey.ScanCode) {

      case SCAN_F8:
        Status = EFI_SUCCESS;
        break;

      case SCAN_ESC:
        Status = EFI_ABORTED;
        break;

      default:
        continue;
      }
    }

    break;
  }

  //
  // Reset the cursor on the screen.
  //
  gST->ConOut->SetCursorPosition (gST->ConOut, 0 , SecRow + 1);

ON_EXIT:
  if (DescendEvent != NULL) {
    gBS->CloseEvent (DescendEvent);
  }
  if (TimeoutEvent != NULL) {
    gBS->CloseEvent (TimeoutEvent);
  }

  return Status;
}


/**
  Select the boot menu by user's input.

  @param[in]  Private         Pointer to PxeBc private data.
  @param[out] Type            The type of the menu.
  @param[in]  UseDefaultItem  Use default item or not.

  @retval EFI_ABORTED     User cancel operation.
  @retval EFI_SUCCESS     Select the boot menu success.
  @retval EFI_NOT_READY   Read the input key from the keybroad has not finish.

**/
EFI_STATUS
PxeBcSelectBootMenu (
  IN  PXEBC_PRIVATE_DATA              *Private,
  OUT UINT16                          *Type,
  IN  BOOLEAN                         UseDefaultItem
  )
{
  EFI_PXE_BASE_CODE_MODE     *Mode;
  PXEBC_DHCP_PACKET_CACHE    *Cache;
  PXEBC_VENDOR_OPTION        *VendorOpt;
  EFI_INPUT_KEY              InputKey;
  UINT32                     OfferType;
  UINT8                      MenuSize;
  UINT8                      MenuNum;
  INT32                      TopRow;
  UINT16                     Select;
  UINT16                     LastSelect;
  UINT8                      Index;
  BOOLEAN                    Finish;
  CHAR8                      Blank[PXEBC_DISPLAY_MAX_LINE];
  PXEBC_BOOT_MENU_ENTRY      *MenuItem;
  PXEBC_BOOT_MENU_ENTRY      *MenuArray[PXEBC_MENU_MAX_NUM];

  Finish    = FALSE;
  Select    = 0;
  Index     = 0;
  *Type     = 0;
  Mode      = Private->PxeBc.Mode;
  Cache     = Mode->ProxyOfferReceived ? &Private->ProxyOffer : &Private->DhcpAck;
  OfferType = Mode->UsingIpv6 ? Cache->Dhcp6.OfferType : Cache->Dhcp4.OfferType;

  //
  // There is no specified DhcpPxe10/ProxyPxe10 for IPv6 in PXE and UEFI spec.
  //
  ASSERT (!Mode->UsingIpv6);
  ASSERT (OfferType == PxeOfferTypeProxyPxe10 || OfferType == PxeOfferTypeDhcpPxe10);

  VendorOpt = &Cache->Dhcp4.VendorOpt;
  if (!IS_VALID_BOOT_MENU (VendorOpt->BitMap)) {
    return EFI_SUCCESS;
  }

  //
  // Display the boot menu on the screen.
  //
  SetMem (Blank, sizeof(Blank), ' ');

  MenuSize  = VendorOpt->BootMenuLen;
  MenuItem  = VendorOpt->BootMenu;

  if (MenuSize == 0) {
    return EFI_DEVICE_ERROR;
  }

  while (MenuSize > 0 && Index < PXEBC_MENU_MAX_NUM) {
    ASSERT (MenuItem != NULL);
    MenuArray[Index]  = MenuItem;
    MenuSize          = (UINT8) (MenuSize - (MenuItem->DescLen + 3));
    MenuItem          = (PXEBC_BOOT_MENU_ENTRY *) ((UINT8 *) MenuItem + MenuItem->DescLen + 3);
    Index++;
  }

  if (UseDefaultItem) {
    ASSERT (MenuArray[0] != NULL);
    CopyMem (Type, &MenuArray[0]->Type, sizeof (UINT16));
    *Type = NTOHS (*Type);
    return EFI_SUCCESS;
  }

  MenuNum = Index;

  for (Index = 0; Index < MenuNum; Index++) {
    ASSERT (MenuArray[Index] != NULL);
    PxeBcDisplayBootItem (MenuArray[Index]->DescStr, MenuArray[Index]->DescLen);
  }

  TopRow = gST->ConOut->Mode->CursorRow - MenuNum;

  //
  // Select the boot item by user in the boot menu.
  //
  do {
    //
    // Highlight selected row.
    //
    gST->ConOut->SetAttribute (gST->ConOut, EFI_TEXT_ATTR (EFI_BLACK, EFI_LIGHTGRAY));
    gST->ConOut->SetCursorPosition (gST->ConOut, 0, TopRow + Select);
    ASSERT (Select < PXEBC_MENU_MAX_NUM);
    ASSERT (MenuArray[Select] != NULL);
    Blank[MenuArray[Select]->DescLen] = 0;
    AsciiPrint ("%a\r", Blank);
    PxeBcDisplayBootItem (MenuArray[Select]->DescStr, MenuArray[Select]->DescLen);
    gST->ConOut->SetCursorPosition (gST->ConOut, 0, TopRow + MenuNum);
    LastSelect = Select;

    while (gST->ConIn->ReadKeyStroke (gST->ConIn, &InputKey) == EFI_NOT_READY) {
      gBS->Stall (10 * TICKS_PER_MS);
    }

    if (InputKey.ScanCode == 0) {
      switch (InputKey.UnicodeChar) {
      case CTRL ('c'):
        InputKey.ScanCode = SCAN_ESC;
        break;

      case CTRL ('j'):  /* linefeed */
      case CTRL ('m'):  /* return */
        Finish = TRUE;
        break;

      case CTRL ('i'):  /* tab */
      case ' ':
      case 'd':
      case 'D':
        InputKey.ScanCode = SCAN_DOWN;
        break;

      case CTRL ('h'):  /* backspace */
      case 'u':
      case 'U':
        InputKey.ScanCode = SCAN_UP;
        break;

      default:
        InputKey.ScanCode = 0;
      }
    }

    switch (InputKey.ScanCode) {
    case SCAN_LEFT:
    case SCAN_UP:
      if (Select != 0) {
        Select--;
      }
      break;

    case SCAN_DOWN:
    case SCAN_RIGHT:
      if (++Select == MenuNum) {
        Select--;
      }
      break;

    case SCAN_PAGE_UP:
    case SCAN_HOME:
      Select = 0;
      break;

    case SCAN_PAGE_DOWN:
    case SCAN_END:
      Select = (UINT16) (MenuNum - 1);
      break;

    case SCAN_ESC:
      return EFI_ABORTED;
    }

    //
    // Unhighlight the last selected row.
    //
    gST->ConOut->SetAttribute (gST->ConOut, EFI_TEXT_ATTR (EFI_LIGHTGRAY, EFI_BLACK));
    gST->ConOut->SetCursorPosition (gST->ConOut, 0, TopRow + LastSelect);
    ASSERT (LastSelect < PXEBC_MENU_MAX_NUM);
    ASSERT (MenuArray[LastSelect] != NULL);
    Blank[MenuArray[LastSelect]->DescLen] = 0;
    AsciiPrint ("%a\r", Blank);
    PxeBcDisplayBootItem (MenuArray[LastSelect]->DescStr, MenuArray[LastSelect]->DescLen);
    gST->ConOut->SetCursorPosition (gST->ConOut, 0, TopRow + MenuNum);
  } while (!Finish);

  //
  // Swap the byte order.
  //
  ASSERT (Select < PXEBC_MENU_MAX_NUM);
  ASSERT (MenuArray[Select] != NULL);
  CopyMem (Type, &MenuArray[Select]->Type, sizeof (UINT16));
  *Type = NTOHS (*Type);

  return EFI_SUCCESS;
}


/**
  Parse out the boot information from the last Dhcp4 reply packet.

  @param[in, out] Private      Pointer to PxeBc private data.
  @param[out]     BufferSize   Size of the boot file to be downloaded.

  @retval EFI_SUCCESS          Successfully parsed out all the boot information.
  @retval Others               Failed to parse out the boot information.

**/
EFI_STATUS
PxeBcDhcp4BootInfo (
  IN OUT PXEBC_PRIVATE_DATA   *Private,
     OUT UINT64               *BufferSize
  )
{
  EFI_PXE_BASE_CODE_PROTOCOL  *PxeBc;
  EFI_PXE_BASE_CODE_MODE      *Mode;
  EFI_STATUS                  Status;
  PXEBC_DHCP4_PACKET_CACHE    *Cache4;
  UINT16                      Value;
  PXEBC_VENDOR_OPTION         *VendorOpt;
  PXEBC_BOOT_SVR_ENTRY        *Entry;

  PxeBc       = &Private->PxeBc;
  Mode        = PxeBc->Mode;
  Status      = EFI_SUCCESS;
  *BufferSize = 0;

  //
  // Get the last received Dhcp4 reply packet.
  //
  if (Mode->PxeReplyReceived) {
    Cache4 = &Private->PxeReply.Dhcp4;
  } else if (Mode->ProxyOfferReceived) {
    Cache4 = &Private->ProxyOffer.Dhcp4;
  } else {
    Cache4 = &Private->DhcpAck.Dhcp4;
  }

  ASSERT (Cache4->OptList[PXEBC_DHCP4_TAG_INDEX_BOOTFILE] != NULL);

  //
  // Parse the boot server address.
  // If prompt/discover is disabled, get the first boot server from the boot servers list.
  // Otherwise, parse the boot server Ipv4 address from next server address field in DHCP header.
  // If all these fields are not available, use option 54 instead.
  //
  VendorOpt = &Cache4->VendorOpt;
  if (IS_DISABLE_PROMPT_MENU (VendorOpt->DiscoverCtrl) && IS_VALID_BOOT_SERVERS (VendorOpt->BitMap)) {
    Entry = VendorOpt->BootSvr;
    if (VendorOpt->BootSvrLen >= sizeof (PXEBC_BOOT_SVR_ENTRY) && Entry->IpCnt > 0) {
      CopyMem (
        &Private->ServerIp,
        &Entry->IpAddr[0],
        sizeof (EFI_IPv4_ADDRESS)
        );
    }
  }
  if (Private->ServerIp.Addr[0] == 0) {
    //
    // ServerIp.Addr[0] equals zero means we failed to get IP address from boot server list.
    // Try to use next server address field.
    //
    CopyMem (
      &Private->ServerIp,
      &Cache4->Packet.Offer.Dhcp4.Header.ServerAddr,
      sizeof (EFI_IPv4_ADDRESS)
      );
  }
  if (Private->ServerIp.Addr[0] == 0) {
    //
    // Still failed , use the IP address from option 54.
    //
    CopyMem (
      &Private->ServerIp,
      Cache4->OptList[PXEBC_DHCP4_TAG_INDEX_SERVER_ID]->Data,
      sizeof (EFI_IPv4_ADDRESS)
      );
  }

  //
  // Parse the boot file name by option.
  //
  Private->BootFileName = Cache4->OptList[PXEBC_DHCP4_TAG_INDEX_BOOTFILE]->Data;

  if (Cache4->OptList[PXEBC_DHCP4_TAG_INDEX_BOOTFILE_LEN] != NULL) {
    //
    // Parse the boot file size by option.
    //
    CopyMem (&Value, Cache4->OptList[PXEBC_DHCP4_TAG_INDEX_BOOTFILE_LEN]->Data, sizeof (Value));
    Value       = NTOHS (Value);
    //
    // The field of boot file size is 512 bytes in unit.
    //
    *BufferSize = 512 * Value;
  } else {
    //
    // Get the bootfile size by tftp command if no option available.
    //
    Status = PxeBc->Mtftp (
                      PxeBc,
                      EFI_PXE_BASE_CODE_TFTP_GET_FILE_SIZE,
                      NULL,
                      FALSE,
                      BufferSize,
                      &Private->BlockSize,
                      &Private->ServerIp,
                      Private->BootFileName,
                      NULL,
                      FALSE
                      );
  }

  //
  // Save the value of boot file size.
  //
  Private->BootFileSize = (UINTN) *BufferSize;

  //
  // Display all the information: boot server address, boot file name and boot file size.
  //
  AsciiPrint ("\n  Server IP address is ");
  PxeBcShowIp4Addr (&Private->ServerIp.v4);
  AsciiPrint ("\n  NBP filename is %a", Private->BootFileName);
  AsciiPrint ("\n  NBP filesize is %d Bytes", Private->BootFileSize);

  return Status;
}


/**
  Parse out the boot information from the last Dhcp6 reply packet.

  @param[in, out] Private      Pointer to PxeBc private data.
  @param[out]     BufferSize   Size of the boot file to be downloaded.

  @retval EFI_SUCCESS          Successfully parsed out all the boot information.
  @retval EFI_BUFFER_TOO_SMALL
  @retval Others               Failed to parse out the boot information.

**/
EFI_STATUS
PxeBcDhcp6BootInfo (
  IN OUT PXEBC_PRIVATE_DATA   *Private,
     OUT UINT64               *BufferSize
  )
{
  EFI_PXE_BASE_CODE_PROTOCOL  *PxeBc;
  EFI_PXE_BASE_CODE_MODE      *Mode;
  EFI_STATUS                  Status;
  PXEBC_DHCP6_PACKET_CACHE    *Cache6;
  UINT16                      Value;

  PxeBc       = &Private->PxeBc;
  Mode        = PxeBc->Mode;
  Status      = EFI_SUCCESS;
  *BufferSize = 0;

  //
  // Get the last received Dhcp6 reply packet.
  //
  if (Mode->PxeReplyReceived) {
    Cache6 = &Private->PxeReply.Dhcp6;
  } else if (Mode->ProxyOfferReceived) {
    Cache6 = &Private->ProxyOffer.Dhcp6;
  } else {
    Cache6 = &Private->DhcpAck.Dhcp6;
  }

  ASSERT (Cache6->OptList[PXEBC_DHCP6_IDX_BOOT_FILE_URL] != NULL);

  //
  // Set the station address to IP layer.
  //
  Status = PxeBcSetIp6Address (Private);
  if (EFI_ERROR (Status)) {
    return Status;
  }


  //
  // Parse (m)tftp server ip address and bootfile name.
  //
  Status = PxeBcExtractBootFileUrl (
             Private,
             &Private->BootFileName,
             &Private->ServerIp.v6,
             (CHAR8 *) (Cache6->OptList[PXEBC_DHCP6_IDX_BOOT_FILE_URL]->Data),
             NTOHS (Cache6->OptList[PXEBC_DHCP6_IDX_BOOT_FILE_URL]->OpLen)
             );
  if (EFI_ERROR (Status)) {
    return Status;
  }

  //
  // Parse the value of boot file size.
  //
  if (Cache6->OptList[PXEBC_DHCP6_IDX_BOOT_FILE_PARAM] != NULL) {
    //
    // Parse it out if have the boot file parameter option.
    //
    Status = PxeBcExtractBootFileParam ((CHAR8 *) Cache6->OptList[PXEBC_DHCP6_IDX_BOOT_FILE_PARAM]->Data, &Value);
    if (EFI_ERROR (Status)) {
      return Status;
    }
    //
    // The field of boot file size is 512 bytes in unit.
    //
    *BufferSize = 512 * Value;
  } else {
    //
    // Send get file size command by tftp if option unavailable.
    //
    Status = PxeBc->Mtftp (
                      PxeBc,
                      EFI_PXE_BASE_CODE_TFTP_GET_FILE_SIZE,
                      NULL,
                      FALSE,
                      BufferSize,
                      &Private->BlockSize,
                      &Private->ServerIp,
                      Private->BootFileName,
                      NULL,
                      FALSE
                      );
  }

  //
  // Save the value of boot file size.
  //
  Private->BootFileSize = (UINTN) *BufferSize;

  //
  // Display all the information: boot server address, boot file name and boot file size.
  //
  AsciiPrint ("\n  Server IP address is ");
  PxeBcShowIp6Addr (&Private->ServerIp.v6);
  AsciiPrint ("\n  NBP filename is %a", Private->BootFileName);
  AsciiPrint ("\n  NBP filesize is %d Bytes", Private->BootFileSize);

  return Status;
}


/**
  Extract the discover information and boot server entry from the
  cached packets if unspecified.

  @param[in]      Private      Pointer to PxeBc private data.
  @param[in]      Type         The type of bootstrap to perform.
  @param[in, out] DiscoverInfo Pointer to EFI_PXE_BASE_CODE_DISCOVER_INFO.
  @param[out]     BootEntry    Pointer to PXEBC_BOOT_SVR_ENTRY.
  @param[out]     SrvList      Pointer to EFI_PXE_BASE_CODE_SRVLIST.

  @retval EFI_SUCCESS       Successfully extracted the information.
  @retval EFI_DEVICE_ERROR  Failed to extract the information.

**/
EFI_STATUS
PxeBcExtractDiscoverInfo (
  IN     PXEBC_PRIVATE_DATA               *Private,
  IN     UINT16                           Type,
  IN OUT EFI_PXE_BASE_CODE_DISCOVER_INFO  **DiscoverInfo,
     OUT PXEBC_BOOT_SVR_ENTRY             **BootEntry,
     OUT EFI_PXE_BASE_CODE_SRVLIST        **SrvList
  )
{
  EFI_PXE_BASE_CODE_MODE          *Mode;
  PXEBC_DHCP4_PACKET_CACHE        *Cache4;
  PXEBC_VENDOR_OPTION             *VendorOpt;
  PXEBC_BOOT_SVR_ENTRY            *Entry;
  BOOLEAN                         IsFound;
  EFI_PXE_BASE_CODE_DISCOVER_INFO *Info;
  UINT16                          Index;

  Mode = Private->PxeBc.Mode;
  Info = *DiscoverInfo;

  if (Mode->UsingIpv6) {
    Info->IpCnt    = 1;
    Info->UseUCast = TRUE;

    Info->SrvList[0].Type              = Type;
    Info->SrvList[0].AcceptAnyResponse = FALSE;

    //
    // There is no vendor options specified in DHCPv6, so take BootFileUrl in the last cached packet.
    //
    CopyMem (&Info->SrvList[0].IpAddr, &Private->ServerIp, sizeof (EFI_IP_ADDRESS));

    *SrvList  = Info->SrvList;
  } else {
    Entry     = NULL;
    IsFound   = FALSE;
    Cache4    = (Mode->ProxyOfferReceived) ? &Private->ProxyOffer.Dhcp4 : &Private->DhcpAck.Dhcp4;
    VendorOpt = &Cache4->VendorOpt;

    if (!Mode->DhcpAckReceived || !IS_VALID_DISCOVER_VENDOR_OPTION (VendorOpt->BitMap)) {
      //
      // Address is not acquired or no discovery options.
      //
      return EFI_INVALID_PARAMETER;
    }

    //
    // Parse the boot server entry from the vendor option in the last cached packet.
    //
    Info->UseMCast    = (BOOLEAN) !IS_DISABLE_MCAST_DISCOVER (VendorOpt->DiscoverCtrl);
    Info->UseBCast    = (BOOLEAN) !IS_DISABLE_BCAST_DISCOVER (VendorOpt->DiscoverCtrl);
    Info->MustUseList = (BOOLEAN) IS_ENABLE_USE_SERVER_LIST (VendorOpt->DiscoverCtrl);
    Info->UseUCast    = (BOOLEAN) IS_VALID_BOOT_SERVERS (VendorOpt->BitMap);

    if (Info->UseMCast) {
      //
      // Get the multicast discover ip address from vendor option if has.
      //
      CopyMem (&Info->ServerMCastIp.v4, &VendorOpt->DiscoverMcastIp, sizeof (EFI_IPv4_ADDRESS));
    }

    Info->IpCnt = 0;

    if (Info->UseUCast) {
      Entry = VendorOpt->BootSvr;

      while (((UINT8) (Entry - VendorOpt->BootSvr)) < VendorOpt->BootSvrLen) {
        if (Entry->Type == HTONS (Type)) {
          IsFound = TRUE;
          break;
        }
        Entry = GET_NEXT_BOOT_SVR_ENTRY (Entry);
      }

      if (!IsFound) {
        return EFI_DEVICE_ERROR;
      }

      Info->IpCnt = Entry->IpCnt;
      if (Info->IpCnt >= 1) {
        *DiscoverInfo = AllocatePool (sizeof (*Info) + (Info->IpCnt - 1) * sizeof (**SrvList));
        if (*DiscoverInfo == NULL) {
          return EFI_OUT_OF_RESOURCES;
        }
        CopyMem (*DiscoverInfo, Info, sizeof (*Info));
        Info = *DiscoverInfo;
      }

      for (Index = 0; Index < Info->IpCnt; Index++) {
        CopyMem (&Info->SrvList[Index].IpAddr, &Entry->IpAddr[Index], sizeof (EFI_IPv4_ADDRESS));
        Info->SrvList[Index].AcceptAnyResponse = !Info->MustUseList;
        Info->SrvList[Index].Type = NTOHS (Entry->Type);
      }
    }

    *BootEntry = Entry;
    *SrvList   = Info->SrvList;
  }

  return EFI_SUCCESS;
}


/**
  Build the discover packet and send out for boot server.

  @param[in]  Private               Pointer to PxeBc private data.
  @param[in]  Type                  PxeBc option boot item type.
  @param[in]  Layer                 Pointer to option boot item layer.
  @param[in]  UseBis                Use BIS or not.
  @param[in]  DestIp                Pointer to the destination address.
  @param[in]  IpCount               The count of the server address.
  @param[in]  SrvList               Pointer to the server address list.

  @retval     EFI_SUCCESS           Successfully discovered boot file.
  @retval     EFI_OUT_OF_RESOURCES  Failed to allocate resource.
  @retval     EFI_NOT_FOUND         Can't get the PXE reply packet.
  @retval     Others                Failed to discover boot file.

**/
EFI_STATUS
PxeBcDiscoverBootServer (
  IN  PXEBC_PRIVATE_DATA                *Private,
  IN  UINT16                            Type,
  IN  UINT16                            *Layer,
  IN  BOOLEAN                           UseBis,
  IN  EFI_IP_ADDRESS                    *DestIp,
  IN  UINT16                            IpCount,
  IN  EFI_PXE_BASE_CODE_SRVLIST         *SrvList
  )
{
  if (Private->PxeBc.Mode->UsingIpv6) {
    return PxeBcDhcp6Discover (
             Private,
             Type,
             Layer,
             UseBis,
             DestIp
             );
  } else {
    return PxeBcDhcp4Discover (
             Private,
             Type,
             Layer,
             UseBis,
             DestIp,
             IpCount,
             SrvList
             );
  }
}


/**
  Discover all the boot information for boot file.

  @param[in, out] Private      Pointer to PxeBc private data.
  @param[out]     BufferSize   Size of the boot file to be downloaded.

  @retval EFI_SUCCESS          Successfully obtained all the boot information .
  @retval EFI_BUFFER_TOO_SMALL The buffer size is not enough for boot file.
  @retval EFI_ABORTED          User cancel current operation.
  @retval Others               Failed to parse out the boot information.

**/
EFI_STATUS
PxeBcDiscoverBootFile (
  IN OUT PXEBC_PRIVATE_DATA   *Private,
     OUT UINT64               *BufferSize
  )
{
  EFI_PXE_BASE_CODE_PROTOCOL  *PxeBc;
  EFI_PXE_BASE_CODE_MODE      *Mode;
  EFI_STATUS                  Status;
  UINT16                      Type;
  UINT16                      Layer;
  BOOLEAN                     UseBis;

  PxeBc = &Private->PxeBc;
  Mode  = PxeBc->Mode;
  Type  = EFI_PXE_BASE_CODE_BOOT_TYPE_BOOTSTRAP;
  Layer = EFI_PXE_BASE_CODE_BOOT_LAYER_INITIAL;

  //
  // Start D.O.R.A/S.A.R.R exchange to acquire station ip address and
  // other pxe boot information.
  //
  Status = PxeBc->Dhcp (PxeBc, TRUE);
  if (EFI_ERROR (Status)) {
    return Status;
  }

  //
  // Select a boot server from boot server list.
  //
  Status = PxeBcSelectBootPrompt (Private);

  if (Status == EFI_SUCCESS) {
    //
    // Choose by user's input.
    //
    Status = PxeBcSelectBootMenu (Private, &Type, FALSE);
  } else if (Status == EFI_TIMEOUT) {
    //
    // Choose by default item.
    //
    Status = PxeBcSelectBootMenu (Private, &Type, TRUE);
  }

  if (!EFI_ERROR (Status)) {

    if (Type == EFI_PXE_BASE_CODE_BOOT_TYPE_BOOTSTRAP) {
      //
      // Local boot(PXE bootstrap server) need abort
      //
      return EFI_ABORTED;
    }

    //
    // Start to discover the boot server to get (m)tftp server ip address, bootfile
    // name and bootfile size.
    //
    UseBis = (BOOLEAN) (Mode->BisSupported && Mode->BisDetected);
    Status = PxeBc->Discover (PxeBc, Type, &Layer, UseBis, NULL);
    if (EFI_ERROR (Status)) {
      return Status;
    }

    if (Mode->PxeReplyReceived && !Mode->ProxyOfferReceived) {
      //
      // Some network boot loader only search the packet in Mode.ProxyOffer to get its server
      // IP address, so we need to store a copy of Mode.PxeReply packet into Mode.ProxyOffer.
      //
      if (Mode->UsingIpv6) {
        CopyMem (
          &Mode->ProxyOffer.Dhcpv6,
          &Mode->PxeReply.Dhcpv6,
          Private->PxeReply.Dhcp6.Packet.Ack.Length
          );
      } else {
        CopyMem (
          &Mode->ProxyOffer.Dhcpv4,
          &Mode->PxeReply.Dhcpv4,
          Private->PxeReply.Dhcp4.Packet.Ack.Length
          );
      }
      Mode->ProxyOfferReceived = TRUE;
    }
  }

  //
  // Parse the boot information.
  //
  if (Mode->UsingIpv6) {
    Status = PxeBcDhcp6BootInfo (Private, BufferSize);
  } else {
    Status = PxeBcDhcp4BootInfo (Private, BufferSize);
  }

  return Status;
}


/**
  Install PxeBaseCodeCallbackProtocol if not installed before.

  @param[in, out] Private           Pointer to PxeBc private data.
  @param[out]     NewMakeCallback   If TRUE, it is a new callback.
                                    Otherwise, it is not new callback.
  @retval EFI_SUCCESS          PxeBaseCodeCallbackProtocol installed succesfully.
  @retval Others               Failed to install PxeBaseCodeCallbackProtocol.

**/
EFI_STATUS
PxeBcInstallCallback (
  IN OUT PXEBC_PRIVATE_DATA   *Private,
     OUT BOOLEAN              *NewMakeCallback
  )
{
  EFI_PXE_BASE_CODE_PROTOCOL  *PxeBc;
  EFI_STATUS                  Status;

  //
  // Check whether PxeBaseCodeCallbackProtocol already installed.
  //
  PxeBc  = &Private->PxeBc;
  Status = gBS->HandleProtocol (
                  Private->Mode.UsingIpv6 ? Private->Ip6Nic->Controller : Private->Ip4Nic->Controller,
                  &gEfiPxeBaseCodeCallbackProtocolGuid,
                  (VOID **) &Private->PxeBcCallback
                  );
  if (Status == EFI_UNSUPPORTED) {

    CopyMem (
      &Private->LoadFileCallback,
      &gPxeBcCallBackTemplate,
      sizeof (EFI_PXE_BASE_CODE_CALLBACK_PROTOCOL)
      );

    //
    // Install a default callback if user didn't offer one.
    //
    Status = gBS->InstallProtocolInterface (
                    Private->Mode.UsingIpv6 ? &Private->Ip6Nic->Controller : &Private->Ip4Nic->Controller,
                    &gEfiPxeBaseCodeCallbackProtocolGuid,
                    EFI_NATIVE_INTERFACE,
                    &Private->LoadFileCallback
                    );

    (*NewMakeCallback) = (BOOLEAN) (Status == EFI_SUCCESS);

    Status = PxeBc->SetParameters (PxeBc, NULL, NULL, NULL, NULL, NewMakeCallback);
    if (EFI_ERROR (Status)) {
      PxeBc->Stop (PxeBc);
      return Status;
    }
  }

  return EFI_SUCCESS;
}


/**
  Uninstall PxeBaseCodeCallbackProtocol.

  @param[in]  Private           Pointer to PxeBc private data.
  @param[in]  NewMakeCallback   If TRUE, it is a new callback.
                                Otherwise, it is not new callback.

**/
VOID
PxeBcUninstallCallback (
  IN PXEBC_PRIVATE_DATA        *Private,
  IN BOOLEAN                   NewMakeCallback
  )
{
  EFI_PXE_BASE_CODE_PROTOCOL   *PxeBc;

  PxeBc = &Private->PxeBc;

  if (NewMakeCallback) {

    NewMakeCallback = FALSE;

    PxeBc->SetParameters (PxeBc, NULL, NULL, NULL, NULL, &NewMakeCallback);

    gBS->UninstallProtocolInterface (
          Private->Mode.UsingIpv6 ? Private->Ip6Nic->Controller : Private->Ip4Nic->Controller,
          &gEfiPxeBaseCodeCallbackProtocolGuid,
          &Private->LoadFileCallback
          );
  }
}


/**
  Download one of boot file in the list, and it's special for IPv6.

  @param[in]      Private           Pointer to PxeBc private data.
  @param[in, out] BufferSize        Size of user buffer for input;
                                    required buffer size for output.
  @param[in]      Buffer            Pointer to user buffer.

  @retval EFI_SUCCESS               Read one of boot file in the list successfully.
  @retval EFI_BUFFER_TOO_SMALL      The buffer size is not enough for boot file.
  @retval EFI_NOT_FOUND             There is no proper boot file available.
  @retval Others                    Failed to download boot file in the list.

**/
EFI_STATUS
PxeBcReadBootFileList (
  IN     PXEBC_PRIVATE_DATA           *Private,
  IN OUT UINT64                       *BufferSize,
  IN     VOID                         *Buffer           OPTIONAL
  )
{
  EFI_STATUS                          Status;
  EFI_PXE_BASE_CODE_PROTOCOL          *PxeBc;

  PxeBc        = &Private->PxeBc;

  //
  // Try to download the boot file if everything is ready.
  //
  if (Buffer != NULL) {
    Status = PxeBc->Mtftp (
                      PxeBc,
                      EFI_PXE_BASE_CODE_TFTP_READ_FILE,
                      Buffer,
                      FALSE,
                      BufferSize,
                      &Private->BlockSize,
                      &Private->ServerIp,
                      Private->BootFileName,
                      NULL,
                      FALSE
                      );


  } else {
    Status      = EFI_BUFFER_TOO_SMALL;
  }

  return Status;
}


/**
  Load boot file into user buffer.

  @param[in]      Private           Pointer to PxeBc private data.
  @param[in, out] BufferSize        Size of user buffer for input;
                                    required buffer size for output.
  @param[in]      Buffer            Pointer to user buffer.

  @retval EFI_SUCCESS          Get all the boot information successfully.
  @retval EFI_BUFFER_TOO_SMALL The buffer size is not enough for boot file.
  @retval EFI_ABORTED          User cancelled the current operation.
  @retval Others               Failed to parse out the boot information.

**/
EFI_STATUS
PxeBcLoadBootFile (
  IN     PXEBC_PRIVATE_DATA           *Private,
  IN OUT UINTN                        *BufferSize,
  IN     VOID                         *Buffer         OPTIONAL
  )
{
  BOOLEAN                             NewMakeCallback;
  UINT64                              RequiredSize;
  UINT64                              CurrentSize;
  EFI_STATUS                          Status;
  EFI_PXE_BASE_CODE_PROTOCOL          *PxeBc;
  EFI_PXE_BASE_CODE_MODE              *PxeBcMode;

  NewMakeCallback = FALSE;
  PxeBc           = &Private->PxeBc;
  PxeBcMode       = &Private->Mode;
  CurrentSize     = *BufferSize;
  RequiredSize    = 0;

  //
  // Install pxebc callback protocol if hasn't been installed yet.
  //
  Status = PxeBcInstallCallback (Private, &NewMakeCallback);
  if (EFI_ERROR(Status)) {
    return Status;
  }

  if (Private->BootFileSize == 0) {
    //
    // Discover the boot information about the bootfile if hasn't.
    //
    Status = PxeBcDiscoverBootFile (Private, &RequiredSize);
    if (EFI_ERROR (Status)) {
      goto ON_EXIT;
    }

    if (PXEBC_IS_SIZE_OVERFLOWED (RequiredSize)) {
      //
      // It's error if the required buffer size is beyond the system scope.
      //
      Status = EFI_DEVICE_ERROR;
      goto ON_EXIT;
    } else if (RequiredSize > 0) {
      //
      // Get the right buffer size of the bootfile required.
      //
      if (CurrentSize < RequiredSize || Buffer == NULL) {
        //
        // It's buffer too small if the size of user buffer is smaller than the required.
        //
        CurrentSize = RequiredSize;
        Status      = EFI_BUFFER_TOO_SMALL;
        goto ON_EXIT;
      }
      CurrentSize = RequiredSize;
    } else if (RequiredSize == 0 && PxeBcMode->UsingIpv6) {
      //
      // Try to download another bootfile in list if failed to get the filesize of the last one.
      // It's special for the case of IPv6.
      //
      Status = PxeBcReadBootFileList (Private, &CurrentSize, Buffer);
      goto ON_EXIT;
    }
  } else if (CurrentSize < Private->BootFileSize || Buffer == NULL ) {
    //
    // It's buffer too small if the size of user buffer is smaller than the required.
    //
    CurrentSize = Private->BootFileSize;
    Status      = EFI_BUFFER_TOO_SMALL;
    goto ON_EXIT;
  }

  //
  // Begin to download the bootfile if everything is ready.
  //
  AsciiPrint ("\n Downloading NBP file...\n");
  if (PxeBcMode->UsingIpv6) {
    Status = PxeBcReadBootFileList (
               Private,
               &CurrentSize,
               Buffer
               );
  } else {
    Status = PxeBc->Mtftp (
                      PxeBc,
                      EFI_PXE_BASE_CODE_TFTP_READ_FILE,
                      Buffer,
                      FALSE,
                      &CurrentSize,
                      &Private->BlockSize,
                      &Private->ServerIp,
                      Private->BootFileName,
                      NULL,
                      FALSE
                      );
  }

ON_EXIT:
  *BufferSize = (UINTN) CurrentSize;
  PxeBcUninstallCallback(Private, NewMakeCallback);

  if (Status == EFI_SUCCESS) {
    AsciiPrint ("\n  NBP file downloaded successfully.\n");
    return EFI_SUCCESS;
  } else if (Status == EFI_BUFFER_TOO_SMALL && Buffer != NULL) {
    AsciiPrint ("\n  PXE-E05: Buffer size is smaller than the requested file.\n");
  } else if (Status == EFI_DEVICE_ERROR) {
    AsciiPrint ("\n  PXE-E07: Network device error.\n");
  } else if (Status == EFI_OUT_OF_RESOURCES) {
    AsciiPrint ("\n  PXE-E09: Could not allocate I/O buffers.\n");
  } else if (Status == EFI_NO_MEDIA) {
    AsciiPrint ("\n  PXE-E12: Could not detect network connection.\n");
  } else if (Status == EFI_NO_RESPONSE) {
    AsciiPrint ("\n  PXE-E16: No valid offer received.\n");
  } else if (Status == EFI_TIMEOUT) {
    AsciiPrint ("\n  PXE-E18: Server response timeout.\n");
  } else if (Status == EFI_ABORTED) {
    AsciiPrint ("\n  PXE-E21: Remote boot cancelled.\n");
  } else if (Status == EFI_ICMP_ERROR) {
    AsciiPrint ("\n  PXE-E22: Client received ICMP error from server.\n");
  } else if (Status == EFI_TFTP_ERROR) {
    AsciiPrint ("\n  PXE-E23: Client received TFTP error from server.\n");
  } else if (Status == EFI_NOT_FOUND) {
    AsciiPrint ("\n  PXE-E53: No boot filename received.\n");
  } else if (Status != EFI_BUFFER_TOO_SMALL) {
    AsciiPrint ("\n  PXE-E99: Unexpected network error.\n");
  }

  return Status;
}