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>

2 files changed, 27 insertions, 28 deletions

diff --git a/BaseTools/Source/Python/GenFds/Capsule.py b/BaseTools/Source/Python/GenFds/Capsule.py
index 93ecee10c6..c98c054771 100644
--- a/BaseTools/Source/Python/GenFds/Capsule.py
+++ b/BaseTools/Source/Python/GenFds/Capsule.py
@@ -141,7 +141,6 @@ class Capsule (CapsuleClassObject) :
             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 @@ class Capsule (CapsuleClassObject) :
                     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 5d5a1e41ea..07cc1981d6 100644
--- a/BaseTools/Source/Python/GenFds/CapsuleData.py
+++ b/BaseTools/Source/Python/GenFds/CapsuleData.py
@@ -21,6 +21,7 @@ import StringIO
 from struct import pack

 import os

 from Common.Misc import SaveFileOnChange

+import uuid

 

 ## base class for capsule data

 #

@@ -183,10 +184,14 @@ class CapsulePayload(CapsuleData):
         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 @@ class CapsulePayload(CapsuleData):
                        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