BaseTools: UpdateImageSize include Image auth info for FMP Auth capsule

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>
Reviewed-by: Liming Gao <liming.gao@intel.com>
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
@@ -141,7 +141,6 @@
             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 = ''

@@ -162,33 +161,14 @@
                     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))

         #

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
@@ -21,6 +21,7 @@
 from struct import pack

 import os

 from Common.Misc import SaveFileOnChange

+import uuid

 

 ## base class for capsule data

 #

@@ -183,10 +184,14 @@
         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)

@@ -216,4 +221,18 @@
                        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