[PATCH] MdeModulePkg/CapsuleApp: Add Fmp->GetImage() support.

[Jiewen Yao <jiewen.yao@intel.com>]

We add Fmp->GetImage() support in CapsuleApp. So that user may call
Fmp->GetImage() in UEFI shell environment.
This is useful to do unit test for FMP which supports GetImage(),
or user wants to get current image, such as Microcode.

Cc: Eric Dong <eric.dong@intel.com>
Cc: Jeff Fan <jeff.fan@intel.com>
Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Jiewen Yao <jiewen.yao@intel.com>
---
 MdeModulePkg/Application/CapsuleApp/AppSupport.c  | 141 +++++++++++++
 MdeModulePkg/Application/CapsuleApp/CapsuleApp.c  |  57 +++++-
 MdeModulePkg/Application/CapsuleApp/CapsuleDump.c | 214 ++++++++++++++++++++
 3 files changed, 411 insertions(+), 1 deletion(-)

diff --git a/MdeModulePkg/Application/CapsuleApp/AppSupport.c b/MdeModulePkg/Application/CapsuleApp/AppSupport.c
index 0a1224d..6aea76a 100644
--- a/MdeModulePkg/Application/CapsuleApp/AppSupport.c
+++ b/MdeModulePkg/Application/CapsuleApp/AppSupport.c
@@ -27,6 +27,9 @@
 #include <Guid/FileInfo.h>
 #include <Guid/Gpt.h>
 
+#define IS_HYPHEN(a)               ((a) == L'-')
+#define IS_NULL(a)                 ((a) == L'\0')
+
 #define MAX_ARG_NUM     11
 
 UINTN  Argc;
@@ -61,6 +64,144 @@ GetArg (
 }
 
 /**
+  Converts a list of string to a specified buffer.
+
+  @param[out] Buf             The output buffer that contains the string.
+  @param[in]  BufferLength    The length of the buffer
+  @param[in]  Str             The input string that contains the hex number
+
+  @retval EFI_SUCCESS    The string was successfully converted to the buffer.
+
+**/
+EFI_STATUS
+StrToBuf (
+  OUT UINT8    *Buf,
+  IN  UINTN    BufferLength,
+  IN  CHAR16   *Str
+  )
+{
+  UINTN       Index;
+  UINTN       StrLength;
+  UINT8       Digit;
+  UINT8       Byte;
+
+  Digit = 0;
+
+  //
+  // Two hex char make up one byte
+  //
+  StrLength = BufferLength * sizeof (CHAR16);
+
+  for(Index = 0; Index < StrLength; Index++, Str++) {
+
+    if ((*Str >= L'a') && (*Str <= L'f')) {
+      Digit = (UINT8) (*Str - L'a' + 0x0A);
+    } else if ((*Str >= L'A') && (*Str <= L'F')) {
+      Digit = (UINT8) (*Str - L'A' + 0x0A);
+    } else if ((*Str >= L'0') && (*Str <= L'9')) {
+      Digit = (UINT8) (*Str - L'0');
+    } else {
+      return EFI_INVALID_PARAMETER;
+    }
+
+    //
+    // For odd characters, write the upper nibble for each buffer byte,
+    // and for even characters, the lower nibble.
+    //
+    if ((Index & 1) == 0) {
+      Byte = (UINT8) (Digit << 4);
+    } else {
+      Byte = Buf[Index / 2];
+      Byte &= 0xF0;
+      Byte = (UINT8) (Byte | Digit);
+    }
+
+    Buf[Index / 2] = Byte;
+  }
+
+  return EFI_SUCCESS;
+}
+
+/**
+  Converts a string to GUID value.
+  Guid Format is xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx
+
+  @param[in]  Str              The registry format GUID string that contains the GUID value.
+  @param[out] Guid             A pointer to the converted GUID value.
+
+  @retval EFI_SUCCESS     The GUID string was successfully converted to the GUID value.
+  @retval EFI_UNSUPPORTED The input string is not in registry format.
+  @return others          Some error occurred when converting part of GUID value.
+
+**/
+EFI_STATUS
+StrToGuid (
+  IN  CHAR16   *Str,
+  OUT EFI_GUID *Guid
+  )
+{
+  //
+  // Get the first UINT32 data
+  //
+  Guid->Data1 = (UINT32) StrHexToUint64  (Str);
+  while (!IS_HYPHEN (*Str) && !IS_NULL (*Str)) {
+    Str ++;
+  }
+
+  if (IS_HYPHEN (*Str)) {
+    Str++;
+  } else {
+    return EFI_UNSUPPORTED;
+  }
+
+  //
+  // Get the second UINT16 data
+  //
+  Guid->Data2 = (UINT16) StrHexToUint64  (Str);
+  while (!IS_HYPHEN (*Str) && !IS_NULL (*Str)) {
+    Str ++;
+  }
+
+  if (IS_HYPHEN (*Str)) {
+    Str++;
+  } else {
+    return EFI_UNSUPPORTED;
+  }
+
+  //
+  // Get the third UINT16 data
+  //
+  Guid->Data3 = (UINT16) StrHexToUint64  (Str);
+  while (!IS_HYPHEN (*Str) && !IS_NULL (*Str)) {
+    Str ++;
+  }
+
+  if (IS_HYPHEN (*Str)) {
+    Str++;
+  } else {
+    return EFI_UNSUPPORTED;
+  }
+
+  //
+  // Get the following 8 bytes data
+  //
+  StrToBuf (&Guid->Data4[0], 2, Str);
+  //
+  // Skip 2 byte hex chars
+  //
+  Str += 2 * 2;
+
+  if (IS_HYPHEN (*Str)) {
+    Str++;
+  } else {
+    return EFI_UNSUPPORTED;
+  }
+  StrToBuf (&Guid->Data4[2], 6, Str);
+
+  return EFI_SUCCESS;
+}
+
+/**
   Return File System Volume containing this shell application.
 
   @return File System Volume containing this shell application.
diff --git a/MdeModulePkg/Application/CapsuleApp/CapsuleApp.c b/MdeModulePkg/Application/CapsuleApp/CapsuleApp.c
index 23672ae..5137259 100644
--- a/MdeModulePkg/Application/CapsuleApp/CapsuleApp.c
+++ b/MdeModulePkg/Application/CapsuleApp/CapsuleApp.c
@@ -86,6 +86,22 @@ DumpFmpData (
   );
 
 /**
+  Dump FMP image data.
+
+  @param[in]  ImageTypeId   The ImageTypeId of the FMP image.
+                            It is used to identify the FMP protocol.
+  @param[in]  ImageIndex    The ImageIndex of the FMP image.
+                            It is the input parameter for FMP->GetImage().
+  @param[in]  ImageName     The file name to hold the output FMP image.
+**/
+VOID
+DumpFmpImage (
+  IN EFI_GUID  *ImageTypeId,
+  IN UINTN     ImageIndex,
+  IN CHAR16    *ImageName
+  );
+
+/**
   Dump ESRT info.
 **/
 VOID
