Congestion tree validation (#84)

* add common/pkg/tree validation

* update noah go.mod

* cleaning and fixes

* fix builder_test.go

* Fix deferred func

* fix even number of vtxos in congestion tree

---------

Co-authored-by: altafan <18440657+altafan@users.noreply.github.com>

asp/internal/core/application/service.go

@@ -227,7 +227,14 @@ func (s *service) startFinalization() {
 		return
 	}
 
+	var changes []domain.RoundEvent
 	defer func() {
+		if len(changes) > 0 {
+			if err := s.repoManager.Events().Save(ctx, round.Id, changes...); err != nil {
+				log.WithError(err).Warn("failed to store new round events")
+			}
+		}
+
 		if round.IsFailed() {
 			s.startRound()
 			return
@@ -240,14 +247,6 @@ func (s *service) startFinalization() {
 		return
 	}
 
-	var changes []domain.RoundEvent
-	defer func() {
-		if err := s.repoManager.Events().Save(ctx, round.Id, changes...); err != nil {
-			log.WithError(err).Warn("failed to store new round events")
-			return
-		}
-	}()
-
 	// TODO: understand how many payments must be popped from the queue and actually registered for the round
 	num := s.paymentRequests.len()
 	if num == 0 {
@@ -274,6 +273,8 @@ func (s *service) startFinalization() {
 		return
 	}
 
+	log.Debugf("pool tx created for round %s", round.Id)
+
 	connectors, forfeitTxs, err := s.builder.BuildForfeitTxs(s.pubkey, signedPoolTx, payments)
 	if err != nil {
 		changes = round.Fail(fmt.Errorf("failed to create connectors and forfeit txs: %s", err))
@@ -281,7 +282,14 @@ func (s *service) startFinalization() {
 		return
 	}
 
-	events, _ := round.StartFinalization(connectors, tree, signedPoolTx)
+	log.Debugf("forfeit transactions created for round %s", round.Id)
+
+	events, err := round.StartFinalization(connectors, tree, signedPoolTx)
+	if err != nil {
+		changes = round.Fail(fmt.Errorf("failed to start finalization: %s", err))
+		log.WithError(err).Warn("failed to start finalization")
+		return
+	}
 	changes = append(changes, events...)
 
 	s.forfeitTxs.push(forfeitTxs)

asp/internal/core/domain/congestion_tree.go

@@ -1,44 +0,0 @@
-package domain
-
-type Node struct {
-	Txid       string
-	Tx         string
-	ParentTxid string
-	Leaf       bool
-}
-
-type CongestionTree [][]Node
-
-func (c CongestionTree) Leaves() []Node {
-	leaves := c[len(c)-1]
-	for _, level := range c[:len(c)-1] {
-		for _, node := range level {
-			if node.Leaf {
-				leaves = append(leaves, node)
-			}
-		}
-	}
-
-	return leaves
-}
-
-func (c CongestionTree) Children(nodeTxid string) []Node {
-	var children []Node
-	for _, level := range c {
-		for _, node := range level {
-			if node.ParentTxid == nodeTxid {
-				children = append(children, node)
-			}
-		}
-	}
-
-	return children
-}
-
-func (c CongestionTree) NumberOfNodes() int {
-	var count int
-	for _, level := range c {
-		count += len(level)
-	}
-	return count
-}

asp/internal/core/domain/events.go

@@ -1,5 +1,7 @@
 package domain
 
+import "github.com/ark-network/ark/common/tree"
+
 type RoundEvent interface {
 	isEvent()
 }
@@ -17,7 +19,7 @@ type RoundStarted struct {
 
 type RoundFinalizationStarted struct {
 	Id                 string
-	CongestionTree     CongestionTree
+	CongestionTree     tree.CongestionTree
 	Connectors         []string
 	UnsignedForfeitTxs []string
 	PoolTx             string

asp/internal/core/domain/round.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"time"
 
+	"github.com/ark-network/ark/common/tree"
 	"github.com/google/uuid"
 )
 
@@ -41,7 +42,7 @@ type Round struct {
 	Txid              string
 	TxHex             string
 	ForfeitTxs        []string
-	CongestionTree    CongestionTree
+	CongestionTree    tree.CongestionTree
 	Connectors        []string
 	DustAmount        uint64
 	Version           uint
@@ -143,11 +144,11 @@ func (r *Round) RegisterPayments(payments []Payment) ([]RoundEvent, error) {
 	return []RoundEvent{event}, nil
 }
 
-func (r *Round) StartFinalization(connectors []string, tree CongestionTree, poolTx string) ([]RoundEvent, error) {
+func (r *Round) StartFinalization(connectors []string, congestionTree tree.CongestionTree, poolTx string) ([]RoundEvent, error) {
 	if len(connectors) <= 0 {
 		return nil, fmt.Errorf("missing list of connectors")
 	}
-	if len(tree) <= 0 {
+	if len(congestionTree) <= 0 {
 		return nil, fmt.Errorf("missing congestion tree")
 	}
 	if len(poolTx) <= 0 {
@@ -162,7 +163,7 @@ func (r *Round) StartFinalization(connectors []string, tree CongestionTree, pool
 
 	event := RoundFinalizationStarted{
 		Id:             r.Id,
-		CongestionTree: tree,
+		CongestionTree: congestionTree,
 		Connectors:     connectors,
 		PoolTx:         poolTx,
 	}

asp/internal/core/domain/round_test.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"testing"
 
+	"github.com/ark-network/ark/common/tree"
 	"github.com/ark-network/ark/internal/core/domain"
 	"github.com/stretchr/testify/require"
 )
@@ -72,7 +73,7 @@ var (
 	emptyTx        = "0200000000000000000000"
 	txid           = "0000000000000000000000000000000000000000000000000000000000000000"
 	pubkey         = "030000000000000000000000000000000000000000000000000000000000000001"
-	congestionTree = domain.CongestionTree{
+	congestionTree = tree.CongestionTree{
 		{
 			{
 				Txid:       txid,
@@ -318,7 +319,7 @@ func testStartFinalization(t *testing.T) {
 			fixtures := []struct {
 				round       *domain.Round
 				connectors  []string
-				tree        domain.CongestionTree
+				tree        tree.CongestionTree
 				poolTx      string
 				expectedErr string
 			}{

asp/internal/core/ports/tx_builder.go

@@ -1,6 +1,7 @@
 package ports
 
 import (
+	"github.com/ark-network/ark/common/tree"
 	"github.com/ark-network/ark/internal/core/domain"
 	"github.com/decred/dcrd/dcrec/secp256k1/v4"
 )
@@ -8,7 +9,7 @@ import (
 type TxBuilder interface {
 	BuildPoolTx(
 		aspPubkey *secp256k1.PublicKey, wallet WalletService, payments []domain.Payment, minRelayFee uint64,
-	) (poolTx string, congestionTree domain.CongestionTree, err error)
+	) (poolTx string, congestionTree tree.CongestionTree, err error)
 	BuildForfeitTxs(
 		aspPubkey *secp256k1.PublicKey, poolTx string, payments []domain.Payment,
 	) (connectors []string, forfeitTxs []string, err error)

asp/internal/infrastructure/db/service_test.go

@@ -6,6 +6,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/ark-network/ark/common/tree"
 	"github.com/ark-network/ark/internal/core/domain"
 	"github.com/ark-network/ark/internal/core/ports"
 	"github.com/ark-network/ark/internal/infrastructure/db"
@@ -21,7 +22,7 @@ const (
 	pubkey   = "0300000000000000000000000000000000000000000000000000000000000000001"
 )
 
-var congestionTree = [][]domain.Node{
+var congestionTree = [][]tree.Node{
 	{
 		{
 			Txid:       txid,

asp/internal/infrastructure/tx-builder/covenant/builder.go

@@ -4,7 +4,7 @@ import (
 	"context"
 	"encoding/hex"
 
-	"github.com/ark-network/ark/common"
+	"github.com/ark-network/ark/common/tree"
 	"github.com/ark-network/ark/internal/core/domain"
 	"github.com/ark-network/ark/internal/core/ports"
 	"github.com/decred/dcrd/dcrec/secp256k1/v4"
@@ -165,7 +165,7 @@ func (b *txBuilder) BuildPoolTx(
 	wallet ports.WalletService,
 	payments []domain.Payment,
 	minRelayFee uint64,
-) (poolTx string, congestionTree domain.CongestionTree, err error) {
+) (poolTx string, congestionTree tree.CongestionTree, err error) {
 	aspScriptBytes, err := p2wpkhScript(aspPubkey, b.net)
 	if err != nil {
 		return
@@ -226,12 +226,12 @@ func (b *txBuilder) BuildPoolTx(
 }
 
 func (b *txBuilder) getLeafTaprootTree(userPubkey, aspPubkey *secp256k1.PublicKey) ([]byte, *taproot.IndexedElementsTapScriptTree, error) {
-	sweepTaprootLeaf, err := sweepTapLeaf(aspPubkey)
+	sweepTaprootLeaf, err := tree.SweepScript(aspPubkey, expirationTime)
 	if err != nil {
 		return nil, nil, err
 	}
 
-	vtxoLeaf, err := common.VtxoScript(userPubkey)
+	vtxoLeaf, err := tree.VtxoScript(userPubkey)
 	if err != nil {
 		return nil, nil, err
 	}
@@ -239,7 +239,7 @@ func (b *txBuilder) getLeafTaprootTree(userPubkey, aspPubkey *secp256k1.PublicKe
 	leafTaprootTree := taproot.AssembleTaprootScriptTree(*vtxoLeaf, *sweepTaprootLeaf)
 	root := leafTaprootTree.RootNode.TapHash()
 
-	unspendableKeyBytes, _ := hex.DecodeString(unspendablePoint)
+	unspendableKeyBytes, _ := hex.DecodeString(tree.UnspendablePoint)
 	unspendableKey, _ := secp256k1.ParsePubKey(unspendableKeyBytes)
 
 	taprootKey := taproot.ComputeTaprootOutputKey(

asp/internal/infrastructure/tx-builder/covenant/builder_test.go

@@ -5,17 +5,16 @@ import (
 	"testing"
 
 	"github.com/ark-network/ark/common"
+	"github.com/ark-network/ark/common/tree"
 	"github.com/ark-network/ark/internal/core/domain"
 	"github.com/ark-network/ark/internal/core/ports"
 	txbuilder "github.com/ark-network/ark/internal/infrastructure/tx-builder/covenant"
-	"github.com/btcsuite/btcd/btcec/v2/schnorr"
 	"github.com/btcsuite/btcd/chaincfg/chainhash"
 	secp256k1 "github.com/decred/dcrd/dcrec/secp256k1/v4"
 	"github.com/stretchr/testify/require"
 	"github.com/vulpemventures/go-elements/network"
 	"github.com/vulpemventures/go-elements/payment"
 	"github.com/vulpemventures/go-elements/psetv2"
-	"github.com/vulpemventures/go-elements/taproot"
 	"github.com/vulpemventures/go-elements/transaction"
 )
 
@@ -124,9 +123,36 @@ func TestBuildCongestionTree(t *testing.T) {
 
 	fixtures := []struct {
 		payments          []domain.Payment
-		expectedNodesNum  int // 2*len(receivers)-1
+		expectedNodesNum  int // 2*len(receivers) -1
 		expectedLeavesNum int
 	}{
+		{
+			payments: []domain.Payment{
+				{
+					Id: "0",
+					Inputs: []domain.Vtxo{
+						{
+							VtxoKey: domain.VtxoKey{
+								Txid: "fd68e3c5796cc7db0a8036d486d5f625b6b2f2c014810ac020e1ac23e82c59d6",
+								VOut: 0,
+							},
+							Receiver: domain.Receiver{
+								Pubkey: "020000000000000000000000000000000000000000000000000000000000000002",
+								Amount: 1100,
+							},
+						},
+					},
+					Receivers: []domain.Receiver{
+						{
+							Pubkey: "020000000000000000000000000000000000000000000000000000000000000002",
+							Amount: 1100,
+						},
+					},
+				},
+			},
+			expectedNodesNum:  1,
+			expectedLeavesNum: 1,
+		},
 		{
 			payments: []domain.Payment{
 				{
@@ -238,6 +264,88 @@ func TestBuildCongestionTree(t *testing.T) {
 			},
 			expectedNodesNum:  5,
 			expectedLeavesNum: 3,
+		}, {
+			payments: []domain.Payment{
+				{
+					Id: "a242cdd8-f3d5-46c0-ae98-94135a2bee3f",
+					Inputs: []domain.Vtxo{
+						{
+							VtxoKey: domain.VtxoKey{
+								Txid: "755c820771284d85ea4bbcc246565b4eddadc44237a7e57a0f9cb78a840d1d41",
+								VOut: 0,
+							},
+							Receiver: domain.Receiver{
+								Pubkey: "02c87e5c1758df5ad42a918ec507b6e8dfcdcebf22f64f58eb4ad5804257d658a5",
+								Amount: 1000,
+							},
+						},
+						{
+							VtxoKey: domain.VtxoKey{
+								Txid: "66a0df86fcdeb84b8877adfe0b2c556dba30305d72ddbd4c49355f6930355357",
+								VOut: 0,
+							},
+							Receiver: domain.Receiver{
+								Pubkey: "02c87e5c1758df5ad42a918ec507b6e8dfcdcebf22f64f58eb4ad5804257d658a5",
+								Amount: 1000,
+							},
+						},
+						{
+							VtxoKey: domain.VtxoKey{
+								Txid: "9913159bc7aa493ca53cbb9cbc88f97ba01137c814009dc7ef520c3fafc67909",
+								VOut: 1,
+							},
+							Receiver: domain.Receiver{
+								Pubkey: "02c87e5c1758df5ad42a918ec507b6e8dfcdcebf22f64f58eb4ad5804257d658a5",
+								Amount: 500,
+							},
+						},
+						{
+							VtxoKey: domain.VtxoKey{
+								Txid: "5e10e77a7cdedc153be5193a4b6055a7802706ded4f2a9efefe86ed2f9a6ae60",
+								VOut: 0,
+							},
+							Receiver: domain.Receiver{
+								Pubkey: "02c87e5c1758df5ad42a918ec507b6e8dfcdcebf22f64f58eb4ad5804257d658a5",
+								Amount: 1000,
+							},
+						},
+						{
+							VtxoKey: domain.VtxoKey{
+								Txid: "5e10e77a7cdedc153be5193a4b6055a7802706ded4f2a9efefe86ed2f9a6ae60",
+								VOut: 1,
+							},
+							Receiver: domain.Receiver{
+								Pubkey: "02c87e5c1758df5ad42a918ec507b6e8dfcdcebf22f64f58eb4ad5804257d658a5",
+								Amount: 1000,
+							},
+						},
+					},
+					Receivers: []domain.Receiver{
+						{
+							Pubkey: "02c87e5c1758df5ad42a918ec507b6e8dfcdcebf22f64f58eb4ad5804257d658a5",
+							Amount: 1000,
+						},
+						{
+							Pubkey: "02c87e5c1758df5ad42a918ec507b6e8dfcdcebf22f64f58eb4ad5804257d658a5",
+							Amount: 1000,
+						},
+						{
+							Pubkey: "02c87e5c1758df5ad42a918ec507b6e8dfcdcebf22f64f58eb4ad5804257d658a5",
+							Amount: 1000,
+						},
+						{
+							Pubkey: "02c87e5c1758df5ad42a918ec507b6e8dfcdcebf22f64f58eb4ad5804257d658a5",
+							Amount: 1000,
+						},
+						{
+							Pubkey: "02c87e5c1758df5ad42a918ec507b6e8dfcdcebf22f64f58eb4ad5804257d658a5",
+							Amount: 500,
+						},
+					},
+				},
+			},
+			expectedNodesNum:  4,
+			expectedLeavesNum: 3,
 		},
 	}
 
@@ -246,54 +354,19 @@ func TestBuildCongestionTree(t *testing.T) {
 	require.NotNil(t, key)
 
 	for _, f := range fixtures {
-		poolTx, tree, err := builder.BuildPoolTx(key, &mockedWalletService{}, f.payments, 30)
-
+		poolTx, congestionTree, err := builder.BuildPoolTx(key, &mockedWalletService{}, f.payments, 30)
 		require.NoError(t, err)
-		require.Equal(t, f.expectedNodesNum, tree.NumberOfNodes())
-		require.Len(t, tree.Leaves(), f.expectedLeavesNum)
+		require.Equal(t, f.expectedNodesNum, congestionTree.NumberOfNodes())
+		require.Len(t, congestionTree.Leaves(), f.expectedLeavesNum)
 
-		poolTransaction, err := transaction.NewTxFromHex(poolTx)
+		// check that the pool tx has the right number of inputs and outputs
+		err = tree.ValidateCongestionTree(
+			congestionTree,
+			poolTx,
+			key,
+			1209344, // 2 weeks - 8 minutes
+		)
 		require.NoError(t, err)
-
-		poolTxID := poolTransaction.TxHash().String()
-
-		// check the root
-		require.Len(t, tree[0], 1)
-		require.Equal(t, poolTxID, tree[0][0].ParentTxid)
-
-		// check the nodes
-		for _, level := range tree {
-			for _, node := range level {
-				pset, err := psetv2.NewPsetFromBase64(node.Tx)
-				require.NoError(t, err)
-
-				require.Len(t, pset.Inputs, 1)
-				require.Len(t, pset.Outputs, 3)
-
-				inputTxID := chainhash.Hash(pset.Inputs[0].PreviousTxid).String()
-				require.Equal(t, node.ParentTxid, inputTxID)
-
-				children := tree.Children(node.Txid)
-				if len(children) > 0 {
-					require.Len(t, children, 2)
-
-					for i, child := range children {
-						childTx, err := psetv2.NewPsetFromBase64(child.Tx)
-						require.NoError(t, err)
-
-						for _, leaf := range childTx.Inputs[0].TapLeafScript {
-							key := leaf.ControlBlock.InternalKey
-							rootHash := leaf.ControlBlock.RootHash(leaf.Script)
-
-							outputScript := taproot.ComputeTaprootOutputKey(key, rootHash)
-							previousScriptKey := pset.Outputs[i].Script[2:]
-							require.Len(t, previousScriptKey, 32)
-							require.Equal(t, schnorr.SerializePubKey(outputScript), previousScriptKey)
-						}
-					}
-				}
-			}
-		}
 	}
 }
 

asp/internal/infrastructure/tx-builder/covenant/forfeit.go

@@ -3,6 +3,7 @@ package txbuilder
 import (
 	"encoding/hex"
 
+	"github.com/ark-network/ark/common/tree"
 	"github.com/btcsuite/btcd/txscript"
 	"github.com/decred/dcrd/dcrec/secp256k1/v4"
 	"github.com/vulpemventures/go-elements/elementsutil"
@@ -51,7 +52,7 @@ func createForfeitTx(
 		return "", err
 	}
 
-	unspendableKeyBytes, _ := hex.DecodeString(unspendablePoint)
+	unspendableKeyBytes, _ := hex.DecodeString(tree.UnspendablePoint)
 	unspendableKey, _ := secp256k1.ParsePubKey(unspendableKeyBytes)
 
 	for _, proof := range vtxoTaprootTree.LeafMerkleProofs {

asp/internal/infrastructure/tx-builder/covenant/tree.go

@@ -1,11 +1,10 @@
 package txbuilder
 
 import (
-	"encoding/binary"
 	"encoding/hex"
 	"fmt"
 
-	"github.com/ark-network/ark/common"
+	"github.com/ark-network/ark/common/tree"
 	"github.com/ark-network/ark/internal/core/domain"
 	"github.com/btcsuite/btcd/btcec/v2/schnorr"
 	"github.com/btcsuite/btcd/chaincfg/chainhash"
@@ -17,112 +16,11 @@ import (
 )
 
 const (
-	OP_INSPECTOUTPUTSCRIPTPUBKEY = 0xd1
-	OP_INSPECTOUTPUTVALUE        = 0xcf
-	OP_PUSHCURRENTINPUTINDEX     = 0xcd
-	unspendablePoint             = "0250929b74c1a04954b78b4b6035e97a5e078a5a0f28ec96d547bfee9ace803ac0"
-	timeDelta                    = 60 * 60 * 24 * 14 // 14 days in seconds
+	expirationTime = 60 * 60 * 24 * 14 // 14 days in seconds
 )
 
 // the private method buildCongestionTree returns a function letting to plug in the pool transaction output as input of the tree's root node
-type pluggableCongestionTree func(outpoint psetv2.InputArgs) (domain.CongestionTree, error)
-
-// withOutput returns an introspection script that checks the script and the amount of the output at the given index
-// verify will add an OP_EQUALVERIFY at the end of the script, otherwise it will add an OP_EQUAL
-func withOutput(index byte, taprootWitnessProgram []byte, amount uint64, verify bool) []byte {
-	amountBuffer := make([]byte, 8)
-	binary.LittleEndian.PutUint64(amountBuffer, amount)
-
-	script := []byte{
-		index,
-		OP_INSPECTOUTPUTSCRIPTPUBKEY,
-		txscript.OP_1,
-		txscript.OP_EQUALVERIFY,
-		txscript.OP_DATA_32,
-	}
-
-	script = append(script, taprootWitnessProgram...)
-	script = append(script, []byte{
-		txscript.OP_EQUALVERIFY,
-	}...)
-	script = append(script, index)
-	script = append(script, []byte{
-		OP_INSPECTOUTPUTVALUE,
-		txscript.OP_1,
-		txscript.OP_EQUALVERIFY,
-		txscript.OP_DATA_8,
-	}...)
-	script = append(script, amountBuffer...)
-	if verify {
-		script = append(script, []byte{
-			txscript.OP_EQUALVERIFY,
-		}...)
-	} else {
-		script = append(script, []byte{
-			txscript.OP_EQUAL,
-		}...)
-	}
-
-	return script
-}
-
-func checksigScript(pubkey *secp256k1.PublicKey) ([]byte, error) {
-	key := schnorr.SerializePubKey(pubkey)
-	return txscript.NewScriptBuilder().AddData(key).AddOp(txscript.OP_CHECKSIG).Script()
-}
-
-// checkSequenceVerifyScript without checksig
-func checkSequenceVerifyScript(seconds uint) ([]byte, error) {
-	sequence, err := common.BIP68Encode(seconds)
-	if err != nil {
-		return nil, err
-	}
-
-	return append(sequence, []byte{
-		txscript.OP_CHECKSEQUENCEVERIFY,
-		txscript.OP_DROP,
-	}...), nil
-}
-
-// checkSequenceVerifyScript + checksig
-func csvChecksigScript(pubkey *secp256k1.PublicKey, seconds uint) ([]byte, error) {
-	script, err := checksigScript(pubkey)
-	if err != nil {
-		return nil, err
-	}
-
-	csvScript, err := checkSequenceVerifyScript(seconds)
-	if err != nil {
-		return nil, err
-	}
-
-	return append(csvScript, script...), nil
-}
-
-// sweepTapLeaf returns a taproot leaf letting the owner of the key to spend the output after a given timeDelta
-func sweepTapLeaf(sweepKey *secp256k1.PublicKey) (*taproot.TapElementsLeaf, error) {
-	sweepScript, err := csvChecksigScript(sweepKey, timeDelta)
-	if err != nil {
-		return nil, err
-	}
-
-	tapLeaf := taproot.NewBaseTapElementsLeaf(sweepScript)
-	return &tapLeaf, nil
-}
-
-// forceSplitCoinTapLeaf returns a taproot leaf that enforces a split into two outputs
-// each output (left and right) will have the given amount and the given taproot key as witness program
-func forceSplitCoinTapLeaf(
-	leftKey, rightKey *secp256k1.PublicKey, leftAmount, rightAmount uint64,
-) taproot.TapElementsLeaf {
-	nextScriptLeft := withOutput(txscript.OP_0, schnorr.SerializePubKey(leftKey), leftAmount, rightKey != nil)
-	branchScript := append([]byte{}, nextScriptLeft...)
-	if rightKey != nil {
-		nextScriptRight := withOutput(txscript.OP_1, schnorr.SerializePubKey(rightKey), rightAmount, false)
-		branchScript = append(branchScript, nextScriptRight...)
-	}
-	return taproot.NewBaseTapElementsLeaf(branchScript)
-}
+type pluggableCongestionTree func(outpoint psetv2.InputArgs) (tree.CongestionTree, error)
 
 func taprootOutputScript(taprootKey *secp256k1.PublicKey) ([]byte, error) {
 	return txscript.NewScriptBuilder().AddOp(txscript.OP_1).AddData(schnorr.SerializePubKey(taprootKey)).Script()
@@ -165,7 +63,7 @@ func buildCongestionTree(
 	receivers []domain.Receiver,
 	feeSatsPerNode uint64,
 ) (pluggableTree pluggableCongestionTree, sharedOutputScript []byte, sharedOutputAmount uint64, err error) {
-	unspendableKeyBytes, err := hex.DecodeString(unspendablePoint)
+	unspendableKeyBytes, err := hex.DecodeString(tree.UnspendablePoint)
 	if err != nil {
 		return nil, nil, 0, err
 	}
@@ -203,7 +101,7 @@ func buildCongestionTree(
 	}
 
 	// compute the shared output script
-	sweepLeaf, err := sweepTapLeaf(aspPublicKey)
+	sweepLeaf, err := tree.VtxoScript(aspPublicKey)
 	if err != nil {
 		return nil, nil, 0, err
 	}
@@ -217,7 +115,7 @@ func buildCongestionTree(
 	var rightAmount uint64
 	var rightKey *secp256k1.PublicKey
 
-	if len(rootPset.Outputs) > 1 {
+	if len(rootPset.Outputs) > 2 {
 		rightAmount = rootPset.Outputs[1].Value
 		rightKey, err = schnorr.ParsePubKey(rootPset.Outputs[1].Script[2:])
 		if err != nil {
@@ -225,7 +123,7 @@ func buildCongestionTree(
 		}
 	}
 
-	goToTreeScript := forceSplitCoinTapLeaf(
+	goToTreeScript := tree.BranchScript(
 		leftKey, rightKey, leftOutput.Value, rightAmount,
 	)
 
@@ -237,7 +135,7 @@ func buildCongestionTree(
 		return nil, nil, 0, err
 	}
 
-	return func(outpoint psetv2.InputArgs) (domain.CongestionTree, error) {
+	return func(outpoint psetv2.InputArgs) (tree.CongestionTree, error) {
 		psets, err := nodes[0].psets(&psetArgs{
 			input:       outpoint,
 			taprootTree: taprootTree,
@@ -253,7 +151,7 @@ func buildCongestionTree(
 			}
 		}
 
-		tree := make(domain.CongestionTree, maxLevel+1)
+		congestionTree := make(tree.CongestionTree, maxLevel+1)
 
 		for _, psetWithLevel := range psets {
 			utx, err := psetWithLevel.pset.UnsignedTx()
@@ -270,7 +168,7 @@ func buildCongestionTree(
 
 			parentTxid := chainhash.Hash(psetWithLevel.pset.Inputs[0].PreviousTxid).String()
 
-			tree[psetWithLevel.level] = append(tree[psetWithLevel.level], domain.Node{
+			congestionTree[psetWithLevel.level] = append(congestionTree[psetWithLevel.level], tree.Node{
 				Txid:       txid,
 				Tx:         psetB64,
 				ParentTxid: parentTxid,
@@ -278,7 +176,7 @@ func buildCongestionTree(
 			})
 		}
 
-		return tree, nil
+		return congestionTree, nil
 	}, outputScript, uint64(rightAmount) + leftOutput.Value + uint64(feeSatsPerNode), nil
 }
 
@@ -349,7 +247,7 @@ func newBranch(
 }
 
 func (n *node) isLeaf() bool {
-	return n.left.isEmpty() && (n.right == nil || n.right.isEmpty())
+	return (n.left == nil || n.left.isEmpty()) && (n.right == nil || n.right.isEmpty())
 }
 
 // is it the final node of the tree
@@ -398,7 +296,7 @@ func (n *node) taprootKey() (*secp256k1.PublicKey, *taproot.IndexedElementsTapSc
 		return n._taprootKey, n._taprootTree, nil
 	}
 
-	sweepTaprootLeaf, err := sweepTapLeaf(n.sweepKey)
+	sweepTaprootLeaf, err := tree.SweepScript(n.sweepKey, expirationTime)
 	if err != nil {
 		return nil, nil, err
 	}
@@ -414,7 +312,7 @@ func (n *node) taprootKey() (*secp256k1.PublicKey, *taproot.IndexedElementsTapSc
 			return nil, nil, err
 		}
 
-		vtxoLeaf, err := common.VtxoScript(pubkey)
+		vtxoLeaf, err := tree.VtxoScript(pubkey)
 		if err != nil {
 			return nil, nil, err
 		}
@@ -443,7 +341,7 @@ func (n *node) taprootKey() (*secp256k1.PublicKey, *taproot.IndexedElementsTapSc
 		return nil, nil, err
 	}
 
-	branchTaprootLeaf := forceSplitCoinTapLeaf(
+	branchTaprootLeaf := tree.BranchScript(
 		leftKey, rightKey, n.left.amount(), n.right.amount(),
 	)
 
@@ -565,10 +463,10 @@ func (n *node) psets(inputArgs *psetArgs, level int) ([]psetWithLevel, error) {
 	}
 
 	nodeResult := []psetWithLevel{
-		{pset, level, n.isLeaf()},
+		{pset, level, n.isLeaf() || (n.left.isEmpty() || n.right.isEmpty())},
 	}
 
-	if n.left.isEmpty() && (n.right == nil || n.right.isEmpty()) {
+	if n.isLeaf() {
 		return nodeResult, nil
 	}
 
@@ -583,37 +481,46 @@ func (n *node) psets(inputArgs *psetArgs, level int) ([]psetWithLevel, error) {
 
 	txID := unsignedTx.TxHash().String()
 
-	_, leftTaprootTree, err := n.left.taprootKey()
-	if err != nil {
-		return nil, err
+	if !n.left.isEmpty() {
+		_, leftTaprootTree, err := n.left.taprootKey()
+		if err != nil {
+			return nil, err
+		}
+
+		psetsLeft, err := n.left.psets(&psetArgs{
+			input: psetv2.InputArgs{
+				Txid:    txID,
+				TxIndex: 0,
+			},
+			taprootTree: leftTaprootTree,
+		}, level+1)
+		if err != nil {
+			return nil, err
+		}
+
+		nodeResult = append(nodeResult, psetsLeft...)
 	}
 
-	psetsLeft, err := n.left.psets(&psetArgs{
-		input: psetv2.InputArgs{
-			Txid:    txID,
-			TxIndex: 0,
-		},
-		taprootTree: leftTaprootTree,
-	}, level+1)
-	if err != nil {
-		return nil, err
+	if !n.right.isEmpty() {
+
+		_, rightTaprootTree, err := n.right.taprootKey()
+		if err != nil {
+			return nil, err
+		}
+
+		psetsRight, err := n.right.psets(&psetArgs{
+			input: psetv2.InputArgs{
+				Txid:    txID,
+				TxIndex: 1,
+			},
+			taprootTree: rightTaprootTree,
+		}, level+1)
+		if err != nil {
+			return nil, err
+		}
+
+		nodeResult = append(nodeResult, psetsRight...)
 	}
 
-	_, rightTaprootTree, err := n.right.taprootKey()
-	if err != nil {
-		return nil, err
-	}
-
-	psetsRight, err := n.right.psets(&psetArgs{
-		input: psetv2.InputArgs{
-			Txid:    txID,
-			TxIndex: 1,
-		},
-		taprootTree: rightTaprootTree,
-	}, level+1)
-	if err != nil {
-		return nil, err
-	}
-
-	return append(nodeResult, append(psetsLeft, psetsRight...)...), nil
+	return nodeResult, nil
 }

asp/internal/infrastructure/tx-builder/dummy/builder.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/hex"
 
+	"github.com/ark-network/ark/common/tree"
 	"github.com/ark-network/ark/internal/core/domain"
 	"github.com/ark-network/ark/internal/core/ports"
 	"github.com/decred/dcrd/dcrec/secp256k1/v4"
@@ -92,7 +93,7 @@ func (b *txBuilder) BuildForfeitTxs(
 func (b *txBuilder) BuildPoolTx(
 	aspPubkey *secp256k1.PublicKey, wallet ports.WalletService, payments []domain.Payment,
 	minRelayFee uint64,
-) (poolTx string, congestionTree domain.CongestionTree, err error) {
+) (poolTx string, congestionTree tree.CongestionTree, err error) {
 	aspScriptBytes, err := p2wpkhScript(aspPubkey, b.net)
 	if err != nil {
 		return "", nil, err

asp/internal/infrastructure/tx-builder/dummy/tree.go

@@ -3,6 +3,7 @@ package txbuilder
 import (
 	"encoding/hex"
 
+	"github.com/ark-network/ark/common/tree"
 	"github.com/ark-network/ark/internal/core/domain"
 	"github.com/btcsuite/btcd/chaincfg/chainhash"
 	"github.com/decred/dcrd/dcrec/secp256k1/v4"
@@ -63,7 +64,7 @@ func buildCongestionTree(
 	net network.Network,
 	poolTxID string,
 	receivers []domain.Receiver,
-) (congestionTree domain.CongestionTree, err error) {
+) (congestionTree tree.CongestionTree, err error) {
 	var nodes []*node
 
 	for _, r := range receivers {
@@ -92,7 +93,7 @@ func buildCongestionTree(
 		}
 	}
 
-	tree := make(domain.CongestionTree, maxLevel+1)
+	congestionTree = make(tree.CongestionTree, maxLevel+1)
 
 	for _, psetWithLevel := range psets {
 		utx, err := psetWithLevel.pset.UnsignedTx()
@@ -109,7 +110,7 @@ func buildCongestionTree(
 
 		parentTxid := chainhash.Hash(psetWithLevel.pset.Inputs[0].PreviousTxid).String()
 
-		tree[psetWithLevel.level] = append(tree[psetWithLevel.level], domain.Node{
+		congestionTree[psetWithLevel.level] = append(congestionTree[psetWithLevel.level], tree.Node{
 			Txid:       txid,
 			Tx:         psetB64,
 			ParentTxid: parentTxid,
@@ -117,7 +118,7 @@ func buildCongestionTree(
 		})
 	}
 
-	return tree, nil
+	return congestionTree, nil
 }
 
 func createTreeLevel(nodes []*node) ([]*node, error) {

asp/internal/interface/grpc/handlers/arkservice.go

@@ -7,6 +7,7 @@ import (
 
 	arkv1 "github.com/ark-network/ark/api-spec/protobuf/gen/ark/v1"
 	"github.com/ark-network/ark/common"
+	"github.com/ark-network/ark/common/tree"
 	"github.com/ark-network/ark/internal/core/application"
 	"github.com/ark-network/ark/internal/core/domain"
 	"github.com/decred/dcrd/dcrec/secp256k1/v4"
@@ -280,10 +281,10 @@ func (v vtxoList) toProto(hrp string, aspKey *secp256k1.PublicKey) []*arkv1.Vtxo
 	return list
 }
 
-// castCongestionTree converts a domain.CongestionTree to a repeated arkv1.TreeLevel
-func castCongestionTree(tree domain.CongestionTree) *arkv1.Tree {
-	levels := make([]*arkv1.TreeLevel, 0, len(tree))
-	for _, level := range tree {
+// castCongestionTree converts a tree.CongestionTree to a repeated arkv1.TreeLevel
+func castCongestionTree(congestionTree tree.CongestionTree) *arkv1.Tree {
+	levels := make([]*arkv1.TreeLevel, 0, len(congestionTree))
+	for _, level := range congestionTree {
 		levelProto := &arkv1.TreeLevel{
 			Nodes: make([]*arkv1.Node, 0, len(level)),
 		}

asp/internal/interface/grpc/handlers/utils.go

@@ -18,7 +18,7 @@ func parseTxs(txs []string) ([]string, error) {
 	}
 	for _, tx := range txs {
 		if _, err := psetv2.NewPsetFromBase64(tx); err != nil {
-			return nil, fmt.Errorf("invalid tx format %s", err)
+			return nil, fmt.Errorf("invalid tx format")
 		}
 	}
 	return txs, nil