[Patch] BaseTools: UpdateImageSize include Image auth info for FMP Auth capsule

[Yonghong Zhu <yonghong.zhu@intel.com>]

Per UEFI spec UpdateImageSize may or may not include Firmware Image
Authentication information. so for FMP auth capsule, UpdateImageSize
should include the Image auth info.

Cc: Liming Gao <liming.gao@intel.com>
Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Yonghong Zhu <yonghong.zhu@intel.com>
---
 BaseTools/Source/Python/GenFds/Capsule.py     | 34 ++++++---------------------
 BaseTools/Source/Python/GenFds/CapsuleData.py | 21 ++++++++++++++++-
 2 files changed, 27 insertions(+), 28 deletions(-)

diff --git a/BaseTools/Source/Python/GenFds/Capsule.py b/BaseTools/Source/Python/GenFds/Capsule.py
index 93ecee1..c98c054 100644
--- a/BaseTools/Source/Python/GenFds/Capsule.py
+++ b/BaseTools/Source/Python/GenFds/Capsule.py
@@ -139,11 +139,10 @@ class Capsule (CapsuleClassObject) :
             PreSize += os.path.getsize(FileName)
             File = open(FileName, 'rb')
             Content.write(File.read())
             File.close()
         for fmp in self.FmpPayloadList:
-            Buffer = fmp.GenCapsuleSubItem()
             if fmp.Certificate_Guid:
                 ExternalTool, ExternalOption = FindExtendTool([], GenFdsGlobalVariable.ArchList, fmp.Certificate_Guid)
                 CmdOption = ''
                 CapInputFile = fmp.ImageFile
                 if not os.path.isabs(fmp.ImageFile):
@@ -160,37 +159,18 @@ class Capsule (CapsuleClassObject) :
                 GenFdsGlobalVariable.CallExternalTool(CmdList, "Failed to generate FMP auth capsule")
                 if uuid.UUID(fmp.Certificate_Guid) == EFI_CERT_TYPE_PKCS7_GUID:
                     dwLength = 4 + 2 + 2 + 16 + os.path.getsize(CapOutputTmp) - os.path.getsize(CapInputFile)
                 else:
                     dwLength = 4 + 2 + 2 + 16 + 16 + 256 + 256
-                Buffer += pack('Q', fmp.MonotonicCount)
-                Buffer += pack('I', dwLength)
-                Buffer += pack('H', WIN_CERT_REVISION)
-                Buffer += pack('H', WIN_CERT_TYPE_EFI_GUID)
-                Buffer += uuid.UUID(fmp.Certificate_Guid).get_bytes_le()
-                if os.path.exists(CapOutputTmp):
-                    TmpFile = open(CapOutputTmp, 'rb')
-                    Buffer += TmpFile.read()
-                    TmpFile.close()
-                    if fmp.VendorCodeFile:
-                        VendorFile = open(fmp.VendorCodeFile, 'rb')
-                        Buffer += VendorFile.read()
-                        VendorFile.close()
-                    FwMgrHdr.write(pack('=Q', PreSize))
-                    PreSize += len(Buffer)
-                    Content.write(Buffer)
+                fmp.ImageFile = CapOutputTmp
+                AuthData = [fmp.MonotonicCount, dwLength, WIN_CERT_REVISION, WIN_CERT_TYPE_EFI_GUID, fmp.Certificate_Guid]
+                Buffer = fmp.GenCapsuleSubItem(AuthData)
             else:
-                ImageFile = open(fmp.ImageFile, 'rb')
-                Buffer += ImageFile.read()
-                ImageFile.close()
-                if fmp.VendorCodeFile:
-                    VendorFile = open(fmp.VendorCodeFile, 'rb')
-                    Buffer += VendorFile.read()
-                    VendorFile.close()
-                FwMgrHdr.write(pack('=Q', PreSize))
-                PreSize += len(Buffer)
-                Content.write(Buffer)
+                Buffer = fmp.GenCapsuleSubItem()
+            FwMgrHdr.write(pack('=Q', PreSize))
+            PreSize += len(Buffer)
+            Content.write(Buffer)
         BodySize = len(FwMgrHdr.getvalue()) + len(Content.getvalue())
         Header.write(pack('=I', HdrSize + BodySize))
         #
         # The real capsule header structure is 28 bytes
         #
diff --git a/BaseTools/Source/Python/GenFds/CapsuleData.py b/BaseTools/Source/Python/GenFds/CapsuleData.py
index 5d5a1e4..07cc198 100644
--- a/BaseTools/Source/Python/GenFds/CapsuleData.py
+++ b/BaseTools/Source/Python/GenFds/CapsuleData.py
@@ -19,10 +19,11 @@ import Ffs
 from GenFdsGlobalVariable import GenFdsGlobalVariable
 import StringIO
 from struct import pack
 import os
 from Common.Misc import SaveFileOnChange
+import uuid
 
 ## base class for capsule data
 #
 #
 class CapsuleData:
@@ -181,14 +182,18 @@ class CapsulePayload(CapsuleData):
         self.ImageFile = None
         self.VendorCodeFile = None
         self.Certificate_Guid = None
         self.MonotonicCount = None
 
-    def GenCapsuleSubItem(self):
+    def GenCapsuleSubItem(self, AuthData=[]):
         if not self.Version:
             self.Version = 0x00000002
         ImageFileSize = os.path.getsize(self.ImageFile)
+        if AuthData:
+            # the ImageFileSize need include the full authenticated info size. From first bytes of MonotonicCount to last bytes of certificate.
+            # the 32 bit is the MonotonicCount, dwLength, wRevision, wCertificateType and CertType
+            ImageFileSize += 32
         VendorFileSize = 0
         if self.VendorCodeFile:
             VendorFileSize = os.path.getsize(self.VendorCodeFile)
 
         #
@@ -214,6 +219,20 @@ class CapsulePayload(CapsuleData):
                        0,
                        ImageFileSize,
                        VendorFileSize,
                        int(self.HardwareInstance, 16)
                        )
+        if AuthData:
+            Buffer += pack('QIHH', AuthData[0], AuthData[1], AuthData[2], AuthData[3])
+            Buffer += uuid.UUID(AuthData[4]).get_bytes_le()
+
+        #
+        # Append file content to the structure
+        #
+        ImageFile = open(self.ImageFile, 'rb')
+        Buffer += ImageFile.read()
+        ImageFile.close()
+        if self.VendorCodeFile:
+            VendorFile = open(self.VendorCodeFile, 'rb')
+            Buffer += VendorFile.read()
+            VendorFile.close()
         return Buffer
-- 
2.6.1.windows.1