MdeModulePkg/CapsuleApp: Enhance Capsule-On-Disk related functions.

[mirror_edk2.git] / MdeModulePkg / Application / CapsuleApp / AppSupport.c

/** @file\r
  A shell application that triggers capsule update process.\r
\r
  Copyright (c) 2016 - 2019, Intel Corporation. All rights reserved.<BR>\r
  SPDX-License-Identifier: BSD-2-Clause-Patent\r
\r
**/\r
\r
#include "CapsuleApp.h"\r
\r
UINTN  Argc;\r
CHAR16 **Argv;\r
EFI_SHELL_PROTOCOL      *mShellProtocol = NULL;\r
\r
/**\r
\r
  This function parse application ARG.\r
\r
  @return Status\r
**/\r
EFI_STATUS\r
GetArg (\r
  VOID\r
  )\r
{\r
  EFI_STATUS                    Status;\r
  EFI_SHELL_PARAMETERS_PROTOCOL *ShellParameters;\r
\r
  Status = gBS->HandleProtocol (\r
                  gImageHandle,\r
                  &gEfiShellParametersProtocolGuid,\r
                  (VOID**)&ShellParameters\r
                  );\r
  if (EFI_ERROR(Status)) {\r
    return Status;\r
  }\r
\r
  Argc = ShellParameters->Argc;\r
  Argv = ShellParameters->Argv;\r
  return EFI_SUCCESS;\r
}\r
\r
/**\r
  Get shell protocol.\r
\r
  @return Pointer to shell protocol.\r
**/\r
EFI_SHELL_PROTOCOL *\r
GetShellProtocol (\r
  VOID\r
  )\r
{\r
  EFI_STATUS            Status;\r
\r
  if (mShellProtocol == NULL) {\r
    Status = gBS->LocateProtocol (\r
                    &gEfiShellProtocolGuid,\r
                    NULL,\r
                    (VOID **) &mShellProtocol\r
                    );\r
    if (EFI_ERROR (Status)) {\r
      mShellProtocol = NULL;\r
    }\r
  }\r
\r
  return mShellProtocol;\r
}\r
\r
/**\r
  Read a file.\r
\r
  @param[in]  FileName        The file to be read.\r
  @param[out] BufferSize      The file buffer size\r
  @param[out] Buffer          The file buffer\r
\r
  @retval EFI_SUCCESS    Read file successfully\r
  @retval EFI_NOT_FOUND  Shell protocol or file not found\r
  @retval others         Read file failed\r
**/\r
EFI_STATUS\r
ReadFileToBuffer (\r
  IN  CHAR16                               *FileName,\r
  OUT UINTN                                *BufferSize,\r
  OUT VOID                                 **Buffer\r
  )\r
{\r
  EFI_STATUS                        Status;\r
  EFI_SHELL_PROTOCOL                *ShellProtocol;\r
  SHELL_FILE_HANDLE                 Handle;\r
  UINT64                            FileSize;\r
  UINTN                             TempBufferSize;\r
  VOID                              *TempBuffer;\r
\r
  ShellProtocol = GetShellProtocol();\r
  if (ShellProtocol == NULL) {\r
    return EFI_NOT_FOUND;\r
  }\r
\r
  //\r
  // Open file by FileName.\r
  //\r
  Status = ShellProtocol->OpenFileByName (\r
                            FileName,\r
                            &Handle,\r
                            EFI_FILE_MODE_READ\r
                            );\r
  if (EFI_ERROR (Status)) {\r
    return Status;\r
  }\r
\r
  //\r
  // Get the file size.\r
  //\r
  Status = ShellProtocol->GetFileSize (Handle, &FileSize);\r
  if (EFI_ERROR (Status)) {\r
    ShellProtocol->CloseFile (Handle);\r
    return Status;\r
  }\r
\r
  TempBufferSize = (UINTN) FileSize;\r
  TempBuffer = AllocateZeroPool (TempBufferSize);\r
  if (TempBuffer == NULL) {\r
    ShellProtocol->CloseFile (Handle);\r
    return EFI_OUT_OF_RESOURCES;\r
  }\r
\r
  //\r
  // Read the file data to the buffer\r
  //\r
  Status = ShellProtocol->ReadFile (\r
                            Handle,\r
                            &TempBufferSize,\r
                            TempBuffer\r
                            );\r
  if (EFI_ERROR (Status)) {\r
    ShellProtocol->CloseFile (Handle);\r
    return Status;\r
  }\r
\r
  ShellProtocol->CloseFile (Handle);\r
\r
  *BufferSize = TempBufferSize;\r
  *Buffer     = TempBuffer;\r
  return EFI_SUCCESS;\r
}\r
\r
/**\r
  Write a file.\r
\r
  @param[in] FileName        The file to be written.\r
  @param[in] BufferSize      The file buffer size\r
  @param[in] Buffer          The file buffer\r
\r
  @retval EFI_SUCCESS    Write file successfully\r
  @retval EFI_NOT_FOUND  Shell protocol not found\r
  @retval others         Write file failed\r
**/\r
EFI_STATUS\r
WriteFileFromBuffer (\r
  IN  CHAR16                               *FileName,\r
  IN  UINTN                                BufferSize,\r
  IN  VOID                                 *Buffer\r
  )\r
{\r
  EFI_STATUS                        Status;\r
  EFI_SHELL_PROTOCOL                *ShellProtocol;\r
  SHELL_FILE_HANDLE                 Handle;\r
  EFI_FILE_INFO                     *FileInfo;\r
  UINTN                             TempBufferSize;\r
\r
  ShellProtocol = GetShellProtocol();\r
  if (ShellProtocol == NULL) {\r
    return EFI_NOT_FOUND;\r
  }\r
\r
  //\r
  // Open file by FileName.\r
  //\r
  Status = ShellProtocol->OpenFileByName (\r
                            FileName,\r
                            &Handle,\r
                            EFI_FILE_MODE_READ | EFI_FILE_MODE_WRITE | EFI_FILE_MODE_CREATE\r
                            );\r
  if (EFI_ERROR (Status)) {\r
    return Status;\r
  }\r
\r
  //\r
  // Empty the file contents.\r
  //\r
  FileInfo = ShellProtocol->GetFileInfo (Handle);\r
  if (FileInfo == NULL) {\r
    ShellProtocol->CloseFile (Handle);\r
    return EFI_DEVICE_ERROR;\r
  }\r
\r
  //\r
  // If the file size is already 0, then it has been empty.\r
  //\r
  if (FileInfo->FileSize != 0) {\r
    //\r
    // Set the file size to 0.\r
    //\r
    FileInfo->FileSize = 0;\r
    Status = ShellProtocol->SetFileInfo (Handle, FileInfo);\r
    if (EFI_ERROR (Status)) {\r
      FreePool (FileInfo);\r
      ShellProtocol->CloseFile (Handle);\r
      return Status;\r
    }\r
  }\r
  FreePool (FileInfo);\r
\r
  //\r
  // Write the file data from the buffer\r
  //\r
  TempBufferSize = BufferSize;\r
  Status = ShellProtocol->WriteFile (\r
                            Handle,\r
                            &TempBufferSize,\r
                            Buffer\r
                            );\r
  if (EFI_ERROR (Status)) {\r
    ShellProtocol->CloseFile (Handle);\r
    return Status;\r
  }\r
\r
  ShellProtocol->CloseFile (Handle);\r
\r
  return EFI_SUCCESS;\r
}\r
\r