From c1d845aa0d891098e5e9b7026f1c271699b64b08 Mon Sep 17 00:00:00 2001
From: Joost Jager <joost.jager@gmail.com>
Date: Wed, 26 Sep 2018 09:46:48 -0700
Subject: [PATCH] cnct: reuse sweep tx logic for commit resolver

Removes duplicate sweep tx code from commit resolver
and delegates generation to UtxoSweeper in the sweep
package.
---
 contractcourt/chain_arbitrator.go   |  4 +++
 contractcourt/contract_resolvers.go | 49 ++++++-----------------------
 server.go                           |  1 +
 sweep/sweeper.go                    | 11 +++++--
 4 files changed, 23 insertions(+), 42 deletions(-)

diff --git a/contractcourt/chain_arbitrator.go b/contractcourt/chain_arbitrator.go
index db401735..5f8b161d 100644
--- a/contractcourt/chain_arbitrator.go
+++ b/contractcourt/chain_arbitrator.go
@@ -3,6 +3,7 @@ package contractcourt
 import (
 	"errors"
 	"fmt"
+	"github.com/lightningnetwork/lnd/sweep"
 	"sync"
 	"sync/atomic"
 
@@ -130,6 +131,9 @@ type ChainArbitratorConfig struct {
 	// DisableChannel disables a channel, resulting in it not being able to
 	// forward payments.
 	DisableChannel func(wire.OutPoint) error
+
+	// Sweeper allows resolvers to sweep their final outputs.
+	Sweeper *sweep.UtxoSweeper
 }
 
 // ChainArbitrator is a sub-system that oversees the on-chain resolution of all
diff --git a/contractcourt/contract_resolvers.go b/contractcourt/contract_resolvers.go
index 3ecc6be9..09a1cecf 100644
--- a/contractcourt/contract_resolvers.go
+++ b/contractcourt/contract_resolvers.go
@@ -5,6 +5,7 @@ import (
 	"crypto/sha256"
 	"encoding/binary"
 	"fmt"
+	"github.com/lightningnetwork/lnd/sweep"
 	"io"
 	"io/ioutil"
 
@@ -1256,46 +1257,16 @@ func (c *commitSweepResolver) Resolve() (ContractResolver, error) {
 	// If the sweep transaction isn't already generated, and the remote
 	// party broadcast the commitment transaction then we'll create it now.
 	case c.sweepTx == nil && !isLocalCommitTx:
-		// Now that the commitment transaction has confirmed, we'll
-		// craft a transaction to sweep this output into the wallet.
-		signDesc := c.commitResolution.SelfOutputSignDesc
+		input := sweep.MakeBaseInput(
+			&c.commitResolution.SelfOutPoint,
+			lnwallet.CommitmentNoDelay,
+			&c.commitResolution.SelfOutputSignDesc)
 
-		// First, we'll estimate the total weight so we can compute
-		// fees properly. We'll use a lax estimate, as this output is
-		// in no immediate danger.
-		feePerKw, err := c.FeeEstimator.EstimateFeePerKW(6)
-		if err != nil {
-			return nil, err
-		}
-
-		log.Debugf("%T(%v): using %v sat/kw for sweep tx", c,
-			c.chanPoint, int64(feePerKw))
-
-		totalWeight := (&lnwallet.TxWeightEstimator{}).
-			AddP2WKHInput().
-			AddP2WKHOutput().Weight()
-		totalFees := feePerKw.FeeForWeight(int64(totalWeight))
-		sweepAmt := signDesc.Output.Value - int64(totalFees)
-
-		c.sweepTx = wire.NewMsgTx(2)
-		c.sweepTx.AddTxIn(&wire.TxIn{
-			PreviousOutPoint: c.commitResolution.SelfOutPoint,
-		})
-		sweepAddr, err := c.NewSweepAddr()
-		if err != nil {
-			return nil, err
-		}
-		c.sweepTx.AddTxOut(&wire.TxOut{
-			PkScript: sweepAddr,
-			Value:    sweepAmt,
-		})
-
-		// With the transaction fully assembled, we can now generate a
-		// valid witness for the transaction.
-		signDesc.SigHashes = txscript.NewTxSigHashes(c.sweepTx)
-		c.sweepTx.TxIn[0].Witness, err = lnwallet.CommitSpendNoDelay(
-			c.Signer, &signDesc, c.sweepTx,
-		)
+		// TODO: Set tx lock time to current block height instead of
+		// zero. Will be taken care of once sweeper implementation is
+		// complete.
+		c.sweepTx, err = c.Sweeper.CreateSweepTx(
+			[]sweep.Input{&input}, 0)
 		if err != nil {
 			return nil, err
 		}
diff --git a/server.go b/server.go
index a3b52904..7ebe5189 100644
--- a/server.go
+++ b/server.go
@@ -695,6 +695,7 @@ func newServer(listenAddrs []net.Addr, chanDB *channeldb.DB, cc *chainControl,
 		DisableChannel: func(op wire.OutPoint) error {
 			return s.announceChanStatus(op, true)
 		},
+		Sweeper: sweeper,
 	}, chanDB)
 
 	s.breachArbiter = newBreachArbiter(&BreachConfig{
diff --git a/sweep/sweeper.go b/sweep/sweeper.go
index 7e2cc916..3a2197e3 100644
--- a/sweep/sweeper.go
+++ b/sweep/sweeper.go
@@ -73,11 +73,17 @@ func (s *UtxoSweeper) CreateSweepTx(inputs []Input,
 	// outputs.
 	csvCount := 0
 	cltvCount := 0
+	unknownCount := 0
 	for i := range inputs {
 		input := inputs[i]
 
 		switch input.WitnessType() {
 
+		// Outputs on a remote commitment transaction that pay directly
+		// to us.
+		case lnwallet.CommitmentNoDelay:
+			weightEstimate.AddP2WKHInput()
+
 		// Outputs on a past commitment transaction that pay directly
 		// to us.
 		case lnwallet.CommitmentTimeLock:
@@ -113,8 +119,7 @@ func (s *UtxoSweeper) CreateSweepTx(inputs []Input,
 			cltvCount++
 
 		default:
-			// TODO: Also add non-timelocked outputs
-
+			unknownCount++
 			log.Warnf("kindergarten output in nursery store "+
 				"contains unexpected witness type: %v",
 				input.WitnessType())
@@ -123,7 +128,7 @@ func (s *UtxoSweeper) CreateSweepTx(inputs []Input,
 	}
 
 	log.Infof("Creating sweep transaction for %v inputs (%v CSV, %v CLTV)",
-		csvCount+cltvCount, csvCount, cltvCount)
+		len(inputs)-unknownCount, csvCount, cltvCount)
 
 	txWeight := int64(weightEstimate.Weight())
 	return s.populateSweepTx(txWeight, currentBlockHeight, inputs)