From af15e4535dcd4b2bfd0799ab36f0a3622687bdd3 Mon Sep 17 00:00:00 2001
From: Carsten Haitzler <carsten.haitzler@foss.arm.com>
Date: Fri, 17 May 2024 11:49:22 +0100
Subject: [PATCH] ArmPkg: Fix timer wrap-around

The timer counter register can wrap around and when this happens,
we'll get misbehavior for any MicroSecondDelay() calls. This adds
handling for that.

Signed-off-by: Carsten Haitzler <carsten.haitzler@foss.arm.com>
---
 ArmPkg/Library/ArmArchTimerLib/ArmArchTimerLib.c | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/ArmPkg/Library/ArmArchTimerLib/ArmArchTimerLib.c b/ArmPkg/Library/ArmArchTimerLib/ArmArchTimerLib.c
index 76f94c916155..cc1be0125b65 100644
--- a/ArmPkg/Library/ArmArchTimerLib/ArmArchTimerLib.c
+++ b/ArmPkg/Library/ArmArchTimerLib/ArmArchTimerLib.c
@@ -62,6 +62,8 @@ MicroSecondDelay (
 {
   UINT64  TimerTicks64;
   UINT64  SystemCounterVal;
+  UINT64  PreviousSystemCounterVal;
+  UINT64  DeltaCounterVal;
 
   // Calculate counter ticks that represent requested delay:
   //  = MicroSeconds x TICKS_PER_MICRO_SEC
@@ -75,13 +77,17 @@ MicroSecondDelay (
                    );
 
   // Read System Counter value
-  SystemCounterVal = ArmGenericTimerGetSystemCount ();
-
-  TimerTicks64 += SystemCounterVal;
+  PreviousSystemCounterVal = ArmGenericTimerGetSystemCount ();
 
   // Wait until delay count expires.
-  while (SystemCounterVal < TimerTicks64) {
+  while (TimerTicks64 > 0) {
     SystemCounterVal = ArmGenericTimerGetSystemCount ();
+    // Get how much we advanced this tick. Wrap around still has delta correct
+    DeltaCounterVal = (SystemCounterVal - PreviousSystemCounterVal)
+                      & (MAX_UINT64 >> 8); // Account for a lesser (minimum) size
+    // Never wrap back around below zero by choosing the min and thus stop at 0
+    TimerTicks64            -= MIN (TimerTicks64, DeltaCounterVal);
+    PreviousSystemCounterVal = SystemCounterVal;
   }
 
   return MicroSeconds;