input+lnwallet: update taproot scripts to accept optional aux leaf

In this commit, we update all the taproot scripts to also accept an
optional aux leaf. This aux leaf can be used to add more redemption
paths for advanced channels, or just as an extra commitment space.

Author: Roasbeef

input/script_utils.go

@@ -673,6 +673,13 @@ type HtlcScriptTree struct {
 	// TimeoutTapLeaf is the tapleaf for the timeout path.
 	TimeoutTapLeaf txscript.TapLeaf
 
+	// AuxLeaf is an auxiliary leaf that can be used to extend the base
+	// HTLC script tree with new spend paths, or just as extra commitment
+	// space. When present, this leaf will always be in the left-most or
+	// right-most area of the tapscript tree.
+	AuxLeaf AuxTapLeaf
+
+	// htlcType is the type of HTLC script this is.
 	htlcType htlcType
 }
 
@@ -748,8 +755,8 @@ var _ TapscriptDescriptor = (*HtlcScriptTree)(nil)
 // senderHtlcTapScriptTree builds the tapscript tree which is used to anchor
 // the HTLC key for HTLCs on the sender's commitment.
 func senderHtlcTapScriptTree(senderHtlcKey, receiverHtlcKey,
-	revokeKey *btcec.PublicKey, payHash []byte,
-	hType htlcType) (*HtlcScriptTree, error) {
+	revokeKey *btcec.PublicKey, payHash []byte, hType htlcType,
+	auxLeaf AuxTapLeaf) (*HtlcScriptTree, error) {
 
 	// First, we'll obtain the tap leaves for both the success and timeout
 	// path.
@@ -766,11 +773,14 @@ func senderHtlcTapScriptTree(senderHtlcKey, receiverHtlcKey,
 		return nil, err
 	}
 
+	tapLeaves := []txscript.TapLeaf{successTapLeaf, timeoutTapLeaf}
+	auxLeaf.WhenSome(func(l txscript.TapLeaf) {
+		tapLeaves = append(tapLeaves, l)
+	})
+
 	// With the two leaves obtained, we'll now make the tapscript tree,
 	// then obtain the root from that
-	tapscriptTree := txscript.AssembleTaprootScriptTree(
-		successTapLeaf, timeoutTapLeaf,
-	)
+	tapscriptTree := txscript.AssembleTaprootScriptTree(tapLeaves...)
 
 	tapScriptRoot := tapscriptTree.RootNode.TapHash()
 
@@ -789,6 +799,7 @@ func senderHtlcTapScriptTree(senderHtlcKey, receiverHtlcKey,
 		},
 		SuccessTapLeaf: successTapLeaf,
 		TimeoutTapLeaf: timeoutTapLeaf,
+		AuxLeaf:        auxLeaf,
 		htlcType:       hType,
 	}, nil
 }
