From 5f53a7aa59d4df1fe4326af18a9240d4dfebc129 Mon Sep 17 00:00:00 2001
From: Yonghong Zhu <yonghong.zhu@intel.com>
Date: Mon, 29 Aug 2016 15:44:59 +0800
Subject: [PATCH] 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>
---
 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 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
-- 
2.39.2