From c94130328a14b6524af19d6e53cc3824e7c24735 Mon Sep 17 00:00:00 2001
From: Jim Posen <jimpo@coinbase.com>
Date: Mon, 2 Oct 2017 18:52:45 -0700
Subject: [PATCH] lnwallet: Extend Utxo struct with AddressType.

The Utxo struct now includes the address type and redeem/witness
scripts. This is necessary for accurate fee estimation.
---
 lnwallet/btcwallet/btcwallet.go | 21 ++++++++++++++++-----
 lnwallet/interface.go           | 12 ++++++++++--
 mock.go                         |  4 +++-
 3 files changed, 29 insertions(+), 8 deletions(-)

diff --git a/lnwallet/btcwallet/btcwallet.go b/lnwallet/btcwallet/btcwallet.go
index 1650cb61..d4209d15 100644
--- a/lnwallet/btcwallet/btcwallet.go
+++ b/lnwallet/btcwallet/btcwallet.go
@@ -353,17 +353,28 @@ func (b *BtcWallet) ListUnspentWitness(minConfs int32) ([]*lnwallet.Utxo, error)
 			return nil, err
 		}
 
-		// TODO(roasbeef): this assumes all p2sh outputs returned by
-		// the wallet are nested p2sh...
-		if txscript.IsPayToWitnessPubKeyHash(pkScript) ||
-			txscript.IsPayToScriptHash(pkScript) {
+		var addressType lnwallet.AddressType
+		if txscript.IsPayToWitnessPubKeyHash(pkScript) {
+			addressType = lnwallet.WitnessPubKey
+		} else if txscript.IsPayToScriptHash(pkScript) {
+			// TODO(roasbeef): This assumes all p2sh outputs returned by the
+			// wallet are nested p2pkh. We can't check the redeem script because
+			// the btcwallet service does not include it.
+			addressType = lnwallet.NestedWitnessPubKey
+		}
+
+		if addressType == lnwallet.WitnessPubKey ||
+			addressType == lnwallet.NestedWitnessPubKey {
+
 			txid, err := chainhash.NewHashFromStr(output.TxID)
 			if err != nil {
 				return nil, err
 			}
 
 			utxo := &lnwallet.Utxo{
-				Value: btcutil.Amount(output.Amount * 1e8),
+				AddressType: addressType,
+				Value:       btcutil.Amount(output.Amount * 1e8),
+				PkScript:    pkScript,
 				OutPoint: wire.OutPoint{
 					Hash:  *txid,
 					Index: output.Vout,
diff --git a/lnwallet/interface.go b/lnwallet/interface.go
index 94d4d0de..d6ac3827 100644
--- a/lnwallet/interface.go
+++ b/lnwallet/interface.go
@@ -20,8 +20,12 @@ var ErrNotMine = errors.New("the passed output doesn't belong to the wallet")
 type AddressType uint8
 
 const (
+	// UnknownAddressType represents an output with an unknown or non-standard
+	// script.
+	UnknownAddressType AddressType = iota
+
 	// WitnessPubKey represents a p2wkh address.
-	WitnessPubKey AddressType = iota
+	WitnessPubKey
 
 	// NestedWitnessPubKey represents a p2sh output which is itself a
 	// nested p2wkh output.
@@ -34,7 +38,11 @@ const (
 // Utxo is an unspent output denoted by its outpoint, and output value of the
 // original output.
 type Utxo struct {
-	Value btcutil.Amount
+	AddressType   AddressType
+	Value         btcutil.Amount
+	PkScript      []byte
+	RedeemScript  []byte
+	WitnessScript []byte
 	wire.OutPoint
 }
 
diff --git a/mock.go b/mock.go
index b6c194a3..63e4893d 100644
--- a/mock.go
+++ b/mock.go
@@ -149,7 +149,9 @@ func (*mockWalletController) SendOutputs(outputs []*wire.TxOut) (*chainhash.Hash
 // need one unspent for the funding transaction.
 func (*mockWalletController) ListUnspentWitness(confirms int32) ([]*lnwallet.Utxo, error) {
 	utxo := &lnwallet.Utxo{
-		Value: btcutil.Amount(10 * btcutil.SatoshiPerBitcoin),
+		AddressType: lnwallet.WitnessPubKey,
+		Value:       btcutil.Amount(10 * btcutil.SatoshiPerBitcoin),
+		PkScript:    make([]byte, 22),
 		OutPoint: wire.OutPoint{
 			Hash:  chainhash.Hash{},
 			Index: 0,