From fb57c30b703ee64415c43102862cfc2c2f2664be Mon Sep 17 00:00:00 2001
From: Star Zeng <star.zeng@intel.com>
Date: Thu, 26 Jul 2018 16:59:43 +0800
Subject: [PATCH] MdeModulePkg CapsuleApp: Check capsule header for -D and -N
 options

Then meaningful error message can be shown when the input image is
unexpected.

Cc: Michael D Kinney <michael.d.kinney@intel.com>
Cc: Jiewen Yao <jiewen.yao@intel.com>
Cc: Yonghong Zhu <yonghong.zhu@intel.com>
Contributed-under: TianoCore Contribution Agreement 1.1
Signed-off-by: Star Zeng <star.zeng@intel.com>
Reviewed-by: Jiewen Yao <jiewen.yao@intel.com>
---
 .../Application/CapsuleApp/CapsuleApp.c       | 101 ++++++++++++------
 .../Application/CapsuleApp/CapsuleDump.c      |  23 ++++
 2 files changed, 90 insertions(+), 34 deletions(-)

diff --git a/MdeModulePkg/Application/CapsuleApp/CapsuleApp.c b/MdeModulePkg/Application/CapsuleApp/CapsuleApp.c
index 2967b0d1dd..894da2f2d9 100644
--- a/MdeModulePkg/Application/CapsuleApp/CapsuleApp.c
+++ b/MdeModulePkg/Application/CapsuleApp/CapsuleApp.c
@@ -362,6 +362,60 @@ GetEsrtFwType (
   return ESRT_FW_TYPE_UNKNOWN;
 }
 
+/**
+  Validate if it is valid capsule header
+
+  This function assumes the caller provided correct CapsuleHeader pointer
+  and CapsuleSize.
+
+  This function validates the fields in EFI_CAPSULE_HEADER.
+
+  @param[in] CapsuleHeader  Points to a capsule header.
+  @param[in] CapsuleSize    Size of the whole capsule image.
+
+**/
+BOOLEAN
+IsValidCapsuleHeader (
+  IN EFI_CAPSULE_HEADER     *CapsuleHeader,
+  IN UINT64                 CapsuleSize
+  )
+{
+  if (CapsuleSize < sizeof (EFI_CAPSULE_HEADER)) {
+    return FALSE;
+  }
+  if (CapsuleHeader->CapsuleImageSize != CapsuleSize) {
+    return FALSE;
+  }
+  if (CapsuleHeader->HeaderSize > CapsuleHeader->CapsuleImageSize) {
+    return FALSE;
+  }
+  if (CapsuleHeader->HeaderSize < sizeof (EFI_CAPSULE_HEADER)) {
+    return FALSE;
+  }
+
+  return TRUE;
+}
+
+/**
+  Return if this CapsuleGuid is a FMP capsule GUID or not.
+
+  @param[in] CapsuleGuid A pointer to EFI_GUID
+
+  @retval TRUE  It is a FMP capsule GUID.
+  @retval FALSE It is not a FMP capsule GUID.
+**/
+BOOLEAN
+IsFmpCapsuleGuid (
+  IN EFI_GUID  *CapsuleGuid
+  )
+{
+  if (CompareGuid(&gEfiFmpCapsuleGuid, CapsuleGuid)) {
+    return TRUE;
+  }
+
+  return FALSE;
+}
+
 /**
   Append a capsule header on top of current image.
   This function follows Windows UEFI Firmware Update Platform document.
@@ -407,15 +461,28 @@ CreateNestedFmp (
     Print(L"CapsuleApp: Capsule image (%s) is not found.\n", CapsuleName);
     goto Done;
   }
+  if (!IsValidCapsuleHeader (CapsuleBuffer, FileSize)) {
+    Print(L"CapsuleApp: Capsule image (%s) is not a valid capsule.\n", CapsuleName);
+    Status = EFI_INVALID_PARAMETER;
+    goto Done;
+  }
+
+  if (!IsFmpCapsuleGuid (&((EFI_CAPSULE_HEADER *) CapsuleBuffer)->CapsuleGuid)) {
+    Print(L"CapsuleApp: Capsule image (%s) is not a FMP capsule.\n", CapsuleName);
+    Status = EFI_INVALID_PARAMETER;
+    goto Done;
+  }
 
   ImageTypeId = GetCapsuleImageTypeId(CapsuleBuffer);
   if (ImageTypeId == NULL) {
     Print(L"CapsuleApp: Capsule ImageTypeId is not found.\n");
+    Status = EFI_INVALID_PARAMETER;
     goto Done;
   }
   FwType = GetEsrtFwType(ImageTypeId);
   if ((FwType != ESRT_FW_TYPE_SYSTEMFIRMWARE) && (FwType != ESRT_FW_TYPE_DEVICEFIRMWARE)) {
     Print(L"CapsuleApp: Capsule FwType is invalid.\n");
+    Status = EFI_INVALID_PARAMETER;
     goto Done;
   }
 
@@ -724,40 +791,6 @@ CleanGatherList (
   }
 }
 
-/**
-  Validate if it is valid capsule header
-
-  This function assumes the caller provided correct CapsuleHeader pointer
-  and CapsuleSize.
-
-  This function validates the fields in EFI_CAPSULE_HEADER.
-
-  @param[in] CapsuleHeader  Points to a capsule header.
-  @param[in] CapsuleSize    Size of the whole capsule image.
-
-**/
-BOOLEAN
-IsValidCapsuleHeader (
-  IN EFI_CAPSULE_HEADER     *CapsuleHeader,
-  IN UINT64                 CapsuleSize
-  )
-{
-  if (CapsuleSize < sizeof (EFI_CAPSULE_HEADER)) {
-    return FALSE;
-  }
-  if (CapsuleHeader->CapsuleImageSize != CapsuleSize) {
-    return FALSE;
-  }
-  if (CapsuleHeader->HeaderSize > CapsuleHeader->CapsuleImageSize) {
-    return FALSE;
-  }
-  if (CapsuleHeader->HeaderSize < sizeof (EFI_CAPSULE_HEADER)) {
-    return FALSE;
-  }
-
-  return TRUE;
-}
-
 /**
   Print APP usage.
 **/
