[PATCH] MdeModulePkg/CdExpressPei: Fix capsule size mismatch issue.

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

Current PeiCdExpress driver will return capsule size to
be block aligned.
It will fail if we check the capsuleImageSize field.

The patch correct the returned capsule size.

Tested-by: Eric Dong <eric.dong@intel.com>

Cc: Feng Tian <feng.tian@intel.com>
Cc: Eric Dong <eric.dong@intel.com>
Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Jiewen Yao <jiewen.yao@intel.com>
---
 MdeModulePkg/Universal/Disk/CdExpressPei/PeiCdExpress.c | 9 +++++----
 MdeModulePkg/Universal/Disk/CdExpressPei/PeiCdExpress.h | 1 +
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/MdeModulePkg/Universal/Disk/CdExpressPei/PeiCdExpress.c b/MdeModulePkg/Universal/Disk/CdExpressPei/PeiCdExpress.c
index f78118a..371ab73 100644
--- a/MdeModulePkg/Universal/Disk/CdExpressPei/PeiCdExpress.c
+++ b/MdeModulePkg/Universal/Disk/CdExpressPei/PeiCdExpress.c
@@ -1,7 +1,7 @@
 /** @file
   Source file for CD recovery PEIM
 
-Copyright (c) 2006 - 2015, Intel Corporation. All rights reserved.<BR>
+Copyright (c) 2006 - 2016, Intel Corporation. All rights reserved.<BR>
 
 This program and the accompanying materials
 are licensed and made available under the terms and conditions
@@ -477,13 +477,14 @@ RetrieveCapsuleFileFromRoot (
     }
 
     PrivateData->CapsuleData[PrivateData->CapsuleCount].CapsuleStartLBA = FileRecord->LocationOfExtent[0];
-    PrivateData->CapsuleData[PrivateData->CapsuleCount].CapsuleSize =
+    PrivateData->CapsuleData[PrivateData->CapsuleCount].CapsuleBlockAlignedSize =
       (
         FileRecord->DataLength[0] /
         PEI_CD_BLOCK_SIZE +
         1
       ) *
       PEI_CD_BLOCK_SIZE;
+    PrivateData->CapsuleData[PrivateData->CapsuleCount].CapsuleSize = FileRecord->DataLength[0];
 
     return EFI_SUCCESS;
   }
@@ -659,7 +660,7 @@ LoadRecoveryCapsule (
                           BlockIo2Ppi,
                           PrivateData->CapsuleData[CapsuleInstance - 1].IndexBlock,
                           PrivateData->CapsuleData[CapsuleInstance - 1].CapsuleStartLBA,
-                          PrivateData->CapsuleData[CapsuleInstance - 1].CapsuleSize,
+                          PrivateData->CapsuleData[CapsuleInstance - 1].CapsuleBlockAlignedSize,
                           Buffer
                           );
   } else {
@@ -668,7 +669,7 @@ LoadRecoveryCapsule (
                           BlockIoPpi,
                           PrivateData->CapsuleData[CapsuleInstance - 1].IndexBlock,
                           PrivateData->CapsuleData[CapsuleInstance - 1].CapsuleStartLBA,
-                          PrivateData->CapsuleData[CapsuleInstance - 1].CapsuleSize,
+                          PrivateData->CapsuleData[CapsuleInstance - 1].CapsuleBlockAlignedSize,
                           Buffer
                           );
   }
diff --git a/MdeModulePkg/Universal/Disk/CdExpressPei/PeiCdExpress.h b/MdeModulePkg/Universal/Disk/CdExpressPei/PeiCdExpress.h
index 5ca2672..91f34dc 100644
--- a/MdeModulePkg/Universal/Disk/CdExpressPei/PeiCdExpress.h
+++ b/MdeModulePkg/Universal/Disk/CdExpressPei/PeiCdExpress.h
@@ -66,6 +66,7 @@ WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
 typedef struct {
   UINTN                           CapsuleStartLBA;
   UINTN                           CapsuleSize;
+  UINTN                           CapsuleBlockAlignedSize;
   UINTN                           IndexBlock;
   EFI_PEI_RECOVERY_BLOCK_IO_PPI   *BlockIo;
   EFI_PEI_RECOVERY_BLOCK_IO2_PPI  *BlockIo2;
-- 
2.7.4.windows.1