tapchannel: tweak HTLC script key internal key to enforce uniqueness

This commit tweaks the internal key of the asset-level script key with
the HTLC index to enforce uniqueness of script keys across multiple
HTLCs with the same payment hash and timeout (MPP shards of the same
payment).

Author: guggero

tapchannel/aux_leaf_signer.go

@@ -3,12 +3,14 @@ package tapchannel
 import (
 	"bytes"
 	"fmt"
+	"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"
@@ -751,3 +753,82 @@ func (v *schnorrSigValidator) validateSchnorrSig(virtualTx *wire.MsgTx,
 
 	return nil
 }
+
+// HtlcIndexAsScriptKeyTweak 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 HtlcIndexAsScriptKeyTweak(index input.HtlcIndex) *secp256k1.ModNScalar {
+	// We need to avoid the tweak being zero, so we always add 1 to the
+	// index. Otherwise, we'd multiply G by zero.
+	index++
+
+	// If we wrapped around from math.MaxUint64 to 0, 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 == 0 {
+		return new(secp256k1.ModNScalar).SetInt(1)
+	}
+
+	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 {
+
+	// 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 := HtlcIndexAsScriptKeyTweak(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.
+	tweakedInternalPubKey := TweakPubKeyWithIndex(tree.InternalKey, index)
+	newTaprootKey := txscript.ComputeTaprootOutputKey(
+		tweakedInternalPubKey, tree.TapscriptRoot,
+	)
+
+	return input.ScriptTree{
+		InternalKey:   tweakedInternalPubKey,
+		TaprootKey:    newTaprootKey,
+		TapscriptTree: tree.TapscriptTree,
+		TapscriptRoot: tree.TapscriptRoot,
+	}
+}

tapchannel/aux_leaf_signer_test.go

