From 616503de3e2974fafc5d398a1e6ab7ccd260f452 Mon Sep 17 00:00:00 2001
From: "Johan T. Halseth" <johanth@gmail.com>
Date: Mon, 16 Nov 2020 12:52:22 +0100
Subject: [PATCH] sweep: add required TxOuts to sweep tx

Now that inputs might have accompanied outputs to be added to the sweep
tx, we add them to the sweep transaction first, and account for it when
calculating the change amount.
---
 sweep/txgenerator.go | 65 +++++++++++++++++++++++++++++++++++++-------
 1 file changed, 55 insertions(+), 10 deletions(-)

diff --git a/sweep/txgenerator.go b/sweep/txgenerator.go
index 784a7ebd..670a2255 100644
--- a/sweep/txgenerator.go
+++ b/sweep/txgenerator.go
@@ -147,15 +147,23 @@ func createSweepTx(inputs []input.Input, outputPkScript []byte,
 	// Track whether any of the inputs require a certain locktime.
 	locktime := int32(-1)
 
-	// Sum up the total value contained in the inputs, and add all inputs
-	// to the sweep transaction. Ensure that for each csvInput, we set the
-	// sequence number properly.
-	var totalSum btcutil.Amount
+	// We start by adding all inputs that commit to an output. We do this
+	// since the input and output index must stay the same for the
+	// signatures to be valid.
+	var (
+		totalInput     btcutil.Amount
+		requiredOutput btcutil.Amount
+	)
 	for _, o := range inputs {
+		if o.RequiredTxOut() == nil {
+			continue
+		}
+
 		sweepTx.AddTxIn(&wire.TxIn{
 			PreviousOutPoint: *o.OutPoint(),
 			Sequence:         o.BlocksToMaturity(),
 		})
+		sweepTx.AddTxOut(o.RequiredTxOut())
 
 		if lt, ok := o.RequiredLockTime(); ok {
 			// If another input commits to a different locktime,
@@ -167,18 +175,44 @@ func createSweepTx(inputs []input.Input, outputPkScript []byte,
 			locktime = int32(lt)
 		}
 
-		totalSum += btcutil.Amount(o.SignDesc().Output.Value)
+		totalInput += btcutil.Amount(o.SignDesc().Output.Value)
+		requiredOutput += btcutil.Amount(o.RequiredTxOut().Value)
 	}
 
-	// Sweep as much possible, after subtracting txn fees.
-	sweepAmt := totalSum - txFee
+	// Sum up the value contained in the remaining inputs, and add them to
+	// the sweep transaction.
+	for _, o := range inputs {
+		if o.RequiredTxOut() != nil {
+			continue
+		}
+
+		sweepTx.AddTxIn(&wire.TxIn{
+			PreviousOutPoint: *o.OutPoint(),
+			Sequence:         o.BlocksToMaturity(),
+		})
+
+		if lt, ok := o.RequiredLockTime(); ok {
+			if locktime != -1 && locktime != int32(lt) {
+				return nil, fmt.Errorf("incompatible locktime")
+			}
+
+			locktime = int32(lt)
+		}
+
+		totalInput += btcutil.Amount(o.SignDesc().Output.Value)
+	}
+
+	// The value remaining after the required output and fees, go to
+	// change. Not that this fee is what we would have to pay in case the
+	// sweep tx has a change output.
+	changeAmt := totalInput - requiredOutput - txFee
 
 	// The txn will sweep the amount after fees to the pkscript generated
 	// above.
-	if sweepAmt >= dustLimit {
+	if changeAmt >= dustLimit {
 		sweepTx.AddTxOut(&wire.TxOut{
 			PkScript: outputPkScript,
-			Value:    int64(sweepAmt),
+			Value:    int64(changeAmt),
 		})
 	}
 
@@ -255,7 +289,12 @@ func getWeightEstimate(inputs []input.Input, feeRate chainfee.SatPerKWeight) (
 	weightEstimate := newWeightEstimator(feeRate)
 
 	// Our sweep transaction will pay to a single segwit p2wkh address,
-	// ensure it contributes to our weight estimate.
+	// ensure it contributes to our weight estimate. If the inputs we add
+	// have required TxOuts, then this will be our change address. Note
+	// that if we have required TxOuts, we might end up creating a sweep tx
+	// without a change output. It is okay to add the change output to the
+	// weight estimate regardless, since the estimated fee will just be
+	// subtracted from this already dust output, and trimmed.
 	weightEstimate.addP2WKHOutput()
 
 	// For each output, use its witness type to determine the estimate
@@ -274,6 +313,12 @@ func getWeightEstimate(inputs []input.Input, feeRate chainfee.SatPerKWeight) (
 			continue
 		}
 
+		// If this input comes with a committed output, add that as
+		// well.
+		if inp.RequiredTxOut() != nil {
+			weightEstimate.addOutput(inp.RequiredTxOut())
+		}
+
 		sweepInputs = append(sweepInputs, inp)
 	}