@@ -127,6 +143,24 @@ WriteFileFromBuffer (
   );
 
 /**
+  Converts a string to GUID value.
+  Guid Format is xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx
+
+  @param[in]  Str              The registry format GUID string that contains the GUID value.
+  @param[out] Guid             A pointer to the converted GUID value.
+
+  @retval EFI_SUCCESS     The GUID string was successfully converted to the GUID value.
+  @retval EFI_UNSUPPORTED The input string is not in registry format.
+  @return others          Some error occurred when converting part of GUID value.
+
+**/
+EFI_STATUS
+StrToGuid (
+  IN  CHAR16   *Str,
+  OUT EFI_GUID *Guid
+  );
+
+/**
 
   This function parse application ARG.
 
@@ -662,6 +696,7 @@ PrintUsage (
   Print(L"  CapsuleApp -G <BMP> -O <Capsule>\n");
   Print(L"  CapsuleApp -N <Capsule> -O <NestedCapsule>\n");
   Print(L"  CapsuleApp -D <Capsule>\n");
+  Print(L"  CapsuleApp -P GET <ImageTypeId> <Index> -O <FileName>\n");
   Print(L"Parameter:\n");
   Print(L"  -S:  Dump capsule report variable (EFI_CAPSULE_REPORT_GUID),\n");
   Print(L"       which is defined in UEFI specification.\n");
@@ -737,7 +772,27 @@ UefiMain (
     return Status;
   }
   if (StrCmp(Argv[1], L"-P") == 0) {
-    DumpFmpData();
+    if (Argc == 2) {
+      DumpFmpData();
+    }
+    if (Argc >= 3) {
+      if (StrCmp(Argv[2], L"GET") == 0) {
+        EFI_GUID  ImageTypeId;
+        UINTN     ImageIndex;
+        //
+        // FMP->GetImage()
+        //
+        Status = StrToGuid(Argv[3], &ImageTypeId);
+        if (EFI_ERROR(Status)) {
+          Print (L"Invalid ImageTypeId - %s\n", Argv[3]);
+          return Status;
+        }
+        ImageIndex = StrDecimalToUintn(Argv[4]);
+        if (StrCmp(Argv[5], L"-O") == 0) {
+          DumpFmpImage(&ImageTypeId, ImageIndex, Argv[6]);
+        }
+      }
+    }
     return EFI_SUCCESS;
   }
   if (StrCmp(Argv[1], L"-E") == 0) {
diff --git a/MdeModulePkg/Application/CapsuleApp/CapsuleDump.c b/MdeModulePkg/Application/CapsuleApp/CapsuleDump.c
index d09b938..3d83ec4 100644
--- a/MdeModulePkg/Application/CapsuleApp/CapsuleDump.c
+++ b/MdeModulePkg/Application/CapsuleApp/CapsuleDump.c
@@ -46,6 +46,22 @@ ReadFileToBuffer (
   );
 
 /**
+  Write a file.
+
+  @param[in] FileName        The file to be written.
+  @param[in] BufferSize      The file buffer size
+  @param[in] Buffer          The file buffer
+
+  @retval EFI_SUCCESS    Write file successfully
+**/
+EFI_STATUS
+WriteFileFromBuffer (
+  IN  CHAR16                               *FileName,
+  IN  UINTN                                BufferSize,
+  IN  VOID                                 *Buffer
+  );
+
+/**
   Dump UX capsule information.
 
   @param[in] CapsuleHeader      The UX capsule header
@@ -738,3 +754,201 @@ DumpFmpData (
 EXIT:
   FreePool(HandleBuffer);
 }
+
+/**
+  Check if the ImageInfo includes the ImageTypeId.
+
+  @param[in] ImageInfo           A pointer to EFI_FIRMWARE_IMAGE_DESCRIPTOR.
+  @param[in] DescriptorCount     The count of EFI_FIRMWARE_IMAGE_DESCRIPTOR.
+  @param[in] DescriptorSize      The size of an individual EFI_FIRMWARE_IMAGE_DESCRIPTOR, in bytes.
+  @param[in] ImageTypeId         A unique GUID identifying the firmware image type.
+
+  @return TRUE  This ImageInfo includes the ImageTypeId
+  @return FALSE This ImageInfo does not include the ImageTypeId
+**/
+BOOLEAN
+IsThisFmpImageInfo (
+  IN EFI_FIRMWARE_IMAGE_DESCRIPTOR   *ImageInfo,
+  IN UINT8                           DescriptorCount,
+  IN UINTN                           DescriptorSize,
+  IN EFI_GUID                        *ImageTypeId
+  )
+{
+  EFI_FIRMWARE_IMAGE_DESCRIPTOR                 *CurrentImageInfo;
+  UINTN                                         Index;
+
+  CurrentImageInfo = ImageInfo;
+  for (Index = 0; Index < DescriptorCount; Index++) {
+    if (CompareGuid (&CurrentImageInfo->ImageTypeId, ImageTypeId)) {
+      return TRUE;
+    }
+    CurrentImageInfo = (EFI_FIRMWARE_IMAGE_DESCRIPTOR *)((UINT8 *)CurrentImageInfo + DescriptorSize);
+  }
+  return FALSE;
+}
+
+/**
+  return the FMP whoes ImageInfo includes the ImageTypeId.
+
+  @param[in] ImageTypeId         A unique GUID identifying the firmware image type.
+
+  @return The FMP whoes ImageInfo includes the ImageTypeId
+**/
+EFI_FIRMWARE_MANAGEMENT_PROTOCOL *
+FindFmpFromImageTypeId (
+  IN EFI_GUID  *ImageTypeId
+  )
+{
+  EFI_STATUS                                    Status;
+  EFI_FIRMWARE_MANAGEMENT_PROTOCOL              *Fmp;
+  EFI_FIRMWARE_MANAGEMENT_PROTOCOL              *TargetFmp;
+  EFI_HANDLE                                    *HandleBuffer;
+  UINTN                                         NumberOfHandles;
+  UINTN                                         Index;
+  EFI_FIRMWARE_IMAGE_DESCRIPTOR                 *FmpImageInfoBuf;
+  UINTN                                         ImageInfoSize;
+  UINT32                                        FmpImageInfoDescriptorVer;
+  UINT8                                         FmpImageInfoCount;
+  UINTN                                         DescriptorSize;
+  UINT32                                        PackageVersion;
+  CHAR16                                        *PackageVersionName;
+
+  Status = gBS->LocateHandleBuffer (
+                  ByProtocol,
+                  &gEfiFirmwareManagementProtocolGuid,
+                  NULL,
+                  &NumberOfHandles,
+                  &HandleBuffer
+                  );
+  if (EFI_ERROR(Status)) {
+    Print(L"FMP protocol - %r\n", EFI_NOT_FOUND);
+    return NULL;
+  }
+
+  TargetFmp = NULL;
+  for (Index = 0; Index < NumberOfHandles; Index++) {
+    Status = gBS->HandleProtocol(
+                    HandleBuffer[Index],
+                    &gEfiFirmwareManagementProtocolGuid,
+                    (VOID **)&Fmp
+                    );
+    if (EFI_ERROR(Status)) {
+      continue;
+    }
+
+    ImageInfoSize = 0;
+    Status = Fmp->GetImageInfo (
+                    Fmp,
+                    &ImageInfoSize,
+                    NULL,
+                    NULL,
+                    NULL,
+                    NULL,
+                    NULL,
+                    NULL
+                    );
+    if (Status != EFI_BUFFER_TOO_SMALL) {
+      continue;
+    }
+
+    FmpImageInfoBuf = NULL;
+    FmpImageInfoBuf = AllocateZeroPool (ImageInfoSize);
+    if (FmpImageInfoBuf == NULL) {
+      FreePool(HandleBuffer);
+      Print(L"Out of resource\n");
+      return NULL;
+    }
+
+    PackageVersionName = NULL;
+    Status = Fmp->GetImageInfo (
+                    Fmp,
+                    &ImageInfoSize,               // ImageInfoSize
+                    FmpImageInfoBuf,              // ImageInfo
+                    &FmpImageInfoDescriptorVer,   // DescriptorVersion
+                    &FmpImageInfoCount,           // DescriptorCount
+                    &DescriptorSize,              // DescriptorSize
+                    &PackageVersion,              // PackageVersion
+                    &PackageVersionName           // PackageVersionName
+                    );
+
+    //
+    // If FMP GetInformation interface failed, skip this resource
+    //
+    if (EFI_ERROR(Status)) {
+      FreePool(FmpImageInfoBuf);
+      continue;
+    }
+
+    if (PackageVersionName != NULL) {
+      FreePool(PackageVersionName);
+    }
+
+    if (IsThisFmpImageInfo (FmpImageInfoBuf, FmpImageInfoCount, DescriptorSize, ImageTypeId)) {
+      TargetFmp = Fmp;
+    }
+    FreePool(FmpImageInfoBuf);
+    if (TargetFmp != NULL) {
+      break;
+    }
+  }
+  FreePool(HandleBuffer);
+  return TargetFmp;
+}
+
+/**
+  Dump FMP image data.
+
+  @param[in]  ImageTypeId   The ImageTypeId of the FMP image.
+                            It is used to identify the FMP protocol.
+  @param[in]  ImageIndex    The ImageIndex of the FMP image.
+                            It is the input parameter for FMP->GetImage().
+  @param[in]  ImageName     The file name to hold the output FMP image.
+**/
+VOID
+DumpFmpImage (
+  IN EFI_GUID  *ImageTypeId,
+  IN UINTN     ImageIndex,
+  IN CHAR16    *ImageName
+  )
+{
+  EFI_STATUS                                    Status;
+  EFI_FIRMWARE_MANAGEMENT_PROTOCOL              *Fmp;
+  VOID                                          *Image;
+  UINTN                                         ImageSize;
+
+  Fmp = FindFmpFromImageTypeId (ImageTypeId);
+  if (Fmp == NULL) {
+    Print(L"No FMP include ImageTypeId %g\n", ImageTypeId);
+    return ;
+  }
+
+  if (ImageIndex > 0xFF) {
+    Print(L"ImageIndex 0x%x too big\n", ImageIndex);
+    return ;
+  }
+
+  Image = Fmp;
+  ImageSize = 0;
+  Status = Fmp->GetImage (Fmp, (UINT8)ImageIndex, Image, &ImageSize);
+  if (Status != EFI_BUFFER_TOO_SMALL) {
+    Print(L"Fmp->GetImage - %r\n", Status);
+    return ;
+  }
+
+  Image = AllocatePool (ImageSize);
+  if (Image == NULL) {
+    Print(L"Allocate FmpImage 0x%x - %r\n", ImageSize, EFI_OUT_OF_RESOURCES);
+    return ;
+  }
+
+  Status = Fmp->GetImage (Fmp, (UINT8)ImageIndex, Image, &ImageSize);
+  if (EFI_ERROR(Status)) {
+    Print(L"Fmp->GetImage - %r\n", Status);
+    return ;
+  }
+
+  Status = WriteFileFromBuffer(ImageName, ImageSize, Image);
+  Print(L"CapsuleApp: Dump %g ImageIndex (0x%x) to %s %r\n", ImageTypeId, ImageIndex, ImageName, Status);
+
+  return ;
+}
-- 
2.7.4.windows.1