@@ -3,13 +3,17 @@ package tapchannel
 import (
 	"bytes"
 	"crypto/rand"
+	"encoding/binary"
 	"encoding/hex"
 	"fmt"
+	"math"
 	"testing"
 
 	"github.com/btcsuite/btcd/btcec/v2"
 	"github.com/btcsuite/btcd/txscript"
+	"github.com/decred/dcrd/dcrec/secp256k1/v4"
 	"github.com/lightninglabs/taproot-assets/asset"
+	"github.com/lightninglabs/taproot-assets/internal/test"
 	cmsg "github.com/lightninglabs/taproot-assets/tapchannelmsg"
 	"github.com/lightningnetwork/lnd/input"
 	"github.com/lightningnetwork/lnd/lnwire"
@@ -163,3 +167,230 @@ func makeCommitSig(t *testing.T, numAssetIDs, numHTLCs int) *lnwire.CommitSig {
 
 	return msg
 }
+
+// TestHtlcIndexAsScriptKeyTweak tests the HtlcIndexAsScriptKeyTweak function.
+func TestHtlcIndexAsScriptKeyTweak(t *testing.T) {
+	var (
+		buf               = make([]byte, 8)
+		maxUint64MinusOne = new(secp256k1.ModNScalar)
+		maxUint64         = new(secp256k1.ModNScalar)
+	)
+	binary.BigEndian.PutUint64(buf, math.MaxUint64-1)
+	_ = maxUint64MinusOne.SetByteSlice(buf)
+
+	binary.BigEndian.PutUint64(buf, math.MaxUint64)
+	_ = maxUint64.SetByteSlice(buf)
+
+	testCases := []struct {
+		name   string
+		index  uint64
+		result *secp256k1.ModNScalar
+	}{
+		{
+			name:   "index 0",
+			index:  0,
+			result: new(secp256k1.ModNScalar).SetInt(1),
+		},
+		{
+			name:  "index math.MaxUint32-1",
+			index: math.MaxUint32 - 1,
+			result: new(secp256k1.ModNScalar).SetInt(
+				math.MaxUint32,
+			),
+		},
+		{
+			name:   "index math.MaxUint64-2",
+			index:  math.MaxUint64 - 2,
+			result: maxUint64MinusOne,
+		},
+		{
+			name:   "index math.MaxUint64-1",
+			index:  math.MaxUint64 - 1,
+			result: maxUint64,
+		},
+		{
+			name:   "index math.MaxUint64, wraps around to 1",
+			index:  math.MaxUint64,
+			result: new(secp256k1.ModNScalar).SetInt(1),
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			tweak := HtlcIndexAsScriptKeyTweak(tc.index)
+			require.Equal(t, tc.result, tweak)
+		})
+	}
+}
+
+// TestTweakPubKeyWithIndex tests the TweakPubKeyWithIndex function.
+func TestTweakPubKeyWithIndex(t *testing.T) {
+	randNum := test.RandInt[uint32]()
+
+	makePubKey := func(tweak uint64) *btcec.PublicKey {
+		var (
+			buf    = make([]byte, 8)
+			scalar = new(secp256k1.ModNScalar)
+		)
+		binary.BigEndian.PutUint64(buf, uint64(randNum)+tweak)
+		_ = scalar.SetByteSlice(buf)
+		return secp256k1.NewPrivateKey(scalar).PubKey()
+	}
+	startKey := makePubKey(0)
+
+	testCases := []struct {
+		name   string
+		pubKey *btcec.PublicKey
+		index  uint64
+		result *btcec.PublicKey
+	}{
+		{
+			name:   "index 0",
+			pubKey: startKey,
+			index:  0,
+			result: makePubKey(1),
+		},
+		{
+			name:   "index 1",
+			pubKey: startKey,
+			index:  1,
+			result: makePubKey(2),
+		},
+		{
+			name:   "index 99",
+			pubKey: startKey,
+			index:  99,
+			result: makePubKey(100),
+		},
+		{
+			name:   "index math.MaxUint32-1",
+			pubKey: startKey,
+			index:  math.MaxUint32 - 1,
+			result: makePubKey(math.MaxUint32),
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			tweakedKey := TweakPubKeyWithIndex(tc.pubKey, tc.index)
+			require.Equal(
+				t, tc.result.SerializeCompressed(),
+				tweakedKey.SerializeCompressed(),
+			)
+		})
+	}
+}
+
+// TestTweakHtlcTree tests the TweakHtlcTree function.
+func TestTweakHtlcTree(t *testing.T) {
+	randTree := txscript.AssembleTaprootScriptTree(
+		test.RandTapLeaf(nil), test.RandTapLeaf(nil),
+		test.RandTapLeaf(nil),
+	)
+	randRoot := randTree.RootNode.TapHash()
+	randNum := test.RandInt[uint32]()
+
+	makePubKey := func(tweak uint64) *btcec.PublicKey {
+		var (
+			buf    = make([]byte, 8)
+			scalar = new(secp256k1.ModNScalar)
+		)
+		binary.BigEndian.PutUint64(buf, uint64(randNum)+tweak)
+		_ = scalar.SetByteSlice(buf)
+		return secp256k1.NewPrivateKey(scalar).PubKey()
+	}
+	makeTaprootKey := func(tweak uint64) *btcec.PublicKey {
+		return txscript.ComputeTaprootOutputKey(
+			makePubKey(tweak), randRoot[:],
+		)
+	}
+	startKey := makePubKey(0)
+	startTaprootKey := makeTaprootKey(0)
+
+	testCases := []struct {
+		name   string
+		tree   input.ScriptTree
+		index  uint64
+		result input.ScriptTree
+	}{
+		{
+			name: "index 0",
+			tree: input.ScriptTree{
+				InternalKey:   startKey,
+				TaprootKey:    startTaprootKey,
+				TapscriptTree: randTree,
+				TapscriptRoot: randRoot[:],
+			},
+			index: 0,
+			result: input.ScriptTree{
+				InternalKey:   makePubKey(1),
+				TaprootKey:    makeTaprootKey(1),
+				TapscriptTree: randTree,
+				TapscriptRoot: randRoot[:],
+			},
+		},
+		{
+			name: "index 1",
+			tree: input.ScriptTree{
+				InternalKey:   startKey,
+				TaprootKey:    startTaprootKey,
+				TapscriptTree: randTree,
+				TapscriptRoot: randRoot[:],
+			},
+			index: 1,
+			result: input.ScriptTree{
+				InternalKey:   makePubKey(2),
+				TaprootKey:    makeTaprootKey(2),
+				TapscriptTree: randTree,
+				TapscriptRoot: randRoot[:],
+			},
+		},
+		{
+			name: "index 99",
+			tree: input.ScriptTree{
+				InternalKey:   startKey,
+				TaprootKey:    startTaprootKey,
+				TapscriptTree: randTree,
+				TapscriptRoot: randRoot[:],
+			},
+			index: 99,
+			result: input.ScriptTree{
+				InternalKey:   makePubKey(100),
+				TaprootKey:    makeTaprootKey(100),
+				TapscriptTree: randTree,
+				TapscriptRoot: randRoot[:],
+			},
+		},
+		{
+			name: "index math.MaxUint32-1",
+			tree: input.ScriptTree{
+				InternalKey:   startKey,
+				TaprootKey:    startTaprootKey,
+				TapscriptTree: randTree,
+				TapscriptRoot: randRoot[:],
+			},
+			index: math.MaxUint32 - 1,
+			result: input.ScriptTree{
+				InternalKey:   makePubKey(math.MaxUint32),
+				TaprootKey:    makeTaprootKey(math.MaxUint32),
+				TapscriptTree: randTree,
+				TapscriptRoot: randRoot[:],
+			},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			tweakedTree := TweakHtlcTree(tc.tree, tc.index)
+			require.Equal(
+				t, tc.result.InternalKey.SerializeCompressed(),
+				tweakedTree.InternalKey.SerializeCompressed(),
+			)
+			require.Equal(
+				t, tc.result.TaprootKey.SerializeCompressed(),
+				tweakedTree.TaprootKey.SerializeCompressed(),
+			)
+			require.Equal(t, tc.result, tweakedTree)
+		})
+	}
+}

