From 89b5f1918ea8c94bdef385d296d61aee60e7f81b Mon Sep 17 00:00:00 2001
From: yukang <moorekang@gmail.com>
Date: Sat, 19 Oct 2024 20:42:06 +0800
Subject: [PATCH] fix tx_builder capacity overflow issue

---
 src/tests/mod.rs      | 27 ++++++++++++++++++++++++++-
 src/tx_builder/mod.rs |  7 ++++++-
 2 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/src/tests/mod.rs b/src/tests/mod.rs
index 1fc11293..ac3d50b6 100644
--- a/src/tests/mod.rs
+++ b/src/tests/mod.rs
@@ -1,4 +1,4 @@
-use std::collections::HashMap;
+use std::{collections::HashMap, u64};
 
 use ckb_dao_utils::pack_dao_data;
 use ckb_hash::blake2b_256;
@@ -182,6 +182,31 @@ fn test_transfer_from_sighash() {
     ctx.verify(tx, FEE_RATE).unwrap();
 }
 
+#[test]
+fn test_transfer_capacity_overflow() {
+    let sender = build_sighash_script(ACCOUNT1_ARG);
+    let receiver = build_sighash_script(ACCOUNT2_ARG);
+    let ctx = init_context(Vec::new(), vec![(sender.clone(), Some(100 * ONE_CKB))]);
+
+    let large_amount: u64 = u64::MAX;
+    let output = CellOutput::new_builder()
+        .capacity((large_amount).pack())
+        .lock(receiver)
+        .build();
+    let builder = CapacityTransferBuilder::new(vec![(output.clone(), Bytes::default())]);
+    let placeholder_witness = WitnessArgs::new_builder()
+        .lock(Some(Bytes::from(vec![0u8; 65])).pack())
+        .build();
+    let balancer =
+        CapacityBalancer::new_simple(sender.clone(), placeholder_witness.clone(), FEE_RATE);
+
+    let unlockers: HashMap<ScriptId, Box<dyn ScriptUnlocker>> = HashMap::default();
+    let mut cell_collector = ctx.to_live_cells_context();
+    let res = builder.build_unlocked(&mut cell_collector, &ctx, &ctx, &ctx, &balancer, &unlockers);
+    assert!(res.is_err());
+    assert!(res.unwrap_err().to_string().contains("capacity not enough"));
+}
+
 #[test]
 fn test_transfer_from_multisig() {
     let lock_args = vec![
diff --git a/src/tx_builder/mod.rs b/src/tx_builder/mod.rs
index 0567f5b5..d6482f85 100644
--- a/src/tx_builder/mod.rs
+++ b/src/tx_builder/mod.rs
@@ -866,7 +866,12 @@ fn rebalance_tx_capacity(
                 need_more_capacity = min_fee - fee;
             }
             Err(TransactionFeeError::CapacityOverflow(delta)) => {
-                need_more_capacity = delta + min_fee;
+                need_more_capacity = delta.checked_add(min_fee).ok_or_else(|| {
+                    BalanceTxCapacityError::CapacityNotEnough(format!(
+                        "need more capacity, value={}",
+                        HumanCapacity(delta)
+                    ))
+                })?;
             }
             Err(err) => {
                 return Err(err.into());