diff --git a/MdeModulePkg/Application/CapsuleApp/CapsuleDump.c b/MdeModulePkg/Application/CapsuleApp/CapsuleDump.c
index 11bf2e1d45..45c3ecd050 100644
--- a/MdeModulePkg/Application/CapsuleApp/CapsuleDump.c
+++ b/MdeModulePkg/Application/CapsuleApp/CapsuleDump.c
@@ -61,6 +61,24 @@ WriteFileFromBuffer (
   IN  VOID                                 *Buffer
   );
 
+/**
+  Validate if it is valid capsule header
+
+  This function assumes the caller provided correct CapsuleHeader pointer
+  and CapsuleSize.
+
+  This function validates the fields in EFI_CAPSULE_HEADER.
+
+  @param[in] CapsuleHeader  Points to a capsule header.
+  @param[in] CapsuleSize    Size of the whole capsule image.
+
+**/
+BOOLEAN
+IsValidCapsuleHeader (
+  IN EFI_CAPSULE_HEADER     *CapsuleHeader,
+  IN UINT64                 CapsuleSize
+  );
+
 /**
   Dump UX capsule information.
 
@@ -248,6 +266,11 @@ DumpCapsule (
     Print(L"CapsuleApp: Capsule (%s) is not found.\n", CapsuleName);
     goto Done;
   }
+  if (!IsValidCapsuleHeader (Buffer, FileSize)) {
+    Print(L"CapsuleApp: Capsule image (%s) is not a valid capsule.\n", CapsuleName);
+    Status = EFI_INVALID_PARAMETER;
+    goto Done;
+  }
 
   CapsuleHeader = Buffer;
   if (CompareGuid(&CapsuleHeader->CapsuleGuid, &gWindowsUxCapsuleGuid)) {
-- 
2.39.2