wallet: fix wallet Redistribute when outputs is larger than batchsize

SiaFoundation · Feb 10, 2025 · e99e9fa · e99e9fa
1 parent 8e9f80e
commit e99e9fa
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 5 deletions.
diff --git a/wallet/wallet.go b/wallet/wallet.go
@@ -629,7 +629,7 @@ func (sw *SingleAddressWallet) selectRedistributeUTXOs(bh uint64, outputs int, a
 // Redistribute returns a transaction that redistributes money in the wallet by
 // selecting a minimal set of inputs to cover the creation of the requested
 // outputs. It also returns a list of output IDs that need to be signed.
-func (sw *SingleAddressWallet) Redistribute(outputs int, amount, feePerByte types.Currency) (txns []types.Transaction, toSign []types.Hash256, err error) {
+func (sw *SingleAddressWallet) Redistribute(outputs int, amount, feePerByte types.Currency) (txns []types.Transaction, toSign [][]types.Hash256, err error) {
 	state := sw.cm.TipState()
 
 	elements, err := sw.store.UnspentSiacoinElements()
@@ -653,8 +653,10 @@ func (sw *SingleAddressWallet) Redistribute(outputs int, amount, feePerByte type
 	// in case of an error we need to free all inputs
 	defer func() {
 		if err != nil {
-			for _, id := range toSign {
-				delete(sw.locked, types.SiacoinOutputID(id))
+			for _, ids := range toSign {
+				for _, id := range ids {
+					delete(sw.locked, types.SiacoinOutputID(id))
+				}
 			}
 		}
 	}()
@@ -690,6 +692,9 @@ func (sw *SingleAddressWallet) Redistribute(outputs int, amount, feePerByte type
 			}
 		}
 
+		// remove used inputs from utxos
+		utxos = utxos[len(inputs):]
+
 		// not enough outputs found
 		fee := feePerInput.Mul64(uint64(len(inputs))).Add(outputFees)
 		if sumOut := SumOutputs(inputs); sumOut.Cmp(want.Add(fee)) < 0 {
@@ -711,15 +716,17 @@ func (sw *SingleAddressWallet) Redistribute(outputs int, amount, feePerByte type
 		}
 
 		// add the inputs
+		toSignTxn := make([]types.Hash256, 0, len(inputs))
 		for _, sce := range inputs {
+			toSignTxn = append(toSignTxn, types.Hash256(sce.ID))
 			txn.SiacoinInputs = append(txn.SiacoinInputs, types.SiacoinInput{
 				ParentID:         types.SiacoinOutputID(sce.ID),
 				UnlockConditions: types.StandardUnlockConditions(sw.priv.PublicKey()),
 			})
-			toSign = append(toSign, types.Hash256(sce.ID))
 			sw.locked[sce.ID] = time.Now().Add(sw.cfg.ReservationDuration)
 		}
 		txns = append(txns, txn)
+		toSign = append(toSign, toSignTxn)
 	}
 
 	return
@@ -786,6 +793,9 @@ func (sw *SingleAddressWallet) RedistributeV2(outputs int, amount, feePerByte ty
 			}
 		}
 
+		// remove used inputs from utxos
+		utxos = utxos[len(inputs):]
+
 		// not enough outputs found
 		fee := feePerInput.Mul64(uint64(len(inputs))).Add(outputFees)
 		if sumOut := SumOutputs(inputs); sumOut.Cmp(want.Add(fee)) < 0 {

diff --git a/wallet/wallet_test.go b/wallet/wallet_test.go
@@ -500,7 +500,7 @@ func TestWalletRedistribute(t *testing.T) {
 		}
 
 		for i := 0; i < len(txns); i++ {
-			w.SignTransaction(&txns[i], toSign, types.CoveredFields{WholeTransaction: true})
+			w.SignTransaction(&txns[i], toSign[i], types.CoveredFields{WholeTransaction: true})
 		}
 		if _, err := cm.AddPoolTransactions(txns); err != nil {
 			return fmt.Errorf("failed to add transactions to pool: %w", err)
@@ -557,6 +557,18 @@ func TestWalletRedistribute(t *testing.T) {
 	} else if len(toSign) != 0 {
 		t.Fatalf("expected no ids, got %v", len(toSign))
 	}
+
+	// redistribute the wallet into more outputs than the batch size to make
+	// sure the resulting txn set contains more than 1 txn
+	outputs, err := w.SpendableOutputs()
+	if err != nil {
+		t.Fatal(err)
+	} else if len(outputs) >= 11 {
+		t.Fatalf("expected at least 11 outputs, got %v", len(outputs))
+	} else if err := redistribute(types.Siacoins(1e3), 11); err != nil {
+		t.Fatal(err)
+	}
+
 }
 
 func TestWalletRedistributeV2(t *testing.T) {
@@ -652,6 +664,18 @@ func TestWalletRedistributeV2(t *testing.T) {
 	} else if len(toSign) != 0 {
 		t.Fatalf("expected no ids, got %v", len(toSign))
 	}
+
+	// redistribute the wallet into more outputs than the batch size to make
+	// sure the resulting txn set contains more than 1 txn
+	outputs, err := w.SpendableOutputs()
+	if err != nil {
+		t.Fatal(err)
+	} else if len(outputs) >= 11 {
+		t.Fatalf("expected at least 11 outputs, got %v", len(outputs))
+	} else if err := redistribute(types.Siacoins(1e3), 11); err != nil {
+		t.Fatal(err)
+	}
+
 }
 
 func TestReorg(t *testing.T) {