From 601328475358a7baa4da682004245a94de51a2c2 Mon Sep 17 00:00:00 2001
From: Sean McGinn <sean.mcginn@intel.com>
Date: Tue, 30 Aug 2022 16:49:14 -0700
Subject: [PATCH] Stop TCO Timer in Beginning of FWU PLD

Stop TCO timer in beginning of FWU PLD
to ensure it is started without error

Signed-off-by: Sean McGinn <sean.mcginn@intel.com>
---
 BootloaderCorePkg/Stage2/Stage2.c            | 10 ++++++----
 PayloadPkg/FirmwareUpdate/FirmwareUpdate.c   |  9 +++++++++
 PayloadPkg/FirmwareUpdate/FirmwareUpdate.inf |  1 +
 3 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/BootloaderCorePkg/Stage2/Stage2.c b/BootloaderCorePkg/Stage2/Stage2.c
index 7a05c3e0..5e72b8b9 100644
--- a/BootloaderCorePkg/Stage2/Stage2.c
+++ b/BootloaderCorePkg/Stage2/Stage2.c
@@ -289,8 +289,9 @@ NormalBootPath (
   PrintStackHeapInfo ();
   DEBUG_CODE_END ();
 
-  // Stop timer so boot retries stop
-  if (PcdGetBool (PcdSblResiliencyEnabled)) {
+  // FWU payload is the only payload in SBL scope, so stop TCO
+  // timer if another payload is set to be launched
+  if (PcdGetBool (PcdSblResiliencyEnabled) && GetBootMode () != BOOT_ON_FLASH_UPDATE) {
     StopTcoTimer ();
     ClearFailedBootCount ();
   }
@@ -362,8 +363,9 @@ S3ResumePath (
 
   AddMeasurePoint (0x31F0);
 
-  // Stop timer so boot retries stop
-  if (PcdGetBool (PcdSblResiliencyEnabled)) {
+  // FWU payload is the only payload in SBL scope, so stop TCO
+  // timer if another payload is set to be launched
+  if (PcdGetBool (PcdSblResiliencyEnabled) && GetBootMode () != BOOT_ON_FLASH_UPDATE) {
     StopTcoTimer ();
     ClearFailedBootCount ();
   }
diff --git a/PayloadPkg/FirmwareUpdate/FirmwareUpdate.c b/PayloadPkg/FirmwareUpdate/FirmwareUpdate.c
index f43b197a..142a4d68 100644
--- a/PayloadPkg/FirmwareUpdate/FirmwareUpdate.c
+++ b/PayloadPkg/FirmwareUpdate/FirmwareUpdate.c
@@ -31,6 +31,7 @@ SPDX-License-Identifier: BSD-2-Clause-Patent
 #include <Guid/OsBootOptionGuid.h>
 #include <Library/BootGuardLib.h>
 #include <Library/WatchDogTimerLib.h>
+#include <Library/TcoTimerLib.h>
 #include "FirmwareUpdateHelper.h"
 
 UINT32   mSblImageBiosRgnOffset;
@@ -1769,6 +1770,14 @@ PayloadMain (
   //
   (VOID) PcdSet32S (PcdFwUpdStatusBase, mSblImageBiosRgnOffset + (FlashMap->RomSize - (~RsvdBase + 1)));
 
+  //
+  // Stop TCO timer here as SPI read/write timings are highly variable
+  //
+  if (PcdGetBool (PcdSblResiliencyEnabled)) {
+    StopTcoTimer ();
+    ClearFailedBootCount ();
+  }
+
   //
   // Perform firmware recovery/update
   //
diff --git a/PayloadPkg/FirmwareUpdate/FirmwareUpdate.inf b/PayloadPkg/FirmwareUpdate/FirmwareUpdate.inf
index a509a371..1ecd5a23 100644
--- a/PayloadPkg/FirmwareUpdate/FirmwareUpdate.inf
+++ b/PayloadPkg/FirmwareUpdate/FirmwareUpdate.inf
@@ -55,6 +55,7 @@
   TimerLib
   BootGuardLib
   WatchDogTimerLib
+  TcoTimerLib
 
 [Guids]
   gLoaderMemoryMapInfoGuid