From 9efd95cda50605c52a9652b5ad01e09e6fd86106 Mon Sep 17 00:00:00 2001
From: Liam <byteslice@airmail.cc>
Date: Tue, 28 Mar 2023 22:28:27 -0400
Subject: kernel: fix unbounded stack usage in atomics

---
 src/core/hle/kernel/k_address_arbiter.cpp    | 68 ++++++++++++++++------------
 src/core/hle/kernel/k_condition_variable.cpp | 27 ++++++-----
 2 files changed, 56 insertions(+), 39 deletions(-)

diff --git a/src/core/hle/kernel/k_address_arbiter.cpp b/src/core/hle/kernel/k_address_arbiter.cpp
index 08c254028..78d43d729 100644
--- a/src/core/hle/kernel/k_address_arbiter.cpp
+++ b/src/core/hle/kernel/k_address_arbiter.cpp
@@ -35,24 +35,30 @@ bool DecrementIfLessThan(Core::System& system, s32* out, KProcessAddress address
 
     // TODO(bunnei): We should call CanAccessAtomic(..) here.
 
-    // Load the value from the address.
-    const s32 current_value =
-        static_cast<s32>(monitor.ExclusiveRead32(current_core, GetInteger(address)));
-
-    // Compare it to the desired one.
-    if (current_value < value) {
-        // If less than, we want to try to decrement.
-        const s32 decrement_value = current_value - 1;
+    s32 current_value{};
+
+    while (true) {
+        // Load the value from the address.
+        current_value =
+            static_cast<s32>(monitor.ExclusiveRead32(current_core, GetInteger(address)));
+
+        // Compare it to the desired one.
+        if (current_value < value) {
+            // If less than, we want to try to decrement.
+            const s32 decrement_value = current_value - 1;
+
+            // Decrement and try to store.
+            if (monitor.ExclusiveWrite32(current_core, GetInteger(address),
+                                         static_cast<u32>(decrement_value))) {
+                break;
+            }
 
-        // Decrement and try to store.
-        if (!monitor.ExclusiveWrite32(current_core, GetInteger(address),
-                                      static_cast<u32>(decrement_value))) {
             // If we failed to store, try again.
-            DecrementIfLessThan(system, out, address, value);
+        } else {
+            // Otherwise, clear our exclusive hold and finish
+            monitor.ClearExclusive(current_core);
+            break;
         }
-    } else {
-        // Otherwise, clear our exclusive hold and finish
-        monitor.ClearExclusive(current_core);
     }
 
     // We're done.
@@ -70,23 +76,29 @@ bool UpdateIfEqual(Core::System& system, s32* out, KProcessAddress address, s32
 
     // TODO(bunnei): We should call CanAccessAtomic(..) here.
 
-    // Load the value from the address.
-    const s32 current_value =
-        static_cast<s32>(monitor.ExclusiveRead32(current_core, GetInteger(address)));
+    s32 current_value{};
 
-    // Compare it to the desired one.
-    if (current_value == value) {
-        // If equal, we want to try to write the new value.
+    // Load the value from the address.
+    while (true) {
+        current_value =
+            static_cast<s32>(monitor.ExclusiveRead32(current_core, GetInteger(address)));
+
+        // Compare it to the desired one.
+        if (current_value == value) {
+            // If equal, we want to try to write the new value.
+
+            // Try to store.
+            if (monitor.ExclusiveWrite32(current_core, GetInteger(address),
+                                         static_cast<u32>(new_value))) {
+                break;
+            }
 
-        // Try to store.
-        if (!monitor.ExclusiveWrite32(current_core, GetInteger(address),
-                                      static_cast<u32>(new_value))) {
             // If we failed to store, try again.
-            UpdateIfEqual(system, out, address, value, new_value);
+        } else {
+            // Otherwise, clear our exclusive hold and finish.
+            monitor.ClearExclusive(current_core);
+            break;
         }
-    } else {
-        // Otherwise, clear our exclusive hold and finish.
-        monitor.ClearExclusive(current_core);
     }
 
     // We're done.
diff --git a/src/core/hle/kernel/k_condition_variable.cpp b/src/core/hle/kernel/k_condition_variable.cpp
index 73017cf99..efbac0e6a 100644
--- a/src/core/hle/kernel/k_condition_variable.cpp
+++ b/src/core/hle/kernel/k_condition_variable.cpp
@@ -33,21 +33,26 @@ bool UpdateLockAtomic(Core::System& system, u32* out, KProcessAddress address, u
     auto& monitor = system.Monitor();
     const auto current_core = system.Kernel().CurrentPhysicalCoreIndex();
 
-    // Load the value from the address.
-    const auto expected = monitor.ExclusiveRead32(current_core, GetInteger(address));
+    u32 expected{};
 
-    // Orr in the new mask.
-    u32 value = expected | new_orr_mask;
+    while (true) {
+        // Load the value from the address.
+        expected = monitor.ExclusiveRead32(current_core, GetInteger(address));
 
-    // If the value is zero, use the if_zero value, otherwise use the newly orr'd value.
-    if (!expected) {
-        value = if_zero;
-    }
+        // Orr in the new mask.
+        u32 value = expected | new_orr_mask;
+
+        // If the value is zero, use the if_zero value, otherwise use the newly orr'd value.
+        if (!expected) {
+            value = if_zero;
+        }
+
+        // Try to store.
+        if (monitor.ExclusiveWrite32(current_core, GetInteger(address), value)) {
+            break;
+        }
 
-    // Try to store.
-    if (!monitor.ExclusiveWrite32(current_core, GetInteger(address), value)) {
         // If we failed to store, try again.
-        return UpdateLockAtomic(system, out, address, if_zero, new_orr_mask);
     }
 
     // We're done.
-- 
cgit v1.2.3