Check for sufficient balances at voting time

mystenmark · mystenmark · commit b04aa4a106b4 · 2025-11-20T08:49:40.000-08:00
diff --git a/crates/sui-core/src/accumulators/balance_read.rs b/crates/sui-core/src/accumulators/balance_read.rs
@@ -0,0 +1,82 @@
+// Copyright (c) Mysten Labs, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+use std::{collections::BTreeMap, sync::Arc};
+
+use sui_types::{
+    accumulator_root::{AccumulatorObjId, AccumulatorValue, U128},
+    base_types::SequenceNumber,
+    error::{SuiErrorKind, SuiResult, UserInputError},
+    storage::ChildObjectResolver,
+};
+
+pub(crate) trait AccountBalanceRead: Send + Sync {
+    fn get_account_balance(
+        &self,
+        account_id: &AccumulatorObjId,
+        // Version of the accumulator root object, used to
+        // bound the version when we look for child account objects.
+        accumulator_version: SequenceNumber,
+    ) -> u128;
+
+    /// Gets latest balance, without a version bound on the accumulator root object.
+    /// Only used for signing time checks / RPC reads, not scheduling.
+    fn get_latest_account_balance(&self, account_id: &AccumulatorObjId) -> u128;
+
+    /// Checks if balances are available in the latest versions of the referenced acccumulator
+    /// objects. This does un-sequenced reads and can only be used on the signing/voting path
+    /// where deterministic results are not required.
+    fn check_balances_available(
+        &self,
+        requested_balances: &BTreeMap<AccumulatorObjId, u64>,
+    ) -> SuiResult {
+        for (object_id, requested_balance) in requested_balances {
+            let actual_balance = self.get_latest_account_balance(object_id);
+
+            if actual_balance < *requested_balance as u128 {
+                return Err(SuiErrorKind::UserInputError {
+                    error: UserInputError::InvalidWithdrawReservation {
+                        error: format!(
+                            "Available balance for object id {} is less than requested: {} < {}",
+                            object_id, actual_balance, requested_balance
+                        ),
+                    },
+                }
+                .into());
+            }
+        }
+
+        Ok(())
+    }
+}
+
+impl AccountBalanceRead for Arc<dyn ChildObjectResolver + Send + Sync> {
+    fn get_account_balance(
+        &self,
+        account_id: &AccumulatorObjId,
+        accumulator_version: SequenceNumber,
+    ) -> u128 {
+        // TODO: The implementation currently relies on the fact that we could
+        // load older versions of child objects. This has two problems:
+        // 1. Aggressive pruning might prune old versions of child objects,
+        // 2. Tidehunter might not continue to support this kinds of reads.
+        // To fix this, we could also read the latest version of the accumulator root object,
+        // and see if the provided accumulator version is already settled.
+        let value: U128 =
+            AccumulatorValue::load_by_id(self.as_ref(), Some(accumulator_version), *account_id)
+                // Expect is safe because at this point we should know that we are dealing with a Balance<T>
+                // object
+                .expect("read cannot fail")
+                .unwrap_or(U128 { value: 0 });
+
+        value.value
+    }
+
+    fn get_latest_account_balance(&self, account_id: &AccumulatorObjId) -> u128 {
+        let value = AccumulatorValue::load_by_id(self.as_ref(), None, *account_id)
+            .expect("read cannot fail")
+            .unwrap_or(U128 { value: 0 });
+
+        value.value
+    }
+}
diff --git a/crates/sui-core/src/accumulators/mod.rs b/crates/sui-core/src/accumulators/mod.rs
@@ -29,6 +29,8 @@ use sui_types::{
 
 use crate::execution_cache::TransactionCacheRead;
 
+pub mod balance_read;
+
 /// Merged value is the value stored inside accumulator objects.
 /// Each mergeable Move type will map to a single variant as its representation.
 ///
diff --git a/crates/sui-core/src/authority.rs b/crates/sui-core/src/authority.rs
@@ -2,6 +2,7 @@
 // Copyright (c) Mysten Labs, Inc.
 // SPDX-License-Identifier: Apache-2.0
 
+use crate::accumulators::balance_read::AccountBalanceRead;
 use crate::checkpoints::CheckpointBuilderError;
 use crate::checkpoints::CheckpointBuilderResult;
 use crate::congestion_tracker::CongestionTracker;
@@ -1024,6 +1025,12 @@ impl AuthorityState {
             self.get_backing_package_store().as_ref(),
         )?;
 
+        let withdraws = tx_data.process_funds_withdrawals_for_signing()?;
+
+        self.execution_cache_trait_pointers
+            .child_object_resolver
+            .check_balances_available(&withdraws)?;
+
         let (input_objects, receiving_objects) = self.input_loader.read_objects_for_signing(
             Some(tx_digest),
             &input_object_kinds,
diff --git a/crates/sui-core/src/execution_scheduler/balance_withdraw_scheduler/eager_scheduler.rs b/crates/sui-core/src/execution_scheduler/balance_withdraw_scheduler/eager_scheduler.rs
@@ -14,10 +14,12 @@ use sui_types::{
 use tokio::sync::oneshot::Sender;
 use tracing::{debug, instrument};
 
-use crate::execution_scheduler::balance_withdraw_scheduler::{
-    BalanceSettlement, ScheduleResult, ScheduleStatus, TxBalanceWithdraw,
-    balance_read::AccountBalanceRead,
-    scheduler::{BalanceWithdrawSchedulerTrait, WithdrawReservations},
+use crate::{
+    accumulators::balance_read::AccountBalanceRead,
+    execution_scheduler::balance_withdraw_scheduler::{
+        BalanceSettlement, ScheduleResult, ScheduleStatus, TxBalanceWithdraw,
+        scheduler::{BalanceWithdrawSchedulerTrait, WithdrawReservations},
+    },
 };
 
 pub(crate) struct EagerBalanceWithdrawScheduler {
diff --git a/crates/sui-core/src/execution_scheduler/balance_withdraw_scheduler/mock_balance_read.rs b/crates/sui-core/src/execution_scheduler/balance_withdraw_scheduler/mock_balance_read.rs
@@ -1,67 +1,26 @@
 // Copyright (c) Mysten Labs, Inc.
 // SPDX-License-Identifier: Apache-2.0
 
-#[cfg(test)]
 use std::collections::BTreeMap;
 use std::sync::Arc;
 
-#[cfg(test)]
 use parking_lot::RwLock;
-#[cfg(test)]
 use sui_types::base_types::ObjectID;
-use sui_types::{
-    accumulator_root::{AccumulatorObjId, AccumulatorValue, U128},
-    base_types::SequenceNumber,
-    storage::ChildObjectResolver,
-};
+use sui_types::{accumulator_root::AccumulatorObjId, base_types::SequenceNumber};
 
-pub(crate) trait AccountBalanceRead: Send + Sync {
-    fn get_account_balance(
-        &self,
-        account_id: &AccumulatorObjId,
-        // Version of the accumulator root object, used to
-        // bound the version when we look for child account objects.
-        accumulator_version: SequenceNumber,
-    ) -> u128;
-}
-
-impl AccountBalanceRead for Arc<dyn ChildObjectResolver + Send + Sync> {
-    fn get_account_balance(
-        &self,
-        account_id: &AccumulatorObjId,
-        accumulator_version: SequenceNumber,
-    ) -> u128 {
-        // TODO: The implementation currently relies on the fact that we could
-        // load older versions of child objects. This has two problems:
-        // 1. Aggressive pruning might prune old versions of child objects,
-        // 2. Tidehunter might not continue to support this kinds of reads.
-        // To fix this, we could also read the latest version of the accumulator root object,
-        // and see if the provided accumulator version is already settled.
-        let value: U128 =
-            AccumulatorValue::load_by_id(self.as_ref(), Some(accumulator_version), *account_id)
-                // Expect is safe because at this point we should know that we are dealing with a Balance<T>
-                // object
-                .expect("read cannot fail")
-                .unwrap_or(U128 { value: 0 });
-
-        value.value
-    }
-}
+use crate::accumulators::balance_read::AccountBalanceRead;
 
 // Mock implementation of a balance account book for testing.
 // Allows setting the balance for a given account at different accumulator versions.
-#[cfg(test)]
 pub(crate) struct MockBalanceRead {
     inner: Arc<RwLock<MockBalanceReadInner>>,
 }
 
-#[cfg(test)]
 struct MockBalanceReadInner {
     cur_version: SequenceNumber,
     balances: BTreeMap<AccumulatorObjId, BTreeMap<SequenceNumber, Option<u128>>>,
 }
 
-#[cfg(test)]
 impl MockBalanceRead {
     pub(crate) fn new(
         init_version: SequenceNumber,
@@ -99,7 +58,6 @@ impl MockBalanceRead {
     }
 }
 
-#[cfg(test)]
 impl MockBalanceReadInner {
     fn settle_balance_changes(
         &mut self,
@@ -144,9 +102,13 @@ impl MockBalanceReadInner {
             .last()
             .and_then(|(_, balance)| *balance)
     }
+
+    fn get_latest_account_balance(&self, account_id: &AccumulatorObjId) -> Option<u128> {
+        let account_balances = self.balances.get(account_id)?;
+        account_balances.values().last().and_then(|b| *b)
+    }
 }
 
-#[cfg(test)]
 impl AccountBalanceRead for MockBalanceRead {
     /// Mimic the behavior of child object read.
     /// Find the balance for the given account at the max version
@@ -161,4 +123,11 @@ impl AccountBalanceRead for MockBalanceRead {
             .get_account_balance(account_id, accumulator_version)
             .unwrap_or_default()
     }
+
+    fn get_latest_account_balance(&self, account_id: &AccumulatorObjId) -> u128 {
+        let inner = self.inner.read();
+        inner
+            .get_latest_account_balance(account_id)
+            .unwrap_or_default()
+    }
 }
diff --git a/crates/sui-core/src/execution_scheduler/balance_withdraw_scheduler/mod.rs b/crates/sui-core/src/execution_scheduler/balance_withdraw_scheduler/mod.rs
@@ -7,7 +7,9 @@ use sui_types::{
     accumulator_root::AccumulatorObjId, base_types::SequenceNumber, digests::TransactionDigest,
 };
 
-mod balance_read;
+#[cfg(test)]
+pub(crate) mod mock_balance_read;
+
 mod eager_scheduler;
 mod naive_scheduler;
 pub(crate) mod scheduler;
diff --git a/crates/sui-core/src/execution_scheduler/balance_withdraw_scheduler/naive_scheduler.rs b/crates/sui-core/src/execution_scheduler/balance_withdraw_scheduler/naive_scheduler.rs
@@ -7,10 +7,12 @@ use sui_types::base_types::SequenceNumber;
 use tokio::sync::watch;
 use tracing::{debug, instrument};
 
-use crate::execution_scheduler::balance_withdraw_scheduler::{
-    BalanceSettlement, ScheduleResult, ScheduleStatus,
-    balance_read::AccountBalanceRead,
-    scheduler::{BalanceWithdrawSchedulerTrait, WithdrawReservations},
+use crate::{
+    accumulators::balance_read::AccountBalanceRead,
+    execution_scheduler::balance_withdraw_scheduler::{
+        BalanceSettlement, ScheduleResult, ScheduleStatus,
+        scheduler::{BalanceWithdrawSchedulerTrait, WithdrawReservations},
+    },
 };
 
 /// A naive implementation of the balance withdraw scheduler that does not attempt to optimize the scheduling.
diff --git a/crates/sui-core/src/execution_scheduler/balance_withdraw_scheduler/scheduler.rs b/crates/sui-core/src/execution_scheduler/balance_withdraw_scheduler/scheduler.rs
@@ -3,10 +3,13 @@
 
 use std::{collections::BTreeMap, sync::Arc};
 
-use crate::execution_scheduler::balance_withdraw_scheduler::{
-    BalanceSettlement, ScheduleResult, ScheduleStatus, TxBalanceWithdraw,
-    balance_read::AccountBalanceRead, eager_scheduler::EagerBalanceWithdrawScheduler,
-    naive_scheduler::NaiveBalanceWithdrawScheduler,
+use crate::{
+    accumulators::balance_read::AccountBalanceRead,
+    execution_scheduler::balance_withdraw_scheduler::{
+        BalanceSettlement, ScheduleResult, ScheduleStatus, TxBalanceWithdraw,
+        eager_scheduler::EagerBalanceWithdrawScheduler,
+        naive_scheduler::NaiveBalanceWithdrawScheduler,
+    },
 };
 use futures::stream::FuturesUnordered;
 use mysten_metrics::monitored_mpsc::{UnboundedReceiver, UnboundedSender, unbounded_channel};
diff --git a/crates/sui-core/src/execution_scheduler/balance_withdraw_scheduler/tests.rs b/crates/sui-core/src/execution_scheduler/balance_withdraw_scheduler/tests.rs
@@ -3,7 +3,7 @@
 
 use crate::execution_scheduler::balance_withdraw_scheduler::ScheduleResult;
 use crate::execution_scheduler::balance_withdraw_scheduler::{
-    BalanceSettlement, ScheduleStatus, TxBalanceWithdraw, balance_read::MockBalanceRead,
+    BalanceSettlement, ScheduleStatus, TxBalanceWithdraw, mock_balance_read::MockBalanceRead,
     scheduler::BalanceWithdrawScheduler,
 };
 use futures::StreamExt;
diff --git a/crates/sui-core/src/execution_scheduler/execution_scheduler_impl.rs b/crates/sui-core/src/execution_scheduler/execution_scheduler_impl.rs
@@ -316,7 +316,7 @@ impl ExecutionScheduler {
         for (cert, env) in &certs {
             let tx_withdraws = cert
                 .transaction_data()
-                .process_funds_withdrawals()
+                .process_funds_withdrawals_for_signing()
                 .expect("Balance withdraws should have already been checked");
             assert!(!tx_withdraws.is_empty());
             let accumulator_version = env
diff --git a/crates/sui-e2e-tests/tests/address_balance_tests.rs b/crates/sui-e2e-tests/tests/address_balance_tests.rs
diff --git a/crates/sui-types/src/accumulator_root.rs b/crates/sui-types/src/accumulator_root.rs
diff --git a/crates/sui-types/src/transaction.rs b/crates/sui-types/src/transaction.rs
diff --git a/crates/sui-types/src/unit_tests/balance_withdraw_tests.rs b/crates/sui-types/src/unit_tests/balance_withdraw_tests.rs