tapchannel/commitment.go

@@ -756,6 +756,13 @@ func CreateAllocations(chanState lnwallet.AuxChanState, ourBalance,
 				"allocation, HTLC is dust")
 		}
 
+		// To ensure uniqueness of the script key across HTLCs with the
+		// same payment hash and timeout (which would be equal
+		// otherwise), we tweak the asset level internal key of the
+		// script key with the HTLC index. We'll ONLY use this for the
+		// asset level, NOT for the BTC level.
+		tweakedTree := TweakHtlcTree(htlcTree, htlc.HtlcIndex)
+
 		allocations = append(allocations, &Allocation{
 			Type:           allocType,
 			Amount:         rfqmsg.Sum(htlc.AssetBalances),
@@ -766,13 +773,13 @@ func CreateAllocations(chanState lnwallet.AuxChanState, ourBalance,
 			NonAssetLeaves: sibling,
 			ScriptKey: asset.ScriptKey{
 				PubKey: asset.NewScriptKey(
-					htlcTree.TaprootKey,
+					tweakedTree.TaprootKey,
 				).PubKey,
 				TweakedScriptKey: &asset.TweakedScriptKey{
 					RawKey: keychain.KeyDescriptor{
-						PubKey: htlcTree.InternalKey,
+						PubKey: tweakedTree.InternalKey,
 					},
-					Tweak: htlcTree.TapscriptRoot,
+					Tweak: tweakedTree.TapscriptRoot,
 				},
 			},
 			SortTaprootKeyBytes: schnorr.SerializePubKey(