input+sweep: rework witness type into an interface

breacharbiter.go

@@ -821,7 +821,7 @@ func (b *breachArbiter) handleBreachHandoff(breachEvent *ContractBreachEvent) {
 type breachedOutput struct {
 	amt         btcutil.Amount
 	outpoint    wire.OutPoint
-	witnessType input.WitnessType
+	witnessType input.StandardWitnessType
 	signDesc    input.SignDescriptor
 	confHeight  uint32
 
@@ -833,7 +833,7 @@ type breachedOutput struct {
 // makeBreachedOutput assembles a new breachedOutput that can be used by the
 // breach arbiter to construct a justice or sweep transaction.
 func makeBreachedOutput(outpoint *wire.OutPoint,
-	witnessType input.WitnessType,
+	witnessType input.StandardWitnessType,
 	secondLevelScript []byte,
 	signDescriptor *input.SignDescriptor,
 	confHeight uint32) breachedOutput {
@@ -883,9 +883,7 @@ func (bo *breachedOutput) CraftInputScript(signer input.Signer, txn *wire.MsgTx,
 
 	// First, we ensure that the witness generation function has been
 	// initialized for this breached output.
-	bo.witnessFunc = bo.witnessType.GenWitnessFunc(
+	bo.witnessFunc = bo.witnessType.WitnessGenerator(signer, bo.SignDesc())
-		signer, bo.SignDesc(),
-	)
 
 	// Now that we have ensured that the witness generation function has
 	// been initialized, we can proceed to execute it and generate the
@@ -993,7 +991,7 @@ func newRetributionInfo(chanPoint *wire.OutPoint,
 		// Using the breachedHtlc's incoming flag, determine the
 		// appropriate witness type that needs to be generated in order
 		// to sweep the HTLC output.
-		var htlcWitnessType input.WitnessType
+		var htlcWitnessType input.StandardWitnessType
 		if breachedHtlc.IsIncoming {
 			htlcWitnessType = input.HtlcAcceptedRevoke
 		} else {
@@ -1051,32 +1049,15 @@ func (b *breachArbiter) createJusticeTx(
 		// Grab locally scoped reference to breached output.
 		inp := &r.breachedOutputs[i]
 
-		// First, select the appropriate estimated witness weight for
+		// First, determine the appropriate estimated witness weight for
 		// the give witness type of this breached output. If the witness
-		// type is unrecognized, we will omit it from the transaction.
+		// weight cannot be estimated, we will omit it from the
-		var witnessWeight int
+		// transaction.
-		switch inp.WitnessType() {
+		witnessWeight, _, err := inp.WitnessType().SizeUpperBound()
-		case input.CommitSpendNoDelayTweakless:
+		if err != nil {
-			fallthrough
+			brarLog.Warnf("could not determine witness weight "+
-		case input.CommitmentNoDelay:
+				"for breached output in retribution info: %v",
-			witnessWeight = input.P2WKHWitnessSize
+				err)
-
-		case input.CommitmentRevoke:
-			witnessWeight = input.ToLocalPenaltyWitnessSize
-
-		case input.HtlcOfferedRevoke:
-			witnessWeight = input.OfferedHtlcPenaltyWitnessSize
-
-		case input.HtlcAcceptedRevoke:
-			witnessWeight = input.AcceptedHtlcPenaltyWitnessSize
-
-		case input.HtlcSecondLevelRevoke:
-			witnessWeight = input.ToLocalPenaltyWitnessSize
-
-		default:
-			brarLog.Warnf("breached output in retribution info "+
-				"contains unexpected witness type: %v",
-				inp.WitnessType())
 			continue
 		}
 		weightEstimate.AddWitnessInput(witnessWeight)
@@ -1555,7 +1536,7 @@ func (bo *breachedOutput) Decode(r io.Reader) error {
 	if _, err := io.ReadFull(r, scratch[:2]); err != nil {
 		return err
 	}
-	bo.witnessType = input.WitnessType(
+	bo.witnessType = input.StandardWitnessType(
 		binary.BigEndian.Uint16(scratch[:2]),
 	)
 

input/input.go

@@ -114,9 +114,7 @@ func NewBaseInput(outpoint *wire.OutPoint, witnessType WitnessType,
 func (bi *BaseInput) CraftInputScript(signer Signer, txn *wire.MsgTx,
 	hashCache *txscript.TxSigHashes, txinIdx int) (*Script, error) {
 
-	witnessFunc := bi.witnessType.GenWitnessFunc(
+	witnessFunc := bi.witnessType.WitnessGenerator(signer, bi.SignDesc())
-		signer, bi.SignDesc(),
-	)
 
 	return witnessFunc(txn, hashCache, txinIdx)
 }

input/witnessgen.go

@@ -7,87 +7,125 @@ import (
 	"github.com/btcsuite/btcd/wire"
 )
 
-// WitnessType determines how an output's witness will be generated. The
+// WitnessGenerator represents a function that is able to generate the final
-// default commitmentTimeLock type will generate a witness that will allow
+// witness for a particular public key script. Additionally, if required, this
-// spending of a time-locked transaction enforced by CheckSequenceVerify.
+// function will also return the sigScript for spending nested P2SH witness
-type WitnessType uint16
+// outputs. This function acts as an abstraction layer, hiding the details of
+// the underlying script.
+type WitnessGenerator func(tx *wire.MsgTx, hc *txscript.TxSigHashes,
+	inputIndex int) (*Script, error)
+
+// WitnessType determines how an output's witness will be generated. This
+// interface can be implemented to be used for custom sweep scripts if the
+// pre-defined StandardWitnessType list doesn't provide a suitable one.
+type WitnessType interface {
+	// String returns a human readable version of the WitnessType.
+	String() string
+
+	// WitnessGenerator will return a WitnessGenerator function that an
+	// output uses to generate the witness and optionally the sigScript for
+	// a sweep transaction.
+	WitnessGenerator(signer Signer,
+		descriptor *SignDescriptor) WitnessGenerator
+
+	// SizeUpperBound returns the maximum length of the witness of this
+	// WitnessType if it would be included in a tx. It also returns if the
+	// output itself is a nested p2sh output, if so then we need to take
+	// into account the extra sigScript data size.
+	SizeUpperBound() (int, bool, error)
+
+	// AddWeightEstimation adds the estimated size of the witness in bytes
+	// to the given weight estimator and returns the number of
+	// CSVs/CLTVs used by the script.
+	AddWeightEstimation(estimator *TxWeightEstimator) (int, int, error)
+}
+
+// StandardWitnessType is a numeric representation of standard pre-defined types
+// of witness configurations.
+type StandardWitnessType uint16
+
+// A compile time check to ensure StandardWitnessType implements the
+// WitnessType interface.
+var _ WitnessType = (StandardWitnessType)(0)
 
 const (
 	// CommitmentTimeLock is a witness that allows us to spend the output of
 	// a commitment transaction after a relative lock-time lockout.
-	CommitmentTimeLock WitnessType = 0
+	CommitmentTimeLock StandardWitnessType = 0
 
 	// CommitmentNoDelay is a witness that allows us to spend a settled
 	// no-delay output immediately on a counterparty's commitment
 	// transaction.
-	CommitmentNoDelay WitnessType = 1
+	CommitmentNoDelay StandardWitnessType = 1
 
 	// CommitmentRevoke is a witness that allows us to sweep the settled
 	// output of a malicious counterparty's who broadcasts a revoked
 	// commitment transaction.
-	CommitmentRevoke WitnessType = 2
+	CommitmentRevoke StandardWitnessType = 2
 
 	// HtlcOfferedRevoke is a witness that allows us to sweep an HTLC which
 	// we offered to the remote party in the case that they broadcast a
 	// revoked commitment state.
-	HtlcOfferedRevoke WitnessType = 3
+	HtlcOfferedRevoke StandardWitnessType = 3
 
 	// HtlcAcceptedRevoke is a witness that allows us to sweep an HTLC
 	// output sent to us in the case that the remote party broadcasts a
 	// revoked commitment state.
-	HtlcAcceptedRevoke WitnessType = 4
+	HtlcAcceptedRevoke StandardWitnessType = 4
 
 	// HtlcOfferedTimeoutSecondLevel is a witness that allows us to sweep
 	// an HTLC output that we extended to a party, but was never fulfilled.
 	// This HTLC output isn't directly on the commitment transaction, but
 	// is the result of a confirmed second-level HTLC transaction. As a
 	// result, we can only spend this after a CSV delay.
-	HtlcOfferedTimeoutSecondLevel WitnessType = 5
+	HtlcOfferedTimeoutSecondLevel StandardWitnessType = 5
 
 	// HtlcAcceptedSuccessSecondLevel is a witness that allows us to sweep
 	// an HTLC output that was offered to us, and for which we have a
 	// payment preimage. This HTLC output isn't directly on our commitment
 	// transaction, but is the result of confirmed second-level HTLC
 	// transaction. As a result, we can only spend this after a CSV delay.
-	HtlcAcceptedSuccessSecondLevel WitnessType = 6
+	HtlcAcceptedSuccessSecondLevel StandardWitnessType = 6
 
 	// HtlcOfferedRemoteTimeout is a witness that allows us to sweep an
 	// HTLC that we offered to the remote party which lies in the
 	// commitment transaction of the remote party. We can spend this output
 	// after the absolute CLTV timeout of the HTLC as passed.
-	HtlcOfferedRemoteTimeout WitnessType = 7
+	HtlcOfferedRemoteTimeout StandardWitnessType = 7
 
 	// HtlcAcceptedRemoteSuccess is a witness that allows us to sweep an
 	// HTLC that was offered to us by the remote party. We use this witness
 	// in the case that the remote party goes to chain, and we know the
 	// pre-image to the HTLC. We can sweep this without any additional
 	// timeout.
-	HtlcAcceptedRemoteSuccess WitnessType = 8
+	HtlcAcceptedRemoteSuccess StandardWitnessType = 8
 
 	// HtlcSecondLevelRevoke is a witness that allows us to sweep an HTLC
 	// from the remote party's commitment transaction in the case that the
 	// broadcast a revoked commitment, but then also immediately attempt to
 	// go to the second level to claim the HTLC.
-	HtlcSecondLevelRevoke WitnessType = 9
+	HtlcSecondLevelRevoke StandardWitnessType = 9
 
 	// WitnessKeyHash is a witness type that allows us to spend a regular
 	// p2wkh output that's sent to an output which is under complete
 	// control of the backing wallet.
-	WitnessKeyHash WitnessType = 10
+	WitnessKeyHash StandardWitnessType = 10
 
 	// NestedWitnessKeyHash is a witness type that allows us to sweep an
 	// output that sends to a nested P2SH script that pays to a key solely
 	// under our control. The witness generated needs to include the
-	NestedWitnessKeyHash WitnessType = 11
+	NestedWitnessKeyHash StandardWitnessType = 11
 
 	// CommitSpendNoDelayTweakless is similar to the CommitSpendNoDelay
 	// type, but it omits the tweak that randomizes the key we need to
 	// spend with a channel peer supplied set of randomness.
-	CommitSpendNoDelayTweakless = 12
+	CommitSpendNoDelayTweakless StandardWitnessType = 12
 )
 
-// Stirng returns a human readable version of the target WitnessType.
+// String returns a human readable version of the target WitnessType.
-func (wt WitnessType) String() string {
+//
+// NOTE: This is part of the WitnessType interface.
+func (wt StandardWitnessType) String() string {
 	switch wt {
 	case CommitmentTimeLock:
 		return "CommitmentTimeLock"
@@ -127,19 +165,13 @@ func (wt WitnessType) String() string {
 	}
 }
 
-// WitnessGenerator represents a function which is able to generate the final
+// WitnessGenerator will return a WitnessGenerator function that an output uses
-// witness for a particular public key script. Additionally, if required, this
-// function will also return the sigScript for spending nested P2SH witness
-// outputs. This function acts as an abstraction layer, hiding the details of
-// the underlying script.
-type WitnessGenerator func(tx *wire.MsgTx, hc *txscript.TxSigHashes,
-	inputIndex int) (*Script, error)
-
-// GenWitnessFunc will return a WitnessGenerator function that an output uses
 // to generate the witness and optionally the sigScript for a sweep
 // transaction. The sigScript will be generated if the witness type warrants
 // one for spending, such as the NestedWitnessKeyHash witness type.
-func (wt WitnessType) GenWitnessFunc(signer Signer,
+//
+// NOTE: This is part of the WitnessType interface.
+func (wt StandardWitnessType) WitnessGenerator(signer Signer,
 	descriptor *SignDescriptor) WitnessGenerator {
 
 	return func(tx *wire.MsgTx, hc *txscript.TxSigHashes,
@@ -262,5 +294,116 @@ func (wt WitnessType) GenWitnessFunc(signer Signer,
 			return nil, fmt.Errorf("unknown witness type: %v", wt)
 		}
 	}
-
+}
+
+// SizeUpperBound returns the maximum length of the witness of this witness
+// type if it would be included in a tx. We also return if the output itself is
+// a nested p2sh output, if so then we need to take into account the extra
+// sigScript data size.
+//
+// NOTE: This is part of the WitnessType interface.
+func (wt StandardWitnessType) SizeUpperBound() (int, bool, error) {
+	switch wt {
+
+	// Outputs on a remote commitment transaction that pay directly to us.
+	case CommitSpendNoDelayTweakless:
+		fallthrough
+	case WitnessKeyHash:
+		fallthrough
+	case CommitmentNoDelay:
+		return P2WKHWitnessSize, false, nil
+
+	// Outputs on a past commitment transaction that pay directly
+	// to us.
+	case CommitmentTimeLock:
+		return ToLocalTimeoutWitnessSize, false, nil
+
+	// Outgoing second layer HTLC's that have confirmed within the
+	// chain, and the output they produced is now mature enough to
+	// sweep.
+	case HtlcOfferedTimeoutSecondLevel:
+		return ToLocalTimeoutWitnessSize, false, nil
+
+	// Incoming second layer HTLC's that have confirmed within the
+	// chain, and the output they produced is now mature enough to
+	// sweep.
+	case HtlcAcceptedSuccessSecondLevel:
+		return ToLocalTimeoutWitnessSize, false, nil
+
+	// An HTLC on the commitment transaction of the remote party,
+	// that has had its absolute timelock expire.
+	case HtlcOfferedRemoteTimeout:
+		return AcceptedHtlcTimeoutWitnessSize, false, nil
+
+	// An HTLC on the commitment transaction of the remote party,
+	// that can be swept with the preimage.
+	case HtlcAcceptedRemoteSuccess:
+		return OfferedHtlcSuccessWitnessSize, false, nil
+
+	// A nested P2SH input that has a p2wkh witness script. We'll mark this
+	// as nested P2SH so the caller can estimate the weight properly
+	// including the sigScript.
+	case NestedWitnessKeyHash:
+		return P2WKHWitnessSize, true, nil
+
+	// The revocation output on a revoked commitment transaction.
+	case CommitmentRevoke:
+		return ToLocalPenaltyWitnessSize, false, nil
+
+	// The revocation output on a revoked HTLC that we offered to the remote
+	// party.
+	case HtlcOfferedRevoke:
+		return OfferedHtlcPenaltyWitnessSize, false, nil
+
+	// The revocation output on a revoked HTLC that was sent to us.
+	case HtlcAcceptedRevoke:
+		return AcceptedHtlcPenaltyWitnessSize, false, nil
+
+	// The revocation output of a second level output of an HTLC.
+	case HtlcSecondLevelRevoke:
+		return ToLocalPenaltyWitnessSize, false, nil
+	}
+
+	return 0, false, fmt.Errorf("unexpected witness type: %v", wt)
+}
+
+// AddWeightEstimation adds the estimated size of the witness in bytes to the
+// given weight estimator and returns the number of CSVs/CLTVs used by the
+// script.
+//
+// NOTE: This is part of the WitnessType interface.
+func (wt StandardWitnessType) AddWeightEstimation(
+	estimator *TxWeightEstimator) (int, int, error) {
+
+	var (
+		csvCount  = 0
+		cltvCount = 0
+	)
+
+	// For fee estimation purposes, we'll now attempt to obtain an
+	// upper bound on the weight this input will add when fully
+	// populated.
+	size, isNestedP2SH, err := wt.SizeUpperBound()
+	if err != nil {
+		return 0, 0, err
+	}
+
+	// If this is a nested P2SH input, then we'll need to factor in
+	// the additional data push within the sigScript.
+	if isNestedP2SH {
+		estimator.AddNestedP2WSHInput(size)
+	} else {
+		estimator.AddWitnessInput(size)
+	}
+
+	switch wt {
+	case CommitmentTimeLock,
+		HtlcOfferedTimeoutSecondLevel,
+		HtlcAcceptedSuccessSecondLevel:
+		csvCount++
+	case HtlcOfferedRemoteTimeout:
+		cltvCount++
+	}
+
+	return csvCount, cltvCount, nil
 }

sweep/txgenerator.go

@@ -56,7 +56,7 @@ func generateInputPartitionings(sweepableInputs []input.Input,
 	// on the signature length, which is not known yet at this point.
 	yields := make(map[wire.OutPoint]int64)
 	for _, input := range sweepableInputs {
-		size, _, err := getInputWitnessSizeUpperBound(input)
+		size, _, err := input.WitnessType().SizeUpperBound()
 		if err != nil {
 			return nil, fmt.Errorf(
 				"failed adding input weight: %v", err)
@@ -126,7 +126,7 @@ func getPositiveYieldInputs(sweepableInputs []input.Input, maxInputs int,
 	for idx, input := range sweepableInputs {
 		// Can ignore error, because it has already been checked when
 		// calculating the yields.
-		size, isNestedP2SH, _ := getInputWitnessSizeUpperBound(input)
+		size, isNestedP2SH, _ := input.WitnessType().SizeUpperBound()
 
 		// Keep a running weight estimate of the input set.
 		if isNestedP2SH {
@@ -251,59 +251,6 @@ func createSweepTx(inputs []input.Input, outputPkScript []byte,
 	return sweepTx, nil
 }
 
-// getInputWitnessSizeUpperBound returns the maximum length of the witness for
-// the given input if it would be included in a tx. We also return if the
-// output itself is a nested p2sh output, if so then we need to take into
-// account the extra sigScript data size.
-func getInputWitnessSizeUpperBound(inp input.Input) (int, bool, error) {
-	switch inp.WitnessType() {
-
-	// Outputs on a remote commitment transaction that pay directly to us.
-	case input.CommitSpendNoDelayTweakless:
-		fallthrough
-	case input.WitnessKeyHash:
-		fallthrough
-	case input.CommitmentNoDelay:
-		return input.P2WKHWitnessSize, false, nil
-
-	// Outputs on a past commitment transaction that pay directly
-	// to us.
-	case input.CommitmentTimeLock:
-		return input.ToLocalTimeoutWitnessSize, false, nil
-
-	// Outgoing second layer HTLC's that have confirmed within the
-	// chain, and the output they produced is now mature enough to
-	// sweep.
-	case input.HtlcOfferedTimeoutSecondLevel:
-		return input.ToLocalTimeoutWitnessSize, false, nil
-
-	// Incoming second layer HTLC's that have confirmed within the
-	// chain, and the output they produced is now mature enough to
-	// sweep.
-	case input.HtlcAcceptedSuccessSecondLevel:
-		return input.ToLocalTimeoutWitnessSize, false, nil
-
-	// An HTLC on the commitment transaction of the remote party,
-	// that has had its absolute timelock expire.
-	case input.HtlcOfferedRemoteTimeout:
-		return input.AcceptedHtlcTimeoutWitnessSize, false, nil
-
-	// An HTLC on the commitment transaction of the remote party,
-	// that can be swept with the preimage.
-	case input.HtlcAcceptedRemoteSuccess:
-		return input.OfferedHtlcSuccessWitnessSize, false, nil
-
-	// A nested P2SH input that has a p2wkh witness script. We'll mark this
-	// as nested P2SH so the caller can estimate the weight properly
-	// including the sigScript.
-	case input.NestedWitnessKeyHash:
-		return input.P2WKHWitnessSize, true, nil
-	}
-
-	return 0, false, fmt.Errorf("unexpected witness type: %v",
-		inp.WitnessType())
-}
-
 // getWeightEstimate returns a weight estimate for the given inputs.
 // Additionally, it returns counts for the number of csv and cltv inputs.
 func getWeightEstimate(inputs []input.Input) ([]input.Input, int64, int, int) {
@@ -328,10 +275,8 @@ func getWeightEstimate(inputs []input.Input) ([]input.Input, int64, int, int) {
 	for i := range inputs {
 		inp := inputs[i]
 
-		// For fee estimation purposes, we'll now attempt to obtain an
+		wt := inp.WitnessType()
-		// upper bound on the weight this input will add when fully
+		inpCsv, inpCltv, err := wt.AddWeightEstimation(&weightEstimate)
-		// populated.
-		size, isNestedP2SH, err := getInputWitnessSizeUpperBound(inp)
 		if err != nil {
 			log.Warn(err)
 
@@ -339,23 +284,8 @@ func getWeightEstimate(inputs []input.Input) ([]input.Input, int64, int, int) {
 			// given.
 			continue
 		}
-
+		csvCount += inpCsv
-		// If this is a nested P2SH input, then we'll need to factor in
+		cltvCount += inpCltv
-		// the additional data push within the sigScript.
-		if isNestedP2SH {
-			weightEstimate.AddNestedP2WSHInput(size)
-		} else {
-			weightEstimate.AddWitnessInput(size)
-		}
-
-		switch inp.WitnessType() {
-		case input.CommitmentTimeLock,
-			input.HtlcOfferedTimeoutSecondLevel,
-			input.HtlcAcceptedSuccessSecondLevel:
-			csvCount++
-		case input.HtlcOfferedRemoteTimeout:
-			cltvCount++
-		}
 		sweepInputs = append(sweepInputs, inp)
 	}
 

utxonursery.go

@@ -1344,7 +1344,7 @@ type kidOutput struct {
 }
 
 func makeKidOutput(outpoint, originChanPoint *wire.OutPoint,
-	blocksToMaturity uint32, witnessType input.WitnessType,
+	blocksToMaturity uint32, witnessType input.StandardWitnessType,
 	signDescriptor *input.SignDescriptor,
 	absoluteMaturity uint32) kidOutput {
 
@@ -1423,7 +1423,7 @@ func (k *kidOutput) Encode(w io.Writer) error {
 		return err
 	}
 
-	byteOrder.PutUint16(scratch[:2], uint16(k.WitnessType()))
+	byteOrder.PutUint16(scratch[:2], uint16(k.witnessType))
 	if _, err := w.Write(scratch[:2]); err != nil {
 		return err
 	}
@@ -1473,7 +1473,7 @@ func (k *kidOutput) Decode(r io.Reader) error {
 	if _, err := r.Read(scratch[:2]); err != nil {
 		return err
 	}
-	k.witnessType = input.WitnessType(byteOrder.Uint16(scratch[:2]))
+	k.witnessType = input.StandardWitnessType(byteOrder.Uint16(scratch[:2]))
 
 	return input.ReadSignDescriptor(r, &k.signDesc)
 }