[PATCH v4 3/4] MdeModulePkg/CapsuleApp: Add functions to support Capsule-on-Disk

[Chen A Chen <chen.a.chen@intel.com>]

BZ:https://bugzilla.tianocore.org/show_bug.cgi?id=1482

This file provide some basic function to support Capsule-on-Disk.

Cc: Jian J Wang <jian.j.wang@intel.com>
Cc: Hao Wu <hao.a.wu@intel.com>
Cc: Zhang Chao B <chao.b.zhang@intel.com>
Contributed-under: TianoCore Contribution Agreement 1.1
Signed-off-by: Chen A Chen <chen.a.chen@intel.com>
Reviewed-by: Jian J Wang <jian.j.wang@intel.com>
---
 .../Application/CapsuleApp/CapsuleOnDisk.c         | 808 +++++++++++++++++++++
 1 file changed, 808 insertions(+)
 create mode 100644 MdeModulePkg/Application/CapsuleApp/CapsuleOnDisk.c

diff --git a/MdeModulePkg/Application/CapsuleApp/CapsuleOnDisk.c b/MdeModulePkg/Application/CapsuleApp/CapsuleOnDisk.c
new file mode 100644
index 0000000000..393b7ae7db
--- /dev/null
+++ b/MdeModulePkg/Application/CapsuleApp/CapsuleOnDisk.c
@@ -0,0 +1,808 @@
+/** @file
+  Process Capsule On Disk.
+
+  Copyright (c) 2019, Intel Corporation. All rights reserved.<BR>
+  This program and the accompanying materials
+  are licensed and made available under the terms and conditions of the BSD License
+  which accompanies this distribution.  The full text of the license may be found at
+  http://opensource.org/licenses/bsd-license.php
+
+  THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS,
+  WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
+
+**/
+#include <Uefi.h>
+#include <Library/BaseLib.h>
+#include <Library/DebugLib.h>
+#include <Library/BaseMemoryLib.h>
+#include <Library/MemoryAllocationLib.h>
+#include <Library/UefiBootServicesTableLib.h>
+#include <Library/UefiRuntimeServicesTableLib.h>
+#include <Library/UefiLib.h>
+#include <Library/PrintLib.h>
+#include <Library/DevicePathLib.h>
+#include <Library/FileHandleLib.h>
+#include <Library/UefiBootManagerLib.h>
+#include <Protocol/SimpleFileSystem.h>
+#include <Protocol/Shell.h>
+#include <Guid/FileInfo.h>
+#include <Guid/GlobalVariable.h>
+#include <Guid/Gpt.h>
+
+EFI_GUID mCapsuleOnDiskBootOptionGuid = { 0x4CC29BB7, 0x2413, 0x40A2, { 0xB0, 0x6D, 0x25, 0x3E, 0x37, 0x10, 0xF5, 0x32 } };
+
+/**
+  Get shell protocol.
+
+  @return Pointer to shell protocol.
+
+**/
+EFI_SHELL_PROTOCOL *
+GetShellProtocol (
+  VOID
+  );
+
+/**
+  Get file name from file path.
+
+  @param  FilePath    File path.
+
+  @return Pointer to file name.
+
+**/
+CHAR16 *
+GetFileNameFromPath (
+  CHAR16                            *FilePath
+  )
+{
+  EFI_STATUS                        Status;
+  EFI_SHELL_PROTOCOL                *ShellProtocol;
+  SHELL_FILE_HANDLE                 Handle;
+  EFI_FILE_INFO                     *FileInfo;
+
+  ShellProtocol = GetShellProtocol ();
+  if (ShellProtocol == NULL) {
+    return NULL;
+  }
+
+  //
+  // Open file by FileName.
+  //
+  Status = ShellProtocol->OpenFileByName (
+                            FilePath,
+                            &Handle,
+                            EFI_FILE_MODE_READ
+                            );
+  if (EFI_ERROR (Status)) {
+    return NULL;
+  }
+
+  //
+  // Get file name from EFI_FILE_INFO.
+  //
+  FileInfo = ShellProtocol->GetFileInfo (Handle);
+  ShellProtocol->CloseFile (Handle);
+  if (FileInfo == NULL) {
+    return NULL;
+  }
+
+  return FileInfo->FileName;
+}
+
+/**
+  Check if the device path is EFI system Partition.
+
+  @param  DevicePath    The ESP device path.
+
+  @retval TRUE    DevicePath is a device path for ESP.
+  @retval FALSE   DevicePath is not a device path for ESP.
+
+**/
+BOOLEAN
+IsEfiSysPartitionDevicePath (
+  EFI_DEVICE_PATH_PROTOCOL   *DevicePath
+  )
+{
+  EFI_STATUS                 Status;
+  EFI_DEVICE_PATH_PROTOCOL   *TempDevicePath;
+  HARDDRIVE_DEVICE_PATH      *Hd;
+  EFI_HANDLE                 Handle;
+
+  //
+  // Check if the device path contains GPT node
+  //
+  TempDevicePath = DevicePath;
+
+  while (!IsDevicePathEnd (TempDevicePath)) {
+    if ((DevicePathType (TempDevicePath) == MEDIA_DEVICE_PATH) &&
+      (DevicePathSubType (TempDevicePath) == MEDIA_HARDDRIVE_DP)) {
+      Hd = (HARDDRIVE_DEVICE_PATH *)TempDevicePath;
+      if (Hd->MBRType == MBR_TYPE_EFI_PARTITION_TABLE_HEADER) {
+        break;
+      }
+    }
+    TempDevicePath = NextDevicePathNode (TempDevicePath);
+  }
+
+  if (!IsDevicePathEnd (TempDevicePath)) {
+    //
+    // Search for EFI system partition protocol on full device path in Boot Option
+    //
+    Status = gBS->LocateDevicePath (&gEfiPartTypeSystemPartGuid, &DevicePath, &Handle);
+    return EFI_ERROR (Status) ? FALSE : TRUE;
+  } else {
+    return FALSE;
+  }
+}
+
+/**
+  Dump all EFI System Partition.
+
+**/
+VOID
+DumpAllEfiSysPartition (
+  VOID
+  )
+{
+  EFI_HANDLE                 *SimpleFileSystemHandles;
+  UINTN                      NumberSimpleFileSystemHandles;
+  UINTN                      Index;
+  EFI_DEVICE_PATH_PROTOCOL   *DevicePath;
+  UINTN                      NumberEfiSystemPartitions;
+  EFI_SHELL_PROTOCOL         *ShellProtocol;
+
+  ShellProtocol = GetShellProtocol ();
+  NumberEfiSystemPartitions = 0;
+
+  Print (L"EFI System Partition list:\n");
+
+  gBS->LocateHandleBuffer (
+         ByProtocol,
+         &gEfiSimpleFileSystemProtocolGuid,
+         NULL,
+         &NumberSimpleFileSystemHandles,
+         &SimpleFileSystemHandles
+         );
+
+  for (Index = 0; Index < NumberSimpleFileSystemHandles; Index++) {
+    DevicePath = DevicePathFromHandle (SimpleFileSystemHandles[Index]);
+    if (IsEfiSysPartitionDevicePath (DevicePath)) {
+      NumberEfiSystemPartitions++;
+      Print(L"    %s\n        %s\n", ShellProtocol->GetMapFromDevicePath (&DevicePath), ConvertDevicePathToText (DevicePath, TRUE, TRUE));
+    }
+  }
+
+  if (NumberEfiSystemPartitions == 0) {
+    Print(L"    No ESP found.\n");
+  }
+}
+
+/**
+  Check if capsule is provisioned.
+
+  @retval TRUE    Capsule is provisioned previously.
+  @retval FALSE   No capsule is provisioned.
+
+**/
+BOOLEAN
+IsCapsuleProvisioned (
+  VOID
+  )
+{
+  EFI_STATUS            Status;
+  UINT64                OsIndication;
+  UINTN                 DataSize;
+
+  OsIndication = 0;
+  DataSize = sizeof(UINT64);
+  Status = gRT->GetVariable (
+                  L"OsIndications",
+                  &gEfiGlobalVariableGuid,
+                  NULL,
+                  &DataSize,
+                  &OsIndication
+                  );
+  if (!EFI_ERROR (Status) &&
+      (OsIndication & EFI_OS_INDICATIONS_FILE_CAPSULE_DELIVERY_SUPPORTED) != 0) {
+    return TRUE;
+  }
+
+  return FALSE;
+}
+
+/**
+  Get one active Efi System Partition.
+
+  @param[out] FsDevicePath   The device path of Fs
+  @param[out] Fs             The file system within EfiSysPartition
+
+  @retval EFI_SUCCESS     Get file system successfully
+  @retval EFI_NOT_FOUND   No valid file system found
+
+**/
+EFI_STATUS
+GetEfiSysPartition (
+  OUT EFI_DEVICE_PATH_PROTOCOL         **FsDevicePath,
+  OUT EFI_SIMPLE_FILE_SYSTEM_PROTOCOL  **Fs
+  )
+{
+  EFI_HANDLE                 *SimpleFileSystemHandles;
+  UINTN                      NumberSimpleFileSystemHandles;
+  UINTN                      Index;
+  EFI_DEVICE_PATH_PROTOCOL   *DevicePath;
+  EFI_STATUS                 Status;
+
+  Status = gBS->LocateHandleBuffer (
+                  ByProtocol,
+                  &gEfiSimpleFileSystemProtocolGuid,
+                  NULL,
+                  &NumberSimpleFileSystemHandles,
+                  &SimpleFileSystemHandles
+                  );
+
+  if (EFI_ERROR (Status)) {
+    return EFI_NOT_FOUND;
+  }
+
+  for (Index = 0; Index < NumberSimpleFileSystemHandles; Index++) {
+    DevicePath = DevicePathFromHandle (SimpleFileSystemHandles[Index]);
+    if (IsEfiSysPartitionDevicePath (DevicePath)) {
+      Status = gBS->HandleProtocol (SimpleFileSystemHandles[Index], &gEfiSimpleFileSystemProtocolGuid, (VOID **)Fs);
+      if (!EFI_ERROR (Status)) {
+        *FsDevicePath = DevicePath;
+        return EFI_SUCCESS;
+      }
+    }
+  }
+
+  return EFI_NOT_FOUND;
+}
+
+/**
+  Check if Active Efi System Partition within GPT is in the device path.
+
+  @param[in]  DevicePath     The device path
+  @param[out] FsDevicePath   The device path of Fs
+  @param[out] Fs             The file system within EfiSysPartition
+
+  @retval EFI_SUCCESS    Get file system successfully
+  @retval EFI_NOT_FOUND  No valid file system found
+  @retval others         Get file system failed
+
+**/
+EFI_STATUS
+GetEfiSysPartitionFromDevPath (
+  IN  EFI_DEVICE_PATH_PROTOCOL        *DevicePath,
+  OUT EFI_DEVICE_PATH_PROTOCOL        **FsDevicePath,
+  OUT EFI_SIMPLE_FILE_SYSTEM_PROTOCOL **Fs
+  )
+{
+  EFI_STATUS                 Status;
+  EFI_DEVICE_PATH_PROTOCOL   *TempDevicePath;
+  HARDDRIVE_DEVICE_PATH      *Hd;
+  EFI_HANDLE                 Handle;
+
+  //
+  // Check if the device path contains GPT node
+  //
+  TempDevicePath = DevicePath;
+  while (!IsDevicePathEnd (TempDevicePath)) {
+    if ((DevicePathType (TempDevicePath) == MEDIA_DEVICE_PATH) &&
+       (DevicePathSubType (TempDevicePath) == MEDIA_HARDDRIVE_DP)) {
+      Hd = (HARDDRIVE_DEVICE_PATH *)TempDevicePath;
+      if (Hd->MBRType == MBR_TYPE_EFI_PARTITION_TABLE_HEADER) {
+        break;
+      }
+    }
+    TempDevicePath = NextDevicePathNode (TempDevicePath);
+  }
+
+  if (!IsDevicePathEnd (TempDevicePath)) {
+    //
+    // Search for EFI system partition protocol on full device path in Boot Option
+    //
+    Status = gBS->LocateDevicePath (&gEfiPartTypeSystemPartGuid, &DevicePath, &Handle);
+
+    //
+    // Search for simple file system on this handler
+    //
+    if (!EFI_ERROR (Status)) {
+      Status = gBS->HandleProtocol (Handle, &gEfiSimpleFileSystemProtocolGuid, (VOID **)Fs);
+      if (!EFI_ERROR (Status)) {
+        *FsDevicePath = DevicePathFromHandle (Handle);
+        return EFI_SUCCESS;
+      }
+    }
+  }
+
+  return EFI_NOT_FOUND;
+}
+
+/**
+  Get SimpleFileSystem from boot option file path.
+
+  @param[in]  DevicePath     The file path of boot option
+  @param[out] FullPath       The full device path of boot device
+  @param[out] Fs             The file system within EfiSysPartition
+
+  @retval EFI_SUCCESS    Get file system successfully
+  @retval EFI_NOT_FOUND  No valid file system found
+  @retval others         Get file system failed
+
+**/
+EFI_STATUS
+EFIAPI
+GetEfiSysPartitionFromBootOptionFilePath (
+  IN  EFI_DEVICE_PATH_PROTOCOL         *DevicePath,
+  OUT EFI_DEVICE_PATH_PROTOCOL         **FullPath,
+  OUT EFI_SIMPLE_FILE_SYSTEM_PROTOCOL  **Fs
+  )
+{
+  EFI_STATUS                        Status;
+  EFI_DEVICE_PATH_PROTOCOL          *CurFullPath;
+  EFI_DEVICE_PATH_PROTOCOL          *PreFullPath;
+  EFI_DEVICE_PATH_PROTOCOL          *FsFullPath;
+
+  CurFullPath = NULL;
+  FsFullPath = NULL;
+  //
+  // Try every full device Path generated from bootoption
+  //
+  do {
+    PreFullPath = CurFullPath;
+    CurFullPath = EfiBootManagerGetNextFullDevicePath (DevicePath, CurFullPath);
+
+    if (PreFullPath != NULL) {
+      FreePool (PreFullPath);
+    }
+
+    if (CurFullPath == NULL) {
+      //
+      // No Active EFI system partition is found in BootOption device path
+      //
+      Status = EFI_NOT_FOUND;
+      break;
+    }
+
+    DEBUG_CODE (
+      CHAR16 *DevicePathStr;
+
+      DevicePathStr = ConvertDevicePathToText (CurFullPath, TRUE, TRUE);
+      if (DevicePathStr != NULL){
+        DEBUG ((DEBUG_INFO, "Full device path %s\n", DevicePathStr));
+        FreePool (DevicePathStr);
+      }
+    );
+
+    Status = GetEfiSysPartitionFromDevPath (CurFullPath, &FsFullPath, Fs);
+  } while (EFI_ERROR (Status));
+
+  if (*Fs != NULL) {
+    *FullPath = FsFullPath;
+    return EFI_SUCCESS;
+  } else {
+    return EFI_NOT_FOUND;
+  }
+}
+
+/**
+  Get a valid SimpleFileSystem within EFI system partition.
+
+  @param[in]  Map             The FS mapping capsule write to
+  @param[out] BootNext        The value of BootNext Variable
+  @param[out] Fs              The file system within EfiSysPartition
+  @param[out] UpdateBootNext  The flag to indicate whether update BootNext Variable
+
+  @retval EFI_SUCCESS    Get FS successfully
+  @retval EFI_NOT_FOUND  No valid FS found
+  @retval others         Get FS failed
+
+**/
+EFI_STATUS
+EFIAPI
+GetUpdateFileSystem (
+  IN  CHAR16                           *Map,
+  OUT UINT16                           *BootNext,
+  OUT EFI_SIMPLE_FILE_SYSTEM_PROTOCOL  **Fs,
+  OUT BOOLEAN                          *UpdateBootNext
+)
+{
+  EFI_STATUS                      Status;
+  CHAR16                          BootOptionName[20];
+  UINTN                           Index;
+  CONST EFI_DEVICE_PATH_PROTOCOL  *MappedDevicePath;
+  EFI_DEVICE_PATH_PROTOCOL        *DevicePath;
+  EFI_DEVICE_PATH_PROTOCOL        *FullPath;
+  UINT16                          *BootNextData;
+  EFI_BOOT_MANAGER_LOAD_OPTION    BootNextOption;
+  EFI_BOOT_MANAGER_LOAD_OPTION    *BootOptionBuffer;
+  UINTN                           BootOptionCount;
+  EFI_SHELL_PROTOCOL              *ShellProtocol;
+  EFI_BOOT_MANAGER_LOAD_OPTION    NewOption;
+
+  MappedDevicePath = NULL;
+  ShellProtocol = GetShellProtocol ();
+
+  //
+  // 1. If Fs is not assigned and there are capsule provisioned before,
+  // Get EFI system partition from BootNext.
+  //
+  if (IsCapsuleProvisioned () && Map == NULL) {
+    Status = GetVariable2 (
+               L"BootNext",
+               &gEfiGlobalVariableGuid,
+               (VOID **)&BootNextData,
+               NULL
+               );
+    if (!EFI_ERROR (Status)) {
+      UnicodeSPrint (BootOptionName, sizeof (BootOptionName), L"Boot%04x", *BootNextData);
+      Status = EfiBootManagerVariableToLoadOption (BootOptionName, &BootNextOption);
+      if (!EFI_ERROR (Status)) {
+        DevicePath = BootNextOption.FilePath;
+        Status = GetEfiSysPartitionFromBootOptionFilePath (DevicePath, &FullPath, Fs);
+        if (!EFI_ERROR (Status)) {
+          *UpdateBootNext = FALSE;
+          Print(L"Get EFI system partition from BootNext : %s\n", BootNextOption.Description);
+          Print(L"%s %s\n", ShellProtocol->GetMapFromDevicePath (&FullPath), ConvertDevicePathToText (FullPath, TRUE, TRUE));
+          return EFI_SUCCESS;
+        }
+      }
+    }
+  }
+
+  //
+  // Check if Map is valid.
+  //
+  if (Map != NULL) {
+    MappedDevicePath = ShellProtocol->GetDevicePathFromMap (Map);
+    if (MappedDevicePath == NULL) {
+      Print(L"'%s' is not a valid mapping.\n", Map);
+      return EFI_INVALID_PARAMETER;
+    } else if (!IsEfiSysPartitionDevicePath (DuplicateDevicePath (MappedDevicePath))) {
+      Print(L"'%s' is not a EFI System Partition.\n", Map);
+      return EFI_INVALID_PARAMETER;
+    }
+  }
+
+  //
+  // 2. Get EFI system partition form boot options.
+  //
+  BootOptionBuffer = EfiBootManagerGetLoadOptions (&BootOptionCount, LoadOptionTypeBoot);
+  if (BootOptionCount == 0 && Map == NULL) {
+    return EFI_NOT_FOUND;
+  }
+
+  for (Index = 0; Index < BootOptionCount; Index++) {
+    //
+    // Get the boot option from the link list
+    //
+    DevicePath  = BootOptionBuffer[Index].FilePath;
+
+    //
+    // Skip inactive or legacy boot options
+    //
+    if ((BootOptionBuffer[Index].Attributes & LOAD_OPTION_ACTIVE) == 0 ||
+        DevicePathType (DevicePath) == BBS_DEVICE_PATH) {
+      continue;
+    }
+
+    DEBUG_CODE (
+      CHAR16 *DevicePathStr;
+
+      DevicePathStr = ConvertDevicePathToText (DevicePath, TRUE, TRUE);
+      if (DevicePathStr != NULL){
+        DEBUG ((DEBUG_INFO, "Try BootOption %s\n", DevicePathStr));
+        FreePool (DevicePathStr);
+      } else {
+        DEBUG ((DEBUG_INFO, "DevicePathToStr failed\n"));
+      }
+    );
+
+    Status = GetEfiSysPartitionFromBootOptionFilePath (DevicePath, &FullPath, Fs);
+    if (!EFI_ERROR (Status)) {
+      if (Map == NULL) {
+        *BootNext = (UINT16) BootOptionBuffer[Index].OptionNumber;
+        *UpdateBootNext = TRUE;
+        Print (L"Found EFI system partition on Boot%04x: %s\n", *BootNext, BootOptionBuffer[Index].Description);
+        Print (L"%s %s\n", ShellProtocol->GetMapFromDevicePath (&FullPath), ConvertDevicePathToText (FullPath, TRUE, TRUE));
+        return EFI_SUCCESS;
+      }
+
+      if (StrnCmp (Map, ShellProtocol->GetMapFromDevicePath (&FullPath), StrLen (Map)) == 0) {
+        *BootNext = (UINT16) BootOptionBuffer[Index].OptionNumber;
+        *UpdateBootNext = TRUE;
+        Print (L"Found Boot Option on %s : %s\n", Map, BootOptionBuffer[Index].Description);
+        return EFI_SUCCESS;
+      }
+    }
+  }
+
+  //
+  // 3. If no ESP is found on boot option, try to find a ESP and create boot option for it.
+  //
+  if (Map != NULL) {
+    //
+    // If map is assigned, try to get ESP from mapped Fs.
+    //
+    DevicePath = DuplicateDevicePath (MappedDevicePath);
+    Status = GetEfiSysPartitionFromDevPath (DevicePath, &FullPath, Fs);
+    if (EFI_ERROR (Status)) {
+      Print (L"Error: Cannot get EFI system partiion from '%s' - %r\n", Map, Status);
+      return EFI_NOT_FOUND;
+    }
+    Print (L"Warning: Cannot find Boot Option on '%s'!\n", Map);
+  } else {
+    Status = GetEfiSysPartition (&DevicePath, Fs);
+    if (EFI_ERROR (Status)) {
+      Print (L"Error: Cannot find a EFI system partition!\n");
+      return EFI_NOT_FOUND;
+    }
+  }
+
+  Print (L"Create Boot option for capsule on disk:\n");
+  Status = EfiBootManagerInitializeLoadOption (
+             &NewOption,
+             LoadOptionNumberUnassigned,
+             LoadOptionTypeBoot,
+             LOAD_OPTION_ACTIVE,
+             L"UEFI Capsule On Disk",
+             DevicePath,
+             (UINT8 *) &mCapsuleOnDiskBootOptionGuid,
+             sizeof(EFI_GUID)
+             );
+  if (!EFI_ERROR (Status)) {
+    Status = EfiBootManagerAddLoadOptionVariable (&NewOption, (UINTN) -1); {
+      if (!EFI_ERROR (Status)) {
+        *UpdateBootNext = TRUE;
+        *BootNext = (UINT16) NewOption.OptionNumber;
+        Print (L"  Boot%04x: %s\n", *BootNext, ConvertDevicePathToText(DevicePath, TRUE, TRUE));
+        return EFI_SUCCESS;
+      }
+    }
+  }
+
+  Print (L"ERROR: Cannot create boot option! - %r\n", Status);
+
+  return EFI_NOT_FOUND;
+}
+
+/**
+  Write files to a given SimpleFileSystem.
+
+  @param[in] Buffer          The buffer array
+  @param[in] BufferSize      The buffer size array
+  @param[in] FileName        The file name array
+  @param[in] BufferNum       The buffer number
+  @param[in] Fs              The SimpleFileSystem handle to be written
+
+  @retval EFI_SUCCESS    Write file successfully
+  @retval EFI_NOT_FOUND  SFS protocol not found
+  @retval others         Write file failed
+
+**/
+EFI_STATUS
+WriteUpdateFile (
+  IN  VOID                                 **Buffer,
+  IN  UINTN                                *BufferSize,
+  IN  CHAR16                               **FileName,
+  IN  UINTN                                BufferNum,
+  IN  EFI_SIMPLE_FILE_SYSTEM_PROTOCOL      *Fs
+)
+{
+  EFI_STATUS                          Status;
+  EFI_FILE                            *Root;
+  EFI_FILE                            *FileHandle;
+  EFI_FILE_PROTOCOL                   *DirHandle;
+  UINT64                              FileInfo;
+  VOID                                *Filebuffer;
+  UINTN                               FileSize;
+  UINTN                               Index;
+
+  DirHandle   = NULL;
+  FileHandle  = NULL;
+  Index       = 0;
+
+  //
+  // Open Root from SFS
+  //
+  Status = Fs->OpenVolume (Fs, &Root);
+  if (EFI_ERROR (Status)) {
+    Print (L"Cannot open volume. Status = %r\n", Status);
+    return EFI_NOT_FOUND;
+  }
+
+  //
+  // Ensure that efi and updatecapsule directories exist
+  //
+  Status = Root->Open (Root, &DirHandle, L"\\EFI", EFI_FILE_MODE_READ | EFI_FILE_MODE_WRITE, 0);
+  if (EFI_ERROR (Status)) {
+    Status = Root->Open (Root, &DirHandle, L"\\EFI", EFI_FILE_MODE_READ | EFI_FILE_MODE_WRITE | EFI_FILE_MODE_CREATE, EFI_FILE_DIRECTORY);
+    if (EFI_ERROR (Status)) {
+      Print(L"Unable to create %s directory\n", L"\\EFI");
+      return EFI_NOT_FOUND;
+    }
+  }
+  Status = Root->Open (Root, &DirHandle, EFI_CAPSULE_FILE_DIRECTORY, EFI_FILE_MODE_READ | EFI_FILE_MODE_WRITE , 0);
+  if (EFI_ERROR (Status)) {
+    Status = Root->Open (Root, &DirHandle, EFI_CAPSULE_FILE_DIRECTORY, EFI_FILE_MODE_READ | EFI_FILE_MODE_WRITE | EFI_FILE_MODE_CREATE, EFI_FILE_DIRECTORY);
+    if (EFI_ERROR (Status)) {
+      Print(L"Unable to create %s directory\n", EFI_CAPSULE_FILE_DIRECTORY);
+      return EFI_NOT_FOUND;
+    }
+  }
+
+  for (Index = 0; Index < BufferNum; Index++) {
+    FileHandle = NULL;
+
+    //
+    // Open UpdateCapsule file
+    //
+    Status = DirHandle->Open (DirHandle, &FileHandle, FileName[Index], EFI_FILE_MODE_CREATE | EFI_FILE_MODE_WRITE | EFI_FILE_MODE_READ, 0);
+    if (EFI_ERROR (Status)) {
+      Print (L"Unable to create %s file\n", FileName[Index]);
+      return EFI_NOT_FOUND;
+    }
+
+    //
+    // Empty the file contents
+    //
+    Status = FileHandleGetSize (FileHandle, &FileInfo);
+    if (EFI_ERROR (Status)) {
+      FileHandleClose (FileHandle);
+      Print (L"Error Reading %s\n", FileName[Index]);
+      return EFI_DEVICE_ERROR;
+    }
+
+    //
+    // If the file size is already 0, then it has been empty.
+    //
+    if (FileInfo != 0) {
+      //
+      // Set the file size to 0.
+      //
+      FileInfo = 0;
+      Status = FileHandleSetSize (FileHandle, FileInfo);
+      if (EFI_ERROR (Status)) {
+        Print (L"Error Deleting %s\n", FileName[Index]);
+        FileHandleClose (FileHandle);
+        return Status;
+      }
+    }
+
+    //
+    // Write Filebuffer to file
+    //
+    Filebuffer = Buffer[Index];
+    FileSize = BufferSize[Index];
+    Status = FileHandleWrite (FileHandle, &FileSize, Filebuffer);
+    if (EFI_ERROR (Status)) {
+      Print (L"Unable to write Capsule Update to %s, Status = %r\n", FileName[Index], Status);
+      return EFI_NOT_FOUND;
+    }
+
+    Print (L"Succeed to write %s\n", FileName[Index]);
+    FileHandleClose (FileHandle);
+  }
+
+  return EFI_SUCCESS;
+}
+
+/**
+  Set capsule status variable.
+
+  @param[in] SetCap     Set or clear the capsule flag.
+
+  @retval EFI_SUCCESS   Succeed to set SetCap variable.
+  @retval others        Fail to set the variable.
+
+**/
+EFI_STATUS
+SetCapsuleStatusVariable (
+  BOOLEAN                       SetCap
+  )
+{
+  EFI_STATUS                    Status;
+  UINT64                        OsIndication;
+  UINTN                         DataSize;
+
+  OsIndication = 0;
+  DataSize = sizeof(UINT64);
+  Status = gRT->GetVariable (
+                  L"OsIndications",
+                  &gEfiGlobalVariableGuid,
+                  NULL,
+                  &DataSize,
+                  &OsIndication
+                  );
+  if (EFI_ERROR (Status)) {
+    OsIndication = 0;
+  }
+  if (SetCap) {
+    OsIndication |= ((UINT64)EFI_OS_INDICATIONS_FILE_CAPSULE_DELIVERY_SUPPORTED);
+  }
+  else {
+    OsIndication &= ~((UINT64)EFI_OS_INDICATIONS_FILE_CAPSULE_DELIVERY_SUPPORTED);
+  }
+  Status = gRT->SetVariable (
+                  L"OsIndications",
+                  &gEfiGlobalVariableGuid,
+                  EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_RUNTIME_ACCESS | EFI_VARIABLE_NON_VOLATILE,
+                  sizeof(UINT64),
+                  &OsIndication
+                  );
+
+  return Status;
+}
+
+/**
+  Process Capsule On Disk.
+
+  @param[in]  CapsuleBuffer       An array of pointer to capsule images
+  @param[in]  CapsuleBufferSize   An array of UINTN to capsule images size
+  @param[in]  FilePath            An array of capsule images file path
+  @param[in]  Map                 File system mapping string
+  @param[in]  CapsuleNum          The count of capsule images
+
+  @retval EFI_SUCCESS       Capsule on disk success.
+  @retval others            Capsule on disk fail.
+
+**/
+EFI_STATUS
+ProcessCapsuleOnDisk (
+  IN VOID                          **CapsuleBuffer,
+  IN UINTN                         *CapsuleBufferSize,
+  IN CHAR16                        **FilePath,
+  IN CHAR16                        *Map,
+  IN UINTN                         CapsuleNum
+  )
+{
+  EFI_STATUS                      Status;
+  UINT16                          BootNext;
+  EFI_SIMPLE_FILE_SYSTEM_PROTOCOL *Fs;
+  BOOLEAN                         UpdateBootNext;
+
+  //
+  // Get a valid file system from boot path
+  //
+  Fs = NULL;
+
+  Status = GetUpdateFileSystem (Map, &BootNext, &Fs, &UpdateBootNext);
+  if (EFI_ERROR (Status)) {
+    Print (L"CapsuleApp: cannot find a valid file system on boot devies. Status = %r\n", Status);
+    return Status;
+  }
+
+  //
+  // Copy capsule image to '\efi\UpdateCapsule\'
+  //
+  Status = WriteUpdateFile (CapsuleBuffer, CapsuleBufferSize, FilePath, CapsuleNum, Fs);
+  if (EFI_ERROR (Status)) {
+    Print (L"CapsuleApp: capsule image could not be copied for update.\n");
+    return Status;
+  }
+
+  //
+  // Set variable then reset
+  //
+  Status = SetCapsuleStatusVariable (TRUE);
+  if (EFI_ERROR (Status)) {
+    Print (L"CapsuleApp: unable to set OSIndication variable.\n");
+    return Status;
+  }
+
+  if (UpdateBootNext) {
+    Status = gRT->SetVariable (
+      L"BootNext",
+      &gEfiGlobalVariableGuid,
+      EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_RUNTIME_ACCESS | EFI_VARIABLE_NON_VOLATILE,
+      sizeof(UINT16),
+      &BootNext
+      );
+    if (EFI_ERROR (Status)){
+      Print (L"CapsuleApp: unable to set BootNext variable.\n");
+      return Status;
+    }
+  }
+
+  return EFI_SUCCESS;
+}
-- 
2.16.2.windows.1