[custom channels]: generate unique script keys for HTLCs

tapchannel/allocation_sort.go

@@ -2,7 +2,8 @@ package tapchannel
 
 import (
 	"bytes"
-	"sort"
+	"cmp"
+	"slices"
 )
 
 // InPlaceAllocationSort performs an in-place sort of output allocations.
@@ -14,51 +15,21 @@ import (
 // transactions, the script does not directly commit to them. Instead, the CLTVs
 // must be supplied separately to act as a tie-breaker, otherwise we may produce
 // invalid HTLC signatures if the receiver produces an alternative ordering
-// during verification.
+// during verification. Because multiple shards of the same MPP payment can be
+// identical in all other fields, we also use the HtlcIndex as a final
+// tie-breaker.
 //
-// NOTE: Commitment and commitment anchor outputs should have a 0 CLTV value.
+// NOTE: Commitment and commitment anchor outputs should have a 0 CLTV and
+// HtlcIndex value.
 func InPlaceAllocationSort(allocations []*Allocation) {
-	sort.Sort(sortableAllocationSlice{allocations})
-}
-
-// sortableAllocationSlice is a slice of allocations and the corresponding CLTV
-// values of any HTLCs. Commitment and commitment anchor outputs should have a
-// CLTV of 0.
-type sortableAllocationSlice struct {
-	allocations []*Allocation
-}
-
-// Len returns the length of the sortableAllocationSlice.
-//
-// NOTE: Part of the sort.Interface interface.
-func (s sortableAllocationSlice) Len() int {
-	return len(s.allocations)
-}
-
-// Swap exchanges the position of outputs i and j.
-//
-// NOTE: Part of the sort.Interface interface.
-func (s sortableAllocationSlice) Swap(i, j int) {
-	s.allocations[i], s.allocations[j] = s.allocations[j], s.allocations[i]
-}
-
-// Less is a modified BIP69 output comparison, that sorts based on value, then
-// pkScript, then CLTV value.
-//
-// NOTE: Part of the sort.Interface interface.
-func (s sortableAllocationSlice) Less(i, j int) bool {
-	allocI, allocJ := s.allocations[i], s.allocations[j]
-
-	if allocI.BtcAmount != allocJ.BtcAmount {
-		return allocI.BtcAmount < allocJ.BtcAmount
-	}
-
-	pkScriptCmp := bytes.Compare(
-		allocI.SortTaprootKeyBytes, allocJ.SortTaprootKeyBytes,
-	)
-	if pkScriptCmp != 0 {
-		return pkScriptCmp < 0
-	}
-
-	return allocI.CLTV < allocJ.CLTV
+	slices.SortFunc(allocations, func(i, j *Allocation) int {
+		return cmp.Or(
+			cmp.Compare(i.BtcAmount, j.BtcAmount),
+			bytes.Compare(
+				i.SortTaprootKeyBytes, j.SortTaprootKeyBytes,
+			),
+			cmp.Compare(i.CLTV, j.CLTV),
+			cmp.Compare(i.HtlcIndex, j.HtlcIndex),
+		)
+	})
 }

tapchannel/allocation_sort_test.go

@@ -38,6 +38,24 @@ func TestInPlaceAllocationSort(t *testing.T) {
 					SortTaprootKeyBytes: []byte("b"),
 					CLTV:                100,
 				},
+				{
+					BtcAmount:           1000,
+					SortTaprootKeyBytes: []byte("b"),
+					CLTV:                100,
+					HtlcIndex:           1,
+				},
+				{
+					BtcAmount:           1000,
+					SortTaprootKeyBytes: []byte("b"),
+					CLTV:                100,
+					HtlcIndex:           9,
+				},
+				{
+					BtcAmount:           1000,
+					SortTaprootKeyBytes: []byte("b"),
+					CLTV:                100,
+					HtlcIndex:           3,
+				},
 				{
 					BtcAmount:           1000,
 					SortTaprootKeyBytes: []byte("a"),
@@ -60,6 +78,24 @@ func TestInPlaceAllocationSort(t *testing.T) {
 					SortTaprootKeyBytes: []byte("b"),
 					CLTV:                100,
 				},
+				{
+					BtcAmount:           1000,
+					SortTaprootKeyBytes: []byte("b"),
+					CLTV:                100,
+					HtlcIndex:           1,
+				},
+				{
+					BtcAmount:           1000,
+					SortTaprootKeyBytes: []byte("b"),
+					CLTV:                100,
+					HtlcIndex:           3,
+				},
+				{
+					BtcAmount:           1000,
+					SortTaprootKeyBytes: []byte("b"),
+					CLTV:                100,
+					HtlcIndex:           9,
+				},
 				{
 					BtcAmount:           2000,
 					SortTaprootKeyBytes: []byte("b"),

tapchannel/aux_leaf_creator.go

@@ -129,6 +129,7 @@ func FetchLeavesFromCommit(chainParams *address.ChainParams,
 			leaf, err := CreateSecondLevelHtlcTx(
 				chanState, com.CommitTx, htlc.Amt.ToSatoshis(),
 				keys, chainParams, htlcOutputs, cltvTimeout,
+				htlc.HtlcIndex,
 			)
 			if err != nil {
 				return lfn.Err[returnType](fmt.Errorf("unable "+
@@ -169,6 +170,7 @@ func FetchLeavesFromCommit(chainParams *address.ChainParams,
 			leaf, err := CreateSecondLevelHtlcTx(
 				chanState, com.CommitTx, htlc.Amt.ToSatoshis(),
 				keys, chainParams, htlcOutputs, cltvTimeout,
+				htlc.HtlcIndex,
 			)
 			if err != nil {
 				return lfn.Err[returnType](fmt.Errorf("unable "+

tapchannel/aux_leaf_signer.go

@@ -3,12 +3,15 @@ package tapchannel
 import (
 	"bytes"
 	"fmt"
+	"math"
+	"math/big"
 	"sync"
 
 	"github.com/btcsuite/btcd/btcec/v2"
 	"github.com/btcsuite/btcd/btcutil/psbt"
 	"github.com/btcsuite/btcd/txscript"
 	"github.com/btcsuite/btcd/wire"
+	"github.com/decred/dcrd/dcrec/secp256k1/v4"
 	"github.com/lightninglabs/taproot-assets/address"
 	"github.com/lightninglabs/taproot-assets/asset"
 	"github.com/lightninglabs/taproot-assets/commitment"
@@ -367,7 +370,7 @@ func verifyHtlcSignature(chainParams *address.ChainParams,
 
 	vPackets, err := htlcSecondLevelPacketsFromCommit(
 		chainParams, chanState, commitTx, baseJob.KeyRing, htlcOutputs,
-		baseJob, htlcTimeout,
+		baseJob, htlcTimeout, baseJob.HTLC.HtlcIndex,
 	)
 	if err != nil {
 		return fmt.Errorf("error generating second level packets: %w",
@@ -511,7 +514,7 @@ func (s *AuxLeafSigner) generateHtlcSignature(chanState lnwallet.AuxChanState,
 
 	vPackets, err := htlcSecondLevelPacketsFromCommit(
 		s.cfg.ChainParams, chanState, commitTx, baseJob.KeyRing,
-		htlcOutputs, baseJob, htlcTimeout,
+		htlcOutputs, baseJob, htlcTimeout, baseJob.HTLC.HtlcIndex,
 	)
 	if err != nil {
 		return lnwallet.AuxSigJobResp{}, fmt.Errorf("error generating "+
@@ -599,12 +602,12 @@ func (s *AuxLeafSigner) generateHtlcSignature(chanState lnwallet.AuxChanState,
 func htlcSecondLevelPacketsFromCommit(chainParams *address.ChainParams,
 	chanState lnwallet.AuxChanState, commitTx *wire.MsgTx,
 	keyRing lnwallet.CommitmentKeyRing, htlcOutputs []*cmsg.AssetOutput,
-	baseJob lnwallet.BaseAuxJob,
-	htlcTimeout fn.Option[uint32]) ([]*tappsbt.VPacket, error) {
+	baseJob lnwallet.BaseAuxJob, htlcTimeout fn.Option[uint32],
+	htlcIndex uint64) ([]*tappsbt.VPacket, error) {
 
 	packets, _, err := CreateSecondLevelHtlcPackets(
 		chanState, commitTx, baseJob.HTLC.Amount.ToSatoshis(),
-		keyRing, chainParams, htlcOutputs, htlcTimeout,
+		keyRing, chainParams, htlcOutputs, htlcTimeout, htlcIndex,
 	)
 	if err != nil {
 		return nil, fmt.Errorf("error creating second level HTLC "+
@@ -751,3 +754,108 @@ func (v *schnorrSigValidator) validateSchnorrSig(virtualTx *wire.MsgTx,
 
 	return nil
 }
+
+// ScriptKeyTweakFromHtlcIndex converts the given HTLC index into a modulo N
+// scalar that can be used to tweak the internal key of the HTLC script key on
+// the asset level. The value of 1 is always added to the index to make sure
+// this value is always non-zero.
+func ScriptKeyTweakFromHtlcIndex(index input.HtlcIndex) *secp256k1.ModNScalar {
+	// If we're at math.MaxUint64, we'd wrap around to 0 if we incremented
+	// by 1, but we need to make sure the tweak is 1 to not cause a
+	// multiplication by zero. This should never happen, as it would mean we
+	// have more than math.MaxUint64 updates in a channel, which exceeds the
+	// protocol's maximum.
+	if index == math.MaxUint64 {
+		return new(secp256k1.ModNScalar).SetInt(1)
+	}
+
+	// We need to avoid the tweak being zero, so we always add 1 to the
+	// index. Otherwise, we'd multiply G by zero.
+	index++
+
+	indexAsBytes := new(big.Int).SetUint64(index).Bytes()
+	indexAsScalar := new(secp256k1.ModNScalar)
+	_ = indexAsScalar.SetByteSlice(indexAsBytes)
+
+	return indexAsScalar
+}
+
+// TweakPubKeyWithIndex tweaks the given internal public key with the given
+// HTLC index. The tweak is derived from the index in a way that never results
+// in a zero tweak. The value of 1 is always added to the index to make sure
+// this value is always non-zero. The public key is tweaked like this:
+//
+//	tweakedKey = key + (index+1) * G
+func TweakPubKeyWithIndex(pubKey *btcec.PublicKey,
+	index input.HtlcIndex) *btcec.PublicKey {
+
+	// Avoid panic if input is nil.
+	if pubKey == nil {
+		return nil
+	}
+
+	// We need to operate on Jacobian points, which is just a different
+	// representation of the public key that allows us to do scalar
+	// multiplication.
+	var (
+		pubKeyJacobian, tweakTimesG, tweakedKey btcec.JacobianPoint
+	)
+	pubKey.AsJacobian(&pubKeyJacobian)
+
+	// Derive the tweak from the HTLC index in a way that never results in
+	// a zero tweak. Then we multiply G by the tweak.
+	tweak := ScriptKeyTweakFromHtlcIndex(index)
+	secp256k1.ScalarBaseMultNonConst(tweak, &tweakTimesG)
+
+	// And finally we add the result to the key to get the tweaked key.
+	secp256k1.AddNonConst(&pubKeyJacobian, &tweakTimesG, &tweakedKey)
+
+	// Convert the tweaked key back to an affine point and create a new
+	// taproot key from it.
+	tweakedKey.ToAffine()
+	return btcec.NewPublicKey(&tweakedKey.X, &tweakedKey.Y)
+}
+
+// TweakHtlcTree tweaks the internal key of the given HTLC script tree with the
+// given index, then returns the tweaked tree with the updated taproot key.
+// The tapscript tree and tapscript root are not modified.
+// The internal key is tweaked like this:
+//
+//	tweakedInternalKey = internalKey + (index+1) * G
+func TweakHtlcTree(tree input.ScriptTree,
+	index input.HtlcIndex) input.ScriptTree {
+
+	// The tapscript tree and root are not modified, only the internal key
+	// is tweaked, which inherently modifies the taproot key.
+	tweakedInternalPubKey := TweakPubKeyWithIndex(tree.InternalKey, index)
+	newTaprootKey := txscript.ComputeTaprootOutputKey(
+		tweakedInternalPubKey, tree.TapscriptRoot,
+	)
+
+	return input.ScriptTree{
+		InternalKey:   tweakedInternalPubKey,
+		TaprootKey:    newTaprootKey,
+		TapscriptTree: tree.TapscriptTree,
+		TapscriptRoot: tree.TapscriptRoot,
+	}
+}
+
+// AddTweakWithIndex adds the given index to the given tweak. If the tweak is
+// empty, the index is used as the tweak directly. The value of 1 is always
+// added to the index to make sure this value is always non-zero.
+func AddTweakWithIndex(maybeTweak []byte, index input.HtlcIndex) []byte {
+	indexTweak := ScriptKeyTweakFromHtlcIndex(index)
+
+	// If we don't already have a tweak, we just use the index as the tweak.
+	if len(maybeTweak) == 0 {
+		return fn.ByteSlice(indexTweak.Bytes())
+	}
+
+	// If we have a tweak, we need to parse/decode it as a scalar, then add
+	// the index as a scalar, and encode it back to a byte slice.
+	tweak := new(secp256k1.ModNScalar)
+	_ = tweak.SetByteSlice(maybeTweak)
+	newTweak := tweak.Add(indexTweak)
+
+	return fn.ByteSlice(newTweak.Bytes())
+}