fix: harden client-side validation of counterparty data

pkg/ark-lib/tree/tx_tree.go

@@ -5,6 +5,7 @@ import (
 	"strings"
 
 	"github.com/btcsuite/btcd/psbt/v2"
+	"github.com/btcsuite/btcd/wire/v2"
 )
 
 // LeafOutput represents the output of a leaf transaction.
@@ -199,6 +200,7 @@ func (t *TxTree) SerializeNode() (*TxTreeNode, error) {
 // It verifies :
 // - every tx is a valid partial transaction.
 // - every tx has exactly one input.
+// - every tx is final, ie. it can be broadcasted as soon as its parent is confirmed
 // - the child txs spend the right parent's output
 // - the sum of the child txs' output amounts matches the parent tx input amount
 func (t *TxTree) Validate() error {
@@ -210,13 +212,28 @@ func (t *TxTree) Validate() error {
 		return fmt.Errorf("unexpected version: %d, expected 3", t.Root.UnsignedTx.Version)
 	}
 
+	// A non-final tx can't be broadcasted until its timelock elapses, which would
+	// hold the unroll path back while the batch output sweep matures.
+	if t.Root.UnsignedTx.LockTime != 0 {
+		return fmt.Errorf("unexpected locktime: %d, expected 0", t.Root.UnsignedTx.LockTime)
+	}
+
 	nbOfOutputs := uint32(len(t.Root.UnsignedTx.TxOut))
 	nbOfInputs := uint32(len(t.Root.UnsignedTx.TxIn))
 
 	if nbOfInputs != 1 {
 		return fmt.Errorf("unexpected number of inputs: %d, expected 1", nbOfInputs)
 	}
 
+	for inputIndex, input := range t.Root.UnsignedTx.TxIn {
+		if input.Sequence != wire.MaxTxInSequenceNum {
+			return fmt.Errorf(
+				"unexpected sequence for input %d: %d, expected %d",
+				inputIndex, input.Sequence, uint32(wire.MaxTxInSequenceNum),
+			)
+		}
+	}
+
 	// The children map can't be bigger than the number of outputs (excluding the P2A).
 	// A tx tree can be "partial" and specify only some of the outputs as children,
 	// that's why we allow len(g.Children) to be less than nbOfOutputs-1

pkg/ark-lib/tree/tx_tree_finality_test.go

@@ -0,0 +1,67 @@
+package tree_test
+
+import (
+	"testing"
+
+	"github.com/arkade-os/arkd/pkg/ark-lib/tree"
+	"github.com/btcsuite/btcd/wire/v2"
+	"github.com/stretchr/testify/require"
+)
+
+// allNodes returns every node of the tree, so a mutation can be applied to the
+// deeper ones too and only surface if Validate() really recurses.
+func allNodes(txTree *tree.TxTree) []*tree.TxTree {
+	nodes := []*tree.TxTree{txTree}
+	for _, child := range txTree.Children {
+		nodes = append(nodes, allNodes(child)...)
+	}
+	return nodes
+}
+
+// TestTxTreeValidateFinality checks that a tree holding a node that bitcoin
+// would not accept right away is rejected. A pre-signed unroll path is only
+// worth something if it can be broadcasted as soon as its parent confirms.
+func TestTxTreeValidateFinality(t *testing.T) {
+	testVectors, err := makeTestVectors()
+	require.NoError(t, err)
+	require.NotEmpty(t, testVectors)
+
+	for _, v := range testVectors {
+		t.Run(v.name, func(t *testing.T) {
+			vtxoTree, err := tree.BuildVtxoTree(
+				rootInput, v.receivers, batchOutSweepClosure[:], vtxoTreeExpiry,
+			)
+			require.NoError(t, err)
+
+			connectorTree, err := tree.BuildConnectorTree(rootInput, v.receivers)
+			require.NoError(t, err)
+
+			for name, txTree := range map[string]*tree.TxTree{
+				"vtxo tree": vtxoTree, "connector tree": connectorTree,
+			} {
+				t.Run(name, func(t *testing.T) {
+					// the tree as built by the honest stack must stay valid
+					require.NoError(t, txTree.Validate())
+
+					for i, node := range allNodes(txTree) {
+						node.Root.UnsignedTx.LockTime = 1
+						require.ErrorContainsf(
+							t, txTree.Validate(), "unexpected locktime",
+							"locktime not validated on node %d", i,
+						)
+						node.Root.UnsignedTx.LockTime = 0
+						require.NoError(t, txTree.Validate())
+
+						node.Root.UnsignedTx.TxIn[0].Sequence = wire.MaxTxInSequenceNum - 1
+						require.ErrorContainsf(
+							t, txTree.Validate(), "unexpected sequence",
+							"sequence not validated on node %d", i,
+						)
+						node.Root.UnsignedTx.TxIn[0].Sequence = wire.MaxTxInSequenceNum
+						require.NoError(t, txTree.Validate())
+					}
+				})
+			}
+		})
+	}
+}

pkg/client-lib/batch-session/handler/default_handler.go

@@ -179,6 +179,12 @@ func (h *defaultHandler) OnTreeSigningStarted(
 		return false, err
 	}
 
+	// the tree has to be verified before we send any nonce or signature for it,
+	// otherwise we pre-sign unroll paths we haven't checked.
+	if err := h.validateVtxoTreeAgainstCommitmentTx(commitmentTx, vtxoTree); err != nil {
+		return false, fmt.Errorf("failed to verify vtxo tree: %s", err)
+	}
+
 	batchOutput := commitmentTx.UnsignedTx.TxOut[0]
 	batchOutputAmount := batchOutput.Value
 
@@ -438,15 +444,11 @@ func (h *defaultHandler) vtxosToForfeit() []clientlib.Vtxo {
 	return withoutRecoverable
 }
 
-func (h *defaultHandler) validateVtxoTree(
-	event clientlib.BatchFinalizationEvent, vtxoTree, connectorTree *tree.TxTree,
+// validateVtxoTreeAgainstCommitmentTx groups the checks needing nothing but the
+// vtxo tree and the commitment tx, so that they can run before the tree is signed.
+func (h *defaultHandler) validateVtxoTreeAgainstCommitmentTx(
+	commitmentPtx *psbt.Packet, vtxoTree *tree.TxTree,
 ) error {
-	commitmentTx := event.Tx
-	commitmentPtx, err := psbt.NewFromRawBytes(strings.NewReader(commitmentTx), true)
-	if err != nil {
-		return err
-	}
-
 	// validate the vtxo tree is well formed
 	if !isOnchainOnly(h.Receivers) {
 		if err := tree.ValidateVtxoTree(
@@ -476,9 +478,23 @@ func (h *defaultHandler) validateVtxoTree(
 	}
 
 	// validate it contains our outputs
-	if err := validateReceivers(
+	return validateReceivers(
 		h.ServerParams.Network, commitmentPtx, h.Receivers, vtxoTree,
-	); err != nil {
+	)
+}
+
+func (h *defaultHandler) validateVtxoTree(
+	event clientlib.BatchFinalizationEvent, vtxoTree, connectorTree *tree.TxTree,
+) error {
+	commitmentTx := event.Tx
+	commitmentPtx, err := psbt.NewFromRawBytes(strings.NewReader(commitmentTx), true)
+	if err != nil {
+		return err
+	}
+
+	// re-run them against the commitment tx of this event, it may not be the one
+	// we were given when the signing session started.
+	if err := h.validateVtxoTreeAgainstCommitmentTx(commitmentPtx, vtxoTree); err != nil {
 		return err
 	}
 

pkg/client-lib/batch-session/handler/default_handler_test.go

@@ -0,0 +1,218 @@
+package batchsessionhandler
+
+import (
+	"context"
+	"encoding/hex"
+	"testing"
+
+	arklib "github.com/arkade-os/arkd/pkg/ark-lib"
+	"github.com/arkade-os/arkd/pkg/ark-lib/script"
+	"github.com/arkade-os/arkd/pkg/ark-lib/tree"
+	clientlib "github.com/arkade-os/arkd/pkg/client-lib"
+	"github.com/btcsuite/btcd/btcec/v2"
+	"github.com/btcsuite/btcd/chainhash/v2"
+	"github.com/btcsuite/btcd/psbt/v2"
+	"github.com/btcsuite/btcd/txscript/v2"
+	"github.com/btcsuite/btcd/wire/v2"
+	"github.com/stretchr/testify/require"
+)
+
+// recordingSignerSession records whether the tree ever reached the musig2 session.
+type recordingSignerSession struct {
+	tree.SignerSession
+	pubkey    string
+	initiated bool
+}
+
+func (s *recordingSignerSession) Init(_ []byte, _ int64, _ *tree.TxTree) error {
+	s.initiated = true
+	return nil
+}
+
+func (s *recordingSignerSession) GetPublicKey() string { return s.pubkey }
+
+func (s *recordingSignerSession) GetNonces() (tree.TreeNonces, error) {
+	return tree.TreeNonces{}, nil
+}
+
+// recordingClient records the nonce submissions so a test can assert nothing was
+// sent for a tree that didn't validate.
+type recordingClient struct {
+	clientlib.Client
+	submittedNonces int
+}
+
+func (c *recordingClient) SubmitTreeNonces(
+	_ context.Context, _, _ string, _ tree.TreeNonces,
+) error {
+	c.submittedNonces++
+	return nil
+}
+
+type treeSigningFixture struct {
+	handler  *defaultHandler
+	client   *recordingClient
+	session  *recordingSignerSession
+	event    clientlib.TreeSigningStartedEvent
+	vtxoTree *tree.TxTree
+}
+
+// newTreeSigningFixture builds the batch a honest server would propose: a
+// commitment tx paying the batch output and the vtxo tree spending it.
+func newTreeSigningFixture(t *testing.T) *treeSigningFixture {
+	t.Helper()
+
+	forfeitPrvkey, err := btcec.NewPrivateKey()
+	require.NoError(t, err)
+	cosignerPrvkey, err := btcec.NewPrivateKey()
+	require.NoError(t, err)
+	vtxoPrvkey, err := btcec.NewPrivateKey()
+	require.NoError(t, err)
+
+	batchExpiry := arklib.RelativeLocktime{Type: arklib.LocktimeTypeBlock, Value: 144}
+
+	sweepClosure := script.CSVMultisigClosure{
+		MultisigClosure: script.MultisigClosure{
+			PubKeys: []*btcec.PublicKey{forfeitPrvkey.PubKey()},
+		},
+		Locktime: batchExpiry,
+	}
+	sweepScript, err := sweepClosure.Script()
+	require.NoError(t, err)
+
+	sweepRoot := txscript.AssembleTaprootScriptTree(
+		txscript.NewBaseTapLeaf(sweepScript),
+	).RootNode.TapHash()
+
+	vtxoPkScript, err := script.P2TRScript(vtxoPrvkey.PubKey())
+	require.NoError(t, err)
+
+	const receiverAmount = 10000
+	leaf := tree.Leaf{
+		Outputs: []tree.LeafOutput{
+			{Amount: receiverAmount, Script: hex.EncodeToString(vtxoPkScript)},
+		},
+		CosignersPublicKeys: []string{
+			hex.EncodeToString(cosignerPrvkey.PubKey().SerializeCompressed()),
+		},
+	}
+
+	batchOutScript, batchOutAmount, err := tree.BuildBatchOutput(
+		[]tree.Leaf{leaf}, sweepRoot[:],
+	)
+	require.NoError(t, err)
+
+	prevoutHash, err := chainhash.NewHashFromStr(
+		"49f8664acc899be91902f8ade781b7eeb9cbe22bdd9efbc36e56195de21bcd12",
+	)
+	require.NoError(t, err)
+
+	commitmentTx, err := psbt.New(
+		[]*wire.OutPoint{{Hash: *prevoutHash, Index: 0}},
+		[]*wire.TxOut{{Value: batchOutAmount, PkScript: batchOutScript}},
+		3, 0, []uint32{wire.MaxTxInSequenceNum},
+	)
+	require.NoError(t, err)
+
+	commitmentTxHash := commitmentTx.UnsignedTx.TxHash()
+	vtxoTree, err := tree.BuildVtxoTree(
+		&wire.OutPoint{Hash: commitmentTxHash, Index: 0},
+		[]tree.Leaf{leaf}, sweepRoot[:], batchExpiry,
+	)
+	require.NoError(t, err)
+
+	encodedCommitmentTx, err := commitmentTx.B64Encode()
+	require.NoError(t, err)
+
+	addr := arklib.Address{
+		HRP:        arklib.BitcoinRegTest.Addr,
+		Signer:     forfeitPrvkey.PubKey(),
+		VtxoTapKey: vtxoPrvkey.PubKey(),
+	}
+	encodedAddr, err := addr.EncodeV0()
+	require.NoError(t, err)
+
+	client := &recordingClient{}
+	session := &recordingSignerSession{
+		pubkey: hex.EncodeToString(cosignerPrvkey.PubKey().SerializeCompressed()),
+	}
+
+	return &treeSigningFixture{
+		handler: &defaultHandler{
+			Args: Args{
+				Client: client,
+				ServerParams: clientlib.ServerParams{
+					Network:       arklib.BitcoinRegTest,
+					ForfeitPubKey: forfeitPrvkey.PubKey(),
+				},
+				Receivers:      []clientlib.Receiver{{To: encodedAddr, Amount: receiverAmount}},
+				SignerSessions: []tree.SignerSession{session},
+			},
+			batchExpiry: batchExpiry,
+		},
+		client:  client,
+		session: session,
+		event: clientlib.TreeSigningStartedEvent{
+			Id:                   "batch-id",
+			UnsignedCommitmentTx: encodedCommitmentTx,
+			CosignersPubkeys:     []string{session.pubkey},
+		},
+		vtxoTree: vtxoTree,
+	}
+}
+
+func TestOnTreeSigningStartedValidatesBeforeSigning(t *testing.T) {
+	// the tree of a honest batch must still reach the musig2 session
+	t.Run("valid", func(t *testing.T) {
+		f := newTreeSigningFixture(t)
+
+		skip, err := f.handler.OnTreeSigningStarted(
+			context.Background(), f.event, f.vtxoTree,
+		)
+		require.NoError(t, err)
+		require.False(t, skip)
+		require.True(t, f.session.initiated)
+		require.Equal(t, 1, f.client.submittedNonces)
+	})
+
+	// a tree not spending the batch output must be rejected before we contribute
+	// anything to the signing session
+	t.Run("not spending the batch output", func(t *testing.T) {
+		f := newTreeSigningFixture(t)
+		f.vtxoTree.Root.UnsignedTx.TxIn[0].PreviousOutPoint.Index = 1
+
+		skip, err := f.handler.OnTreeSigningStarted(
+			context.Background(), f.event, f.vtxoTree,
+		)
+		require.ErrorContains(t, err, "failed to verify vtxo tree")
+		require.False(t, skip)
+		require.False(t, f.session.initiated)
+		require.Zero(t, f.client.submittedNonces)
+	})
+
+	// same for a tree we couldn't broadcast when we need to unroll
+	t.Run("non final", func(t *testing.T) {
+		f := newTreeSigningFixture(t)
+		f.vtxoTree.Root.UnsignedTx.TxIn[0].Sequence = wire.MaxTxInSequenceNum - 1
+
+		_, err := f.handler.OnTreeSigningStarted(
+			context.Background(), f.event, f.vtxoTree,
+		)
+		require.ErrorContains(t, err, "unexpected sequence")
+		require.False(t, f.session.initiated)
+		require.Zero(t, f.client.submittedNonces)
+	})
+
+	// and for a tree not paying what we asked for
+	t.Run("wrong receiver amount", func(t *testing.T) {
+		f := newTreeSigningFixture(t)
+		f.handler.Receivers[0].Amount++
+
+		_, err := f.handler.OnTreeSigningStarted(
+			context.Background(), f.event, f.vtxoTree,
+		)
+		require.ErrorContains(t, err, "offchain send output not found")
+		require.False(t, f.session.initiated)
+		require.Zero(t, f.client.submittedNonces)
+	})
+}