MdeModulePkg/DxeCapsuleLibFmp: Add more check for the UX capsule

Contributed-under: TianoCore Contribution Agreement 1.1
Signed-off-by: Ruiyu Ni <ruiyu.ni@intel.com>
Reviewed-by: Jiewen Yao <jiewen.yao@intel.com>

diff --git a/MdeModulePkg/Library/DxeCapsuleLibFmp/DxeCapsuleLib.c b/MdeModulePkg/Library/DxeCapsuleLibFmp/DxeCapsuleLib.c
index 15dbc0021600c1c3e1ae968c1b9d2e2e3b6d07bb..555c5971d00d231cf9d4a381faba346e33fd06e6 100644 (file)
--- a/MdeModulePkg/Library/DxeCapsuleLibFmp/DxeCapsuleLib.c
+++ b/MdeModulePkg/Library/DxeCapsuleLibFmp/DxeCapsuleLib.c
@@ -330,8 +330,25 @@ DisplayCapsuleImage (
   UINTN                         Width;\r
   EFI_GRAPHICS_OUTPUT_PROTOCOL  *GraphicsOutput;\r
 \r
   UINTN                         Width;\r
   EFI_GRAPHICS_OUTPUT_PROTOCOL  *GraphicsOutput;\r
 \r

-  ImagePayload = (DISPLAY_DISPLAY_PAYLOAD *)(CapsuleHeader + 1);\r
-  PayloadSize = CapsuleHeader->CapsuleImageSize - sizeof(EFI_CAPSULE_HEADER);\r
+  //\r
+  // UX capsule doesn't have extended header entries.\r
+  //\r
+  if (CapsuleHeader->HeaderSize != sizeof (EFI_CAPSULE_HEADER)) {\r
+    return EFI_UNSUPPORTED;\r
+  }\r
+  ImagePayload = (DISPLAY_DISPLAY_PAYLOAD *)((UINTN) CapsuleHeader + CapsuleHeader->HeaderSize);\r
+  //\r
+  // (CapsuleImageSize > HeaderSize) is guaranteed by IsValidCapsuleHeader().\r
+  //\r
+  PayloadSize = CapsuleHeader->CapsuleImageSize - CapsuleHeader->HeaderSize;\r
+\r
+  //\r
+  // Make sure the image payload at least contain the DISPLAY_DISPLAY_PAYLOAD header.\r
+  // Further size check is performed by the logic translating BMP to GOP BLT.\r
+  //\r
+  if (PayloadSize <= sizeof (DISPLAY_DISPLAY_PAYLOAD)) {\r
+    return EFI_INVALID_PARAMETER;\r
+  }\r

 \r
   if (ImagePayload->Version != 1) {\r
     return EFI_UNSUPPORTED;\r
 \r
   if (ImagePayload->Version != 1) {\r
     return EFI_UNSUPPORTED;\r