From e5cbd82343a534bb7fea62b3de908a415429746d Mon Sep 17 00:00:00 2001
From: Daniil Polyakov <arjentix@gmail.com>
Date: Sat, 6 Dec 2025 21:37:51 +0300
Subject: [PATCH] fix: add overflow check for balance sum validation

---
 nssa/core/src/program.rs | 42 ++++++++++++++++++++++++++++++++++++++--
 1 file changed, 40 insertions(+), 2 deletions(-)

diff --git a/nssa/core/src/program.rs b/nssa/core/src/program.rs
index 2d38fa4..04e91c1 100644
--- a/nssa/core/src/program.rs
+++ b/nssa/core/src/program.rs
@@ -204,8 +204,19 @@ pub fn validate_execution(
     }
 
     // 7. Total balance is preserved
-    let total_balance_pre_states: u128 = pre_states.iter().map(|pre| pre.account.balance).sum();
-    let total_balance_post_states: u128 = post_states.iter().map(|post| post.account.balance).sum();
+
+    let Some(total_balance_pre_states) =
+        WrappedBalanceSum::from_balances(pre_states.iter().map(|pre| pre.account.balance))
+    else {
+        return false;
+    };
+
+    let Some(total_balance_post_states) =
+        WrappedBalanceSum::from_balances(post_states.iter().map(|post| post.account.balance))
+    else {
+        return false;
+    };
+
     if total_balance_pre_states != total_balance_post_states {
         return false;
     }
@@ -213,6 +224,33 @@ pub fn validate_execution(
     true
 }
 
+/// Representation of a number as `lo + hi * 2^128`.
+#[derive(PartialEq, Eq)]
+struct WrappedBalanceSum {
+    lo: u128,
+    hi: u128,
+}
+
+impl WrappedBalanceSum {
+    /// Constructs a [`WrappedBalanceSum`] from an iterator of balances.
+    ///
+    /// Returns [`None`] if balance sum overflows `lo + hi * 2^128` representation, which is not
+    /// expected in practical scenarios.
+    fn from_balances(balances: impl Iterator<Item = u128>) -> Option<Self> {
+        let mut wrapped = WrappedBalanceSum { lo: 0, hi: 0 };
+
+        for balance in balances {
+            let (new_sum, did_overflow) = wrapped.lo.overflowing_add(balance);
+            if did_overflow {
+                wrapped.hi = wrapped.hi.checked_add(1)?;
+            }
+            wrapped.lo = new_sum;
+        }
+
+        Some(wrapped)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;