From 75de4e829aae7fd2ad900a2391e554ff397b8434 Mon Sep 17 00:00:00 2001 From: Dario Anongba Varela Date: Fri, 21 Aug 2026 14:37:52 +0200 Subject: [PATCH] tree: add asset-aware materialization context Asset-bearing tree nodes use distinct taproot tweaks and carry units independently from carrier satoshis. Store this data in an outpoint-keyed context so extracted paths sign deterministically without changing Bitcoin-only trees. Keep per-node tweak selection inside Tree signing and retain the existing public signer API. Copy MuSig2 participant slices before aggregation because the local signer may sort caller-owned state. --- lib/tree/asset_tree_context.go | 162 ++++++++++++++++++++++ lib/tree/asset_tree_context_test.go | 172 +++++++++++++++++++++++ lib/tree/leaf.go | 3 + lib/tree/node.go | 45 ++++-- lib/tree/node_test.go | 41 +++++- lib/tree/signing.go | 22 ++- lib/tree/signing_tweak_test.go | 203 ++++++++++++++++++++++++++++ lib/tree/tree.go | 16 ++- lib/tree/tree_structure.go | 39 +++++- 9 files changed, 680 insertions(+), 23 deletions(-) create mode 100644 lib/tree/asset_tree_context.go create mode 100644 lib/tree/asset_tree_context_test.go create mode 100644 lib/tree/signing_tweak_test.go diff --git a/lib/tree/asset_tree_context.go b/lib/tree/asset_tree_context.go new file mode 100644 index 000000000..d8fa951fb --- /dev/null +++ b/lib/tree/asset_tree_context.go @@ -0,0 +1,162 @@ +package tree + +import ( + "fmt" + "math" + + "github.com/btcsuite/btcd/wire/v2" +) + +// AssetTreeContext contains the data needed to materialize and sign an asset +// tree. Builders populate the context before sharing the tree. +type AssetTreeContext struct { + // Nodes do not have inputs until materialization, so the structure pass + // records amounts by node first. + amountsByNode map[*Node]uint64 + + // Input outpoints remain stable when paths are extracted or serialized. + amountsByInput map[wire.OutPoint]uint64 + + leafRootsByInput map[wire.OutPoint][]byte + tweaksByInput map[wire.OutPoint][]byte + packagesByInput map[wire.OutPoint][]byte + assetRef string +} + +// NewAssetTreeContext returns an empty asset tree context. +func NewAssetTreeContext() *AssetTreeContext { + return &AssetTreeContext{ + amountsByNode: make(map[*Node]uint64), + amountsByInput: make(map[wire.OutPoint]uint64), + leafRootsByInput: make(map[wire.OutPoint][]byte), + tweaksByInput: make(map[wire.OutPoint][]byte), + packagesByInput: make(map[wire.OutPoint][]byte), + } +} + +// SetNodeAssetAmount records a node's subtree asset amount. +func (c *AssetTreeContext) SetNodeAssetAmount(node *Node, amount uint64) { + c.amountsByNode[node] = amount + if node.Input != (wire.OutPoint{}) { + // Extracted paths contain cloned nodes, but preserve their + // inputs. + c.amountsByInput[node.Input] = amount + } +} + +// NodeAssetAmount returns a node's subtree asset amount. +func (c *AssetTreeContext) NodeAssetAmount(node *Node) uint64 { + if amount, ok := c.amountsByNode[node]; ok { + return amount + } + + // Path extraction clones nodes but preserves their inputs. + return c.amountsByInput[node.Input] +} + +// SetSigningTweak records the taproot tweak for a tree transaction. +func (c *AssetTreeContext) SetSigningTweak(input wire.OutPoint, tweak []byte) { + c.tweaksByInput[input] = append([]byte(nil), tweak...) +} + +// SigningTweak returns the taproot tweak for a tree transaction. +func (c *AssetTreeContext) SigningTweak(input wire.OutPoint) []byte { + return append([]byte(nil), c.tweaksByInput[input]...) +} + +// SetSealedPackage records a tree transaction's sealed transfer package. +func (c *AssetTreeContext) SetSealedPackage(input wire.OutPoint, pkg []byte) { + c.packagesByInput[input] = append([]byte(nil), pkg...) +} + +// SealedPackage returns a tree transaction's sealed transfer package. +func (c *AssetTreeContext) SealedPackage(input wire.OutPoint) []byte { + return append([]byte(nil), c.packagesByInput[input]...) +} + +// SetAssetRef records the tree's asset reference. +func (c *AssetTreeContext) SetAssetRef(ref string) { + c.assetRef = ref +} + +// AssetRef returns the tree's asset reference. +func (c *AssetTreeContext) AssetRef() string { + return c.assetRef +} + +// IsEmpty reports whether the context carries no asset data at all. +func (c *AssetTreeContext) IsEmpty() bool { + if c == nil { + return true + } + + return len(c.amountsByNode) == 0 && len(c.amountsByInput) == 0 && + len(c.leafRootsByInput) == 0 && len(c.tweaksByInput) == 0 && + len(c.packagesByInput) == 0 && c.assetRef == "" +} + +// tweakLookup returns the signing tweak for a node's input. +func (c *AssetTreeContext) tweakLookup() func(*Node) []byte { + return func(node *Node) []byte { + return c.SigningTweak(node.Input) + } +} + +// aggregateAssetAmounts records the asset amount of every subtree. +func aggregateAssetAmounts(node *Node, leafAssets map[*Node]uint64, + ctx *AssetTreeContext) (uint64, error) { + + if len(node.Children) == 0 { + amount := leafAssets[node] + ctx.SetNodeAssetAmount(node, amount) + + return amount, nil + } + + var total uint64 + for _, child := range node.Children { + childAmount, err := aggregateAssetAmounts( + child, leafAssets, ctx, + ) + if err != nil { + return 0, err + } + + if childAmount > math.MaxUint64-total { + return 0, fmt.Errorf("asset amount overflow " + + "aggregating subtree totals") + } + total += childAmount + } + + ctx.SetNodeAssetAmount(node, total) + + return total, nil +} + +// SetLeafAssetRoot records an asset leaf's commitment root. +func (c *AssetTreeContext) SetLeafAssetRoot(input wire.OutPoint, root []byte) { + if c == nil || len(root) == 0 { + return + } + + c.leafRootsByInput[input] = append([]byte(nil), root...) +} + +// LeafAssetRoot returns an asset leaf's commitment root. +func (c *AssetTreeContext) LeafAssetRoot(input wire.OutPoint) []byte { + if c == nil { + return nil + } + + return append([]byte(nil), c.leafRootsByInput[input]...) +} + +// LeafAssetRoot returns an asset leaf's commitment root. +func (t *Tree) LeafAssetRoot(input wire.OutPoint) []byte { + if t == nil { + return nil + } + + return t.AssetContext.LeafAssetRoot(input) +} diff --git a/lib/tree/asset_tree_context_test.go b/lib/tree/asset_tree_context_test.go new file mode 100644 index 000000000..da18b30a1 --- /dev/null +++ b/lib/tree/asset_tree_context_test.go @@ -0,0 +1,172 @@ +package tree + +import ( + "bytes" + "math" + "testing" + + "github.com/btcsuite/btcd/btcec/v2" + "github.com/btcsuite/btcd/wire/v2" + "github.com/stretchr/testify/require" +) + +// assetTestLeaves returns leaves with the given asset amounts. +func assetTestLeaves(t *testing.T, amounts ...uint64) []LeafDescriptor { + t.Helper() + + leaves := make([]LeafDescriptor, len(amounts)) + for i, amount := range amounts { + _, owner := createTestKey(t) + leaves[i] = LeafDescriptor{ + PkScript: []byte{ + 0x51, + byte(i), + }, + Amount: 10_000, + CoSignerKey: owner, + AssetAmount: amount, + } + } + + return leaves +} + +// TestBuildStructureAssetAggregation tests subtree amount calculation. +func TestBuildStructureAssetAggregation(t *testing.T) { + t.Parallel() + + _, operatorKey := createTestKey(t) + cfg := StructureConfig{OperatorKey: operatorKey, Radix: 2} + + t.Run("btc-only trees have no asset context", func(t *testing.T) { + t.Parallel() + + structure, err := BuildStructure( + assetTestLeaves(t, 0, 0, 0), cfg, + ) + require.NoError(t, err) + require.Nil(t, structure.AssetContext) + }) + + t.Run("subtree totals cover every node", func(t *testing.T) { + t.Parallel() + + structure, err := BuildStructure( + assetTestLeaves(t, 800, 200, 500), cfg, + ) + require.NoError(t, err) + require.NotNil(t, structure.AssetContext) + + assetCtx := structure.AssetContext + require.EqualValues( + t, 1_500, assetCtx.NodeAssetAmount(structure.Root), + ) + + // Every node's total must equal the sum of its children, + // and each leaf must carry one of the input amounts. + var leafTotal uint64 + for node := range structure.Root.NodesIter() { + if len(node.Children) == 0 { + leafTotal += assetCtx.NodeAssetAmount(node) + continue + } + + var childSum uint64 + for _, child := range node.Children { + childSum += assetCtx.NodeAssetAmount(child) + } + require.Equal( + t, childSum, assetCtx.NodeAssetAmount(node), + ) + } + require.EqualValues(t, 1_500, leafTotal) + }) + + t.Run("mixed asset and btc leaves", func(t *testing.T) { + t.Parallel() + + structure, err := BuildStructure( + assetTestLeaves(t, 300, 0), cfg, + ) + require.NoError(t, err) + require.NotNil(t, structure.AssetContext) + require.EqualValues( + t, 300, + structure.AssetContext.NodeAssetAmount(structure.Root), + ) + }) + + t.Run("overflow rejected", func(t *testing.T) { + t.Parallel() + + _, err := BuildStructure( + assetTestLeaves(t, math.MaxUint64, 1), cfg, + ) + require.ErrorContains(t, err, "overflow") + }) +} + +// TestAssetTreeContextTweakLookup tests outpoint lookup and byte ownership. +func TestAssetTreeContextTweakLookup(t *testing.T) { + t.Parallel() + + assetCtx := NewAssetTreeContext() + input := wire.OutPoint{Index: 3} + input.Hash[0] = 0xAB + + tweak := bytes.Repeat([]byte{0x07}, 32) + assetCtx.SetSigningTweak(input, tweak) + tweak[0] = 0xFF + + lookup := assetCtx.tweakLookup() + got := lookup(&Node{Input: input}) + require.Equal(t, bytes.Repeat([]byte{0x07}, 32), got) + got[0] = 0xEE + require.Equal(t, byte(0x07), lookup(&Node{Input: input})[0]) + require.Nil(t, lookup(&Node{Input: wire.OutPoint{Index: 9}})) + + pkg := []byte{0x01, 0x02} + assetCtx.SetSealedPackage(input, pkg) + pkg[0] = 0xEE + gotPkg := assetCtx.SealedPackage(input) + require.Equal(t, []byte{0x01, 0x02}, gotPkg) + gotPkg[0] = 0xDD + require.Equal(t, byte(0x01), assetCtx.SealedPackage(input)[0]) + + root := []byte{0x03, 0x04} + assetCtx.SetLeafAssetRoot(input, root) + root[0] = 0xCC + gotRoot := assetCtx.LeafAssetRoot(input) + require.Equal(t, []byte{0x03, 0x04}, gotRoot) + gotRoot[0] = 0xBB + require.Equal(t, byte(0x03), assetCtx.LeafAssetRoot(input)[0]) + require.False(t, assetCtx.IsEmpty()) +} + +// TestAssetTreeContextEmpty tests empty context detection. +func TestAssetTreeContextEmpty(t *testing.T) { + t.Parallel() + + var nilContext *AssetTreeContext + require.True(t, nilContext.IsEmpty()) + + assetCtx := NewAssetTreeContext() + require.True(t, assetCtx.IsEmpty()) + + assetCtx.SetLeafAssetRoot(wire.OutPoint{Index: 1}, []byte{0x01}) + require.False(t, assetCtx.IsEmpty()) +} + +// TestAssetTreeContextAmountsByIdentity tests node identity keying. +func TestAssetTreeContextAmountsByIdentity(t *testing.T) { + t.Parallel() + + assetCtx := NewAssetTreeContext() + _, key := createTestKey(t) + nodeA := &Node{CoSigners: []*btcec.PublicKey{key}} + nodeB := &Node{CoSigners: []*btcec.PublicKey{key}} + + assetCtx.SetNodeAssetAmount(nodeA, 42) + require.EqualValues(t, 42, assetCtx.NodeAssetAmount(nodeA)) + require.Zero(t, assetCtx.NodeAssetAmount(nodeB)) +} diff --git a/lib/tree/leaf.go b/lib/tree/leaf.go index 9ae16ac8c..d1e50fefb 100644 --- a/lib/tree/leaf.go +++ b/lib/tree/leaf.go @@ -19,4 +19,7 @@ type LeafDescriptor struct { // CoSignerKey is the public key of the leaf owner who must participate // in signing this leaf's transaction along with the operator. CoSignerKey *btcec.PublicKey + + // AssetAmount is the number of asset units in the leaf output. + AssetAmount uint64 } diff --git a/lib/tree/node.go b/lib/tree/node.go index 918ee35bf..fde3c6642 100644 --- a/lib/tree/node.go +++ b/lib/tree/node.go @@ -216,10 +216,35 @@ func NewBranchNode(input wire.OutPoint, groups [][]LeafDescriptor, }, nil } -// ComputeFinalKey computes the final aggregated public key for signing. -// This is a helper function that aggregates the cosigners and applies the -// taproot tweak with the given sweep tapscript root. It handles both -// single-key and multi-key cases. +// ComputeInternalKey returns the untweaked MuSig2 aggregate key. +func ComputeInternalKey(cosigners []*btcec.PublicKey) (*btcec.PublicKey, + error) { + + if len(cosigners) == 0 { + return nil, fmt.Errorf("no cosigners provided") + } + + if len(cosigners) == 1 { + return cosigners[0], nil + } + + // AggregateKeys sorts its input in place. + aggKey, _, _, err := musig2.AggregateKeys( + copyCosigners(cosigners), true, + ) + if err != nil { + return nil, fmt.Errorf("failed to aggregate keys: %w", err) + } + + return aggKey.PreTweakedKey, nil +} + +// copyCosigners returns a shallow copy of the cosigner set. +func copyCosigners(cosigners []*btcec.PublicKey) []*btcec.PublicKey { + return append([]*btcec.PublicKey(nil), cosigners...) +} + +// ComputeFinalKey returns the tweaked aggregate key for the cosigners. func ComputeFinalKey(cosigners []*btcec.PublicKey, sweepTapscriptRoot []byte) (*btcec.PublicKey, error) { @@ -236,9 +261,10 @@ func ComputeFinalKey(cosigners []*btcec.PublicKey, ), nil } - // Multi-key case: use MuSig2 aggregation with taproot tweak. + // AggregateKeys sorts its input in place. aggKey, _, _, err := musig2.AggregateKeys( - cosigners, true, musig2.WithTaprootKeyTweak(sweepTapscriptRoot), + copyCosigners(cosigners), true, + musig2.WithTaprootKeyTweak(sweepTapscriptRoot), ) if err != nil { return nil, fmt.Errorf("failed to aggregate keys: %w", err) @@ -828,14 +854,17 @@ func (n *Node) SigHash(prevOutFetcher txscript.PrevOutputFetcher) ([]byte, ) } -// NewSignerSession creates a new MuSig2 signing session for this node. +// NewSignerSession creates a MuSig2 signing session for this node. func (n *Node) NewSignerSession(signerKey *keychain.KeyDescriptor, signer input.MuSig2Signer, sweepTapscriptRoot []byte) ( *input.MuSig2SessionInfo, error) { + // A local signer may sort the cosigner slice in place. + cosigners := copyCosigners(n.CoSigners) + return signer.MuSig2CreateSession( input.MuSig2Version100RC2, signerKey.KeyLocator, - n.CoSigners, &input.MuSig2Tweaks{ + cosigners, &input.MuSig2Tweaks{ TaprootTweak: sweepTapscriptRoot, }, nil, nil, ) diff --git a/lib/tree/node_test.go b/lib/tree/node_test.go index 81180b928..2227277ac 100644 --- a/lib/tree/node_test.go +++ b/lib/tree/node_test.go @@ -1,7 +1,9 @@ package tree import ( + "bytes" "fmt" + "sort" "testing" "github.com/btcsuite/btcd/btcec/v2" @@ -1070,12 +1072,10 @@ func TestNodeConstructors(t *testing.T) { require.Equal(t, leaf.PkScript, node.Outputs[0].PkScript) require.Equal(t, int64(0), node.Outputs[1].Value) - // Verify cosigners (owner and operator). - // Note: musig2.AggregateKeys sorts keys, so we can't - // assume order. - require.Len(t, node.CoSigners, 2) - require.Contains(t, node.CoSigners, ownerKey) - require.Contains(t, node.CoSigners, operatorKey) + require.Equal( + t, []*btcec.PublicKey{ownerKey, operatorKey}, + node.CoSigners, + ) // Verify it's a leaf. require.True(t, node.IsLeaf()) @@ -1496,3 +1496,32 @@ func TestPrettyPrint(t *testing.T) { require.Contains(t, output, "K1") }) } + +// TestComputeKeysDoNotReorderCosigners tests cosigner slice ownership. +func TestComputeKeysDoNotReorderCosigners(t *testing.T) { + t.Parallel() + + // Three keys whose serialized order is deliberately not sorted. + var cosigners []*btcec.PublicKey + for len(cosigners) < 3 { + key, err := btcec.NewPrivateKey() + require.NoError(t, err) + cosigners = append(cosigners, key.PubKey()) + } + sort.Slice(cosigners, func(i, j int) bool { + return bytes.Compare( + cosigners[i].SerializeCompressed(), + cosigners[j].SerializeCompressed(), + ) > 0 + }) + + original := append([]*btcec.PublicKey(nil), cosigners...) + + _, err := ComputeInternalKey(cosigners) + require.NoError(t, err) + require.Equal(t, original, cosigners, "internal key reordered the set") + + _, err = ComputeFinalKey(cosigners, make([]byte, 32)) + require.NoError(t, err) + require.Equal(t, original, cosigners, "final key reordered the set") +} diff --git a/lib/tree/signing.go b/lib/tree/signing.go index 898b6b276..f474061d5 100644 --- a/lib/tree/signing.go +++ b/lib/tree/signing.go @@ -177,6 +177,17 @@ func NewSignerSession(signer input.MuSig2Signer, prevOuts txscript.PrevOutputFetcher, tree *Node) (*SignerSession, error) { + return newSignerSession( + signer, signerKey, sweepTapscriptRoot, prevOuts, tree, nil, + ) +} + +// newSignerSession creates a signing session with optional per-node tweaks. +func newSignerSession(signer input.MuSig2Signer, + signerKey *keychain.KeyDescriptor, sweepTapscriptRoot []byte, + prevOuts txscript.PrevOutputFetcher, tree *Node, + tweakLookup func(*Node) []byte) (*SignerSession, error) { + // Validate inputs. if signer == nil { return nil, fmt.Errorf("signer cannot be nil") @@ -207,8 +218,17 @@ func NewSignerSession(signer input.MuSig2Signer, return fmt.Errorf("failed to create tx: %w", err) } + taprootTweak := sweepTapscriptRoot + if tweakLookup != nil { + taprootTweak = tweakLookup(node) + if len(taprootTweak) == 0 { + return fmt.Errorf("missing taproot tweak "+ + "for node %s", node.Input) + } + } + session, err := node.NewTxSignerSession( - signer, sweepTapscriptRoot, signerKey, prevOuts, + signer, taprootTweak, signerKey, prevOuts, ) if err != nil { return fmt.Errorf("failed to create tx session: %w", diff --git a/lib/tree/signing_tweak_test.go b/lib/tree/signing_tweak_test.go new file mode 100644 index 000000000..2dbe6e1c0 --- /dev/null +++ b/lib/tree/signing_tweak_test.go @@ -0,0 +1,203 @@ +package tree + +import ( + "bytes" + "sort" + "testing" + + "github.com/btcsuite/btcd/btcec/v2" + "github.com/btcsuite/btcd/btcec/v2/schnorr/musig2" + "github.com/btcsuite/btcd/txscript/v2" + "github.com/btcsuite/btcd/wire/v2" + "github.com/lightningnetwork/lnd/input" + "github.com/lightningnetwork/lnd/keychain" + "github.com/stretchr/testify/require" +) + +// recordingMuSig2Signer records each session's taproot tweak. +type recordingMuSig2Signer struct { + *mockMuSig2Signer + + tweaks [][]byte +} + +// MuSig2CreateSession records the taproot tweak before creating a session. +func (r *recordingMuSig2Signer) MuSig2CreateSession(version input.MuSig2Version, + keyLoc keychain.KeyLocator, signers []*btcec.PublicKey, + tweaks *input.MuSig2Tweaks, otherNonces [][musig2.PubNonceSize]byte, + localNonces *musig2.Nonces) (*input.MuSig2SessionInfo, error) { + + r.tweaks = append( + r.tweaks, + append( + []byte(nil), tweaks.TaprootTweak..., + ), + ) + + return r.mockMuSig2Signer.MuSig2CreateSession( + version, keyLoc, signers, tweaks, otherNonces, localNonces, + ) +} + +// containsTweak reports whether the recorded tweaks include the given one. +func containsTweak(recorded [][]byte, want []byte) bool { + for _, tweak := range recorded { + if bytes.Equal(tweak, want) { + return true + } + } + + return false +} + +// newTweakTestPath builds a two-node path cosigned by pubKey. +func newTweakTestPath(t *testing.T, pubKey *btcec.PublicKey) (*Node, + wire.OutPoint, wire.OutPoint) { + + t.Helper() + + leaf := createSimpleLeaf( + "tweak-leaf", 1_000, []*btcec.PublicKey{pubKey}, + ) + root := createSimpleLeaf( + "tweak-root", 1_000, []*btcec.PublicKey{pubKey}, + ) + root.Children = map[uint32]*Node{0: leaf} + + rootTXID, err := root.TXID() + require.NoError(t, err) + leaf.Input = wire.OutPoint{Hash: rootTXID, Index: 0} + + return root, root.Input, leaf.Input +} + +// TestSignerSessionTweakLookup tests per-node signing tweaks. +func TestSignerSessionTweakLookup(t *testing.T) { + t.Parallel() + + privKey, pubKey := createTestKey(t) + sweepRoot := bytes.Repeat([]byte{0x01}, 32) + rootTweak := bytes.Repeat([]byte{0xAA}, 32) + leafTweak := bytes.Repeat([]byte{0xBB}, 32) + + t.Run("lookup covers every node", func(t *testing.T) { + t.Parallel() + + root, rootOutpoint, leafOutpoint := newTweakTestPath(t, pubKey) + tweaksByInput := map[wire.OutPoint][]byte{ + rootOutpoint: rootTweak, + leafOutpoint: leafTweak, + } + + signer := &recordingMuSig2Signer{ + mockMuSig2Signer: newMockMuSig2Signer(privKey), + } + fetcher, err := root.PrevOutputFetcher( + &wire.TxOut{ + Value: 5_000, + }, + ) + require.NoError(t, err) + + session, err := newSignerSession( + signer, &keychain.KeyDescriptor{PubKey: pubKey}, + sweepRoot, fetcher, root, + func(node *Node) []byte { + return tweaksByInput[node.Input] + }, + ) + require.NoError(t, err) + require.NotNil(t, session) + + require.Len(t, signer.tweaks, 2) + require.True(t, containsTweak(signer.tweaks, rootTweak)) + require.True(t, containsTweak(signer.tweaks, leafTweak)) + require.False(t, containsTweak(signer.tweaks, sweepRoot)) + }) + + t.Run("uncovered node rejected", func(t *testing.T) { + t.Parallel() + + root, rootOutpoint, _ := newTweakTestPath(t, pubKey) + tweaksByInput := map[wire.OutPoint][]byte{ + rootOutpoint: rootTweak, + } + + signer := &recordingMuSig2Signer{ + mockMuSig2Signer: newMockMuSig2Signer(privKey), + } + fetcher, err := root.PrevOutputFetcher( + &wire.TxOut{ + Value: 5_000, + }, + ) + require.NoError(t, err) + + _, err = newSignerSession( + signer, &keychain.KeyDescriptor{PubKey: pubKey}, + sweepRoot, fetcher, root, + func(node *Node) []byte { + return tweaksByInput[node.Input] + }, + ) + require.ErrorContains(t, err, "missing taproot tweak") + }) +} + +// TestComputeInternalKey tests untweaked MuSig2 key aggregation. +func TestComputeInternalKey(t *testing.T) { + t.Parallel() + + _, err := ComputeInternalKey(nil) + require.ErrorContains(t, err, "no cosigners") + + _, single := createTestKey(t) + got, err := ComputeInternalKey([]*btcec.PublicKey{single}) + require.NoError(t, err) + require.True(t, single.IsEqual(got)) + + cosigners := make([]*btcec.PublicKey, 3) + for i := range cosigners { + _, cosigners[i] = createTestKey(t) + } + tweak := bytes.Repeat([]byte{0x42}, 32) + + internal, err := ComputeInternalKey(cosigners) + require.NoError(t, err) + final, err := ComputeFinalKey(cosigners, tweak) + require.NoError(t, err) + + derived := txscript.ComputeTaprootOutputKey(internal, tweak) + require.Equal( + t, final.SerializeCompressed()[1:], + derived.SerializeCompressed()[1:], + ) +} + +// TestNodeSignerSessionDoesNotReorderCosigners tests cosigner slice ownership. +func TestNodeSignerSessionDoesNotReorderCosigners(t *testing.T) { + t.Parallel() + + privKey, pubKey := createTestKey(t) + _, secondKey := createTestKey(t) + _, thirdKey := createTestKey(t) + cosigners := []*btcec.PublicKey{pubKey, secondKey, thirdKey} + sort.Slice(cosigners, func(i, j int) bool { + return bytes.Compare( + cosigners[i].SerializeCompressed(), + cosigners[j].SerializeCompressed(), + ) > 0 + }) + original := append([]*btcec.PublicKey(nil), cosigners...) + + node := &Node{CoSigners: cosigners} + _, err := node.NewSignerSession( + &keychain.KeyDescriptor{ + PubKey: pubKey, + }, + newMockMuSig2Signer(privKey), + nil, + ) + require.NoError(t, err) + require.Equal(t, original, node.CoSigners) +} diff --git a/lib/tree/tree.go b/lib/tree/tree.go index 85b1b6045..d32173cf7 100644 --- a/lib/tree/tree.go +++ b/lib/tree/tree.go @@ -41,6 +41,10 @@ type Tree struct { // For VTXO trees, this is the operator's sweep script. // For connector trees, this is nil (no sweep script). SweepTapscriptRoot []byte + + // AssetContext contains the state of an asset tree. It is nil for + // Bitcoin-only trees. + AssetContext *AssetTreeContext } // NewTree constructs a transaction tree from the given leaves using BFS. @@ -132,6 +136,7 @@ func (t *Tree) ExtractPathForCoSigners(targetKeys ...*btcec.PublicKey) (*Tree, BatchOutpoint: t.BatchOutpoint, BatchOutput: t.BatchOutput, SweepTapscriptRoot: t.SweepTapscriptRoot, + AssetContext: t.AssetContext, }, nil } @@ -164,6 +169,7 @@ func (t *Tree) ExtractPathForIndices(leafIndices ...int) (*Tree, error) { BatchOutpoint: t.BatchOutpoint, BatchOutput: t.BatchOutput, SweepTapscriptRoot: t.SweepTapscriptRoot, + AssetContext: t.AssetContext, }, nil } @@ -309,8 +315,6 @@ func (t *Tree) PrettyPrint() string { } // NewTreeSignerSession creates a TreeSignerSession for this tree. -// This is a convenience wrapper that sets up the session with the tree's -// context. func (t *Tree) NewTreeSignerSession(wallet input.MuSig2Signer, signerKey *keychain.KeyDescriptor) (*SignerSession, error) { @@ -321,8 +325,14 @@ func (t *Tree) NewTreeSignerSession(wallet input.MuSig2Signer, "fetcher: %w", err) } - return NewSignerSession( + var tweakLookup func(*Node) []byte + if t.AssetContext != nil { + tweakLookup = t.AssetContext.tweakLookup() + } + + return newSignerSession( wallet, signerKey, t.SweepTapscriptRoot, prevOutFetcher, t.Root, + tweakLookup, ) } diff --git a/lib/tree/tree_structure.go b/lib/tree/tree_structure.go index ab0cd75f8..076b37fcb 100644 --- a/lib/tree/tree_structure.go +++ b/lib/tree/tree_structure.go @@ -26,6 +26,9 @@ type Structure struct { // Root is the root node of the tree structure. Root *Node + // AssetContext contains subtree amounts for an asset tree. + AssetContext *AssetTreeContext + // LeafScriptMap maps leaf node pointers to their output scripts // (pkscript). This is populated during structure building and used by // the BTC materializer. We keep BTC leaf data out of Node to keep the @@ -62,18 +65,35 @@ func BuildStructure(leaves []LeafDescriptor, // Create leaf scripts map for BTC path (maps leaf nodes to pkscripts). leafScripts := make(map[*Node][]byte) + // Track asset amounts separately from the Bitcoin tree nodes. + leafAssets := make(map[*Node]uint64) + // Build recursively from leaves up. root, err := buildStructureRecursive( leaves, cfg.OperatorKey, cfg.Radix, weightFn, leafScripts, + leafAssets, ) if err != nil { return nil, err } - return &Structure{ + structure := &Structure{ Root: root, LeafScriptMap: leafScripts, - }, nil + } + + if len(leafAssets) > 0 { + assetCtx := NewAssetTreeContext() + if _, err := aggregateAssetAmounts( + root, leafAssets, assetCtx, + ); err != nil { + return nil, err + } + + structure.AssetContext = assetCtx + } + + return structure, nil } // buildStructureRecursive recursively builds the tree structure. For a single @@ -81,11 +101,14 @@ func BuildStructure(leaves []LeafDescriptor, // creates a branch node with children built recursively. func buildStructureRecursive(leaves []LeafDescriptor, operatorKey *btcec.PublicKey, radix int, weightFn PartitionWeightFunc, - leafScripts map[*Node][]byte) (*Node, error) { + leafScripts map[*Node][]byte, + leafAssets map[*Node]uint64) (*Node, error) { // Base case: single leaf becomes a leaf node. if len(leaves) == 1 { - return buildLeafStructure(leaves[0], operatorKey, leafScripts) + return buildLeafStructure( + leaves[0], operatorKey, leafScripts, leafAssets, + ) } // Partition leaves into groups. @@ -103,6 +126,7 @@ func buildStructureRecursive(leaves []LeafDescriptor, child, err := buildStructureRecursive( group, operatorKey, radix, weightFn, leafScripts, + leafAssets, ) if err != nil { return nil, err @@ -135,7 +159,8 @@ func buildStructureRecursive(leaves []LeafDescriptor, // buildLeafStructure creates the structure for a leaf node and populates // the leafScripts map (for BTC trees). func buildLeafStructure(leaf LeafDescriptor, operatorKey *btcec.PublicKey, - leafScripts map[*Node][]byte) (*Node, error) { + leafScripts map[*Node][]byte, + leafAssets map[*Node]uint64) (*Node, error) { if leaf.CoSignerKey == nil { return nil, fmt.Errorf("leaf cosigner key cannot be nil") @@ -162,5 +187,9 @@ func buildLeafStructure(leaf LeafDescriptor, operatorKey *btcec.PublicKey, leafScripts[node] = leaf.PkScript } + if leafAssets != nil && leaf.AssetAmount > 0 { + leafAssets[node] = leaf.AssetAmount + } + return node, nil }