@@ -822,8 +833,8 @@ func senderHtlcTapScriptTree(senderHtlcKey, receiverHtlcKey,
 // The top level keyspend key is the revocation key, which allows a defender to
 // unilaterally spend the created output.
 func SenderHTLCScriptTaproot(senderHtlcKey, receiverHtlcKey,
-	revokeKey *btcec.PublicKey, payHash []byte,
-	localCommit bool) (*HtlcScriptTree, error) {
+	revokeKey *btcec.PublicKey, payHash []byte, localCommit bool,
+	auxLeaf AuxTapLeaf) (*HtlcScriptTree, error) {
 
 	var hType htlcType
 	if localCommit {
@@ -836,8 +847,8 @@ func SenderHTLCScriptTaproot(senderHtlcKey, receiverHtlcKey,
 	// tree that includes the top level output script, as well as the two
 	// tap leaf paths.
 	return senderHtlcTapScriptTree(
-		senderHtlcKey, receiverHtlcKey, revokeKey, payHash,
-		hType,
+		senderHtlcKey, receiverHtlcKey, revokeKey, payHash, hType,
+		auxLeaf,
 	)
 }
 
@@ -1307,8 +1318,8 @@ func ReceiverHtlcTapLeafSuccess(receiverHtlcKey *btcec.PublicKey,
 // receiverHtlcTapScriptTree builds the tapscript tree which is used to anchor
 // the HTLC key for HTLCs on the receiver's commitment.
 func receiverHtlcTapScriptTree(senderHtlcKey, receiverHtlcKey,
-	revokeKey *btcec.PublicKey, payHash []byte,
-	cltvExpiry uint32, hType htlcType) (*HtlcScriptTree, error) {
+	revokeKey *btcec.PublicKey, payHash []byte, cltvExpiry uint32,
+	hType htlcType, auxLeaf AuxTapLeaf) (*HtlcScriptTree, error) {
 
 	// First, we'll obtain the tap leaves for both the success and timeout
 	// path.
@@ -1325,11 +1336,14 @@ func receiverHtlcTapScriptTree(senderHtlcKey, receiverHtlcKey,
 		return nil, err
 	}
 
+	tapLeaves := []txscript.TapLeaf{timeoutTapLeaf, successTapLeaf}
+	auxLeaf.WhenSome(func(l txscript.TapLeaf) {
+		tapLeaves = append(tapLeaves, l)
+	})
+
 	// With the two leaves obtained, we'll now make the tapscript tree,
 	// then obtain the root from that
-	tapscriptTree := txscript.AssembleTaprootScriptTree(
-		timeoutTapLeaf, successTapLeaf,
-	)
+	tapscriptTree := txscript.AssembleTaprootScriptTree(tapLeaves...)
 
 	tapScriptRoot := tapscriptTree.RootNode.TapHash()
 
@@ -1348,6 +1362,7 @@ func receiverHtlcTapScriptTree(senderHtlcKey, receiverHtlcKey,
 		},
 		SuccessTapLeaf: successTapLeaf,
 		TimeoutTapLeaf: timeoutTapLeaf,
+		AuxLeaf:        auxLeaf,
 		htlcType:       hType,
 	}, nil
 }
@@ -1382,7 +1397,8 @@ func receiverHtlcTapScriptTree(senderHtlcKey, receiverHtlcKey,
 // the tap leaf are returned.
 func ReceiverHTLCScriptTaproot(cltvExpiry uint32,
 	senderHtlcKey, receiverHtlcKey, revocationKey *btcec.PublicKey,
-	payHash []byte, ourCommit bool) (*HtlcScriptTree, error) {
+	payHash []byte, ourCommit bool, auxLeaf AuxTapLeaf) (*HtlcScriptTree,
+	error) {
 
 	var hType htlcType
 	if ourCommit {
@@ -1396,7 +1412,7 @@ func ReceiverHTLCScriptTaproot(cltvExpiry uint32,
 	// tap leaf paths.
 	return receiverHtlcTapScriptTree(
 		senderHtlcKey, receiverHtlcKey, revocationKey, payHash,
-		cltvExpiry, hType,
+		cltvExpiry, hType, auxLeaf,
 	)
 }
 
@@ -1625,9 +1641,9 @@ func TaprootSecondLevelTapLeaf(delayKey *btcec.PublicKey,
 }
 
 // SecondLevelHtlcTapscriptTree construct the indexed tapscript tree needed to
-// generate the taptweak to create the final output and also control block.
-func SecondLevelHtlcTapscriptTree(delayKey *btcec.PublicKey,
-	csvDelay uint32) (*txscript.IndexedTapScriptTree, error) {
+// generate the tap tweak to create the final output and also control block.
+func SecondLevelHtlcTapscriptTree(delayKey *btcec.PublicKey, csvDelay uint32,
+	auxLeaf AuxTapLeaf) (*txscript.IndexedTapScriptTree, error) {
 
 	// First grab the second level leaf script we need to create the top
 	// level output.
@@ -1636,9 +1652,14 @@ func SecondLevelHtlcTapscriptTree(delayKey *btcec.PublicKey,
 		return nil, err
 	}
 
+	tapLeaves := []txscript.TapLeaf{secondLevelTapLeaf}
+	auxLeaf.WhenSome(func(l txscript.TapLeaf) {
+		tapLeaves = append(tapLeaves, l)
+	})
+
 	// Now that we have the sole second level script, we can create the
 	// tapscript tree that commits to both the leaves.
-	return txscript.AssembleTaprootScriptTree(secondLevelTapLeaf), nil
+	return txscript.AssembleTaprootScriptTree(tapLeaves...), nil
 }
 
 // TaprootSecondLevelHtlcScript is the uniform script that's used as the output
@@ -1658,12 +1679,12 @@ func SecondLevelHtlcTapscriptTree(delayKey *btcec.PublicKey,
 //
 // The keyspend path require knowledge of the top level revocation private key.
 func TaprootSecondLevelHtlcScript(revokeKey, delayKey *btcec.PublicKey,
-	csvDelay uint32) (*btcec.PublicKey, error) {
+	csvDelay uint32, auxLeaf AuxTapLeaf) (*btcec.PublicKey, error) {
 
 	// First, we'll make the tapscript tree that commits to the redemption
 	// path.
 	tapScriptTree, err := SecondLevelHtlcTapscriptTree(
-		delayKey, csvDelay,
+		delayKey, csvDelay, auxLeaf,
 	)
 	if err != nil {
 		return nil, err
@@ -1688,17 +1709,21 @@ type SecondLevelScriptTree struct {
 
 	// SuccessTapLeaf is the tapleaf for the redemption path.
 	SuccessTapLeaf txscript.TapLeaf
+
+	// AuxLeaf is an optional leaf that can be used to extend the script
+	// tree.
+	AuxLeaf AuxTapLeaf
 }
 
 // TaprootSecondLevelScriptTree constructs the tapscript tree used to spend the
 // second level HTLC output.
 func TaprootSecondLevelScriptTree(revokeKey, delayKey *btcec.PublicKey,
-	csvDelay uint32) (*SecondLevelScriptTree, error) {
+	csvDelay uint32, auxLeaf AuxTapLeaf) (*SecondLevelScriptTree, error) {
 
 	// First, we'll make the tapscript tree that commits to the redemption
 	// path.
 	tapScriptTree, err := SecondLevelHtlcTapscriptTree(
-		delayKey, csvDelay,
+		delayKey, csvDelay, auxLeaf,
 	)
 	if err != nil {
 		return nil, err
@@ -1719,6 +1744,7 @@ func TaprootSecondLevelScriptTree(revokeKey, delayKey *btcec.PublicKey,
 			InternalKey:   revokeKey,
 		},
 		SuccessTapLeaf: tapScriptTree.LeafMerkleProofs[0].TapLeaf,
+		AuxLeaf:        auxLeaf,
 	}, nil
 }
 
@@ -2095,6 +2121,12 @@ type CommitScriptTree struct {
 	// RevocationLeaf is the leaf used to spend the output with the
 	// revocation key signature.
 	RevocationLeaf txscript.TapLeaf
+
+	// AuxLeaf is an auxiliary leaf that can be used to extend the base
+	// commitment script tree with new spend paths, or just as extra
+	// commitment space. When present, this leaf will always be in the
+	// left-most or right-most area of the tapscript tree.
+	AuxLeaf AuxTapLeaf
 }
 
 // A compile time check to ensure CommitScriptTree implements the
@@ -2154,8 +2186,9 @@ func (c *CommitScriptTree) CtrlBlockForPath(path ScriptPath,
 
 // NewLocalCommitScriptTree returns a new CommitScript tree that can be used to
 // create and spend the commitment output for the local party.
-func NewLocalCommitScriptTree(csvTimeout uint32,
-	selfKey, revokeKey *btcec.PublicKey) (*CommitScriptTree, error) {
+func NewLocalCommitScriptTree(csvTimeout uint32, selfKey,
+	revokeKey *btcec.PublicKey, auxLeaf AuxTapLeaf) (*CommitScriptTree,
+	error) {
 
 	// First, we'll need to construct the tapLeaf that'll be our delay CSV
 	// clause.
@@ -2175,9 +2208,13 @@ func NewLocalCommitScriptTree(csvTimeout uint32,
 	// the two leaves, and then obtain a root from that.
 	delayTapLeaf := txscript.NewBaseTapLeaf(delayScript)
 	revokeTapLeaf := txscript.NewBaseTapLeaf(revokeScript)
-	tapScriptTree := txscript.AssembleTaprootScriptTree(
-		delayTapLeaf, revokeTapLeaf,
-	)
+
+	tapLeaves := []txscript.TapLeaf{delayTapLeaf, revokeTapLeaf}
+	auxLeaf.WhenSome(func(l txscript.TapLeaf) {
+		tapLeaves = append(tapLeaves, l)
+	})
+
+	tapScriptTree := txscript.AssembleTaprootScriptTree(tapLeaves...)
 	tapScriptRoot := tapScriptTree.RootNode.TapHash()
 
 	// Now that we have our root, we can arrive at the final output script
@@ -2195,6 +2232,7 @@ func NewLocalCommitScriptTree(csvTimeout uint32,
 		},
 		SettleLeaf:     delayTapLeaf,
 		RevocationLeaf: revokeTapLeaf,
+		AuxLeaf:        auxLeaf,
 	}, nil
 }
 
@@ -2264,7 +2302,7 @@ func TaprootCommitScriptToSelf(csvTimeout uint32,
 	selfKey, revokeKey *btcec.PublicKey) (*btcec.PublicKey, error) {
 
 	commitScriptTree, err := NewLocalCommitScriptTree(
-		csvTimeout, selfKey, revokeKey,
+		csvTimeout, selfKey, revokeKey, NoneTapLeaf(),
 	)
 	if err != nil {
 		return nil, err
@@ -2593,7 +2631,7 @@ func CommitScriptToRemoteConfirmed(key *btcec.PublicKey) ([]byte, error) {
 // NewRemoteCommitScriptTree constructs a new script tree for the remote party
 // to sweep their funds after a hard coded 1 block delay.
 func NewRemoteCommitScriptTree(remoteKey *btcec.PublicKey,
-) (*CommitScriptTree, error) {
+	auxLeaf AuxTapLeaf) (*CommitScriptTree, error) {
 
 	// First, construct the remote party's tapscript they'll use to sweep
 	// their outputs.
@@ -2609,10 +2647,16 @@ func NewRemoteCommitScriptTree(remoteKey *btcec.PublicKey,
 		return nil, err
 	}
 
+	tapLeaf := txscript.NewBaseTapLeaf(remoteScript)
+
+	tapLeaves := []txscript.TapLeaf{tapLeaf}
+	auxLeaf.WhenSome(func(l txscript.TapLeaf) {
+		tapLeaves = append(tapLeaves, l)
+	})
+
 	// With this script constructed, we'll map that into a tapLeaf, then
 	// make a new tapscript root from that.
-	tapLeaf := txscript.NewBaseTapLeaf(remoteScript)
-	tapScriptTree := txscript.AssembleTaprootScriptTree(tapLeaf)
+	tapScriptTree := txscript.AssembleTaprootScriptTree(tapLeaves...)
 	tapScriptRoot := tapScriptTree.RootNode.TapHash()
 
 	// Now that we have our root, we can arrive at the final output script
@@ -2629,6 +2673,7 @@ func NewRemoteCommitScriptTree(remoteKey *btcec.PublicKey,
 			InternalKey:   &TaprootNUMSKey,
 		},
 		SettleLeaf: tapLeaf,
+		AuxLeaf:    auxLeaf,
 	}, nil
 }
 
@@ -2645,9 +2690,9 @@ func NewRemoteCommitScriptTree(remoteKey *btcec.PublicKey,
 //	<remotepubkey> OP_CHECKSIG
 //	1 OP_CHECKSEQUENCEVERIFY OP_DROP
 func TaprootCommitScriptToRemote(remoteKey *btcec.PublicKey,
-) (*btcec.PublicKey, error) {
+	auxLeaf AuxTapLeaf) (*btcec.PublicKey, error) {
 
-	commitScriptTree, err := NewRemoteCommitScriptTree(remoteKey)
+	commitScriptTree, err := NewRemoteCommitScriptTree(remoteKey, auxLeaf)
 	if err != nil {
 		return nil, err
 	}