Simulation framework for multi-client server testing

.gitignore

@@ -2,4 +2,9 @@
 .vscode/
 
 *.wasm
-wasm_exec.js
+wasm_exec.js
+
+.idea
+
+build
+log

common/bitcointree/musig2.go

@@ -2,6 +2,7 @@ package bitcointree
 
 import (
 	"bytes"
+	"encoding/binary"
 	"errors"
 	"fmt"
 	"io"
@@ -21,9 +22,6 @@ import (
 var (
 	ErrCongestionTreeNotSet = errors.New("congestion tree not set")
 	ErrAggregateKeyNotSet   = errors.New("aggregate key not set")
-
-	columnSeparator = byte('|')
-	rowSeparator    = byte('/')
 )
 
 type Musig2Nonce struct {
@@ -37,12 +35,15 @@ func (n *Musig2Nonce) Encode(w io.Writer) error {
 
 func (n *Musig2Nonce) Decode(r io.Reader) error {
 	bytes := make([]byte, 66)
-	_, err := r.Read(bytes)
+	bytesRead, err := io.ReadFull(r, bytes)
 	if err != nil {
 		return err
 	}
+	if bytesRead != 66 {
+		return fmt.Errorf("expected to read 66 bytes, but read %d", bytesRead)
+	}
 
-	n.PubNonce = [66]byte(bytes)
+	copy(n.PubNonce[:], bytes)
 	return nil
 }
 
@@ -585,57 +586,61 @@ type readable interface {
 func encodeMatrix[T writable](matrix [][]T) ([]byte, error) {
 	var buf bytes.Buffer
 
+	// Write number of rows
+	if err := binary.Write(&buf, binary.LittleEndian, uint32(len(matrix))); err != nil {
+		return nil, err
+	}
+
+	// For each row, write its length and then its elements
 	for _, row := range matrix {
+		// Write row length
+		if err := binary.Write(&buf, binary.LittleEndian, uint32(len(row))); err != nil {
+			return nil, err
+		}
+		// Write row data
 		for _, cell := range row {
-			if err := buf.WriteByte(columnSeparator); err != nil {
-				return nil, err
-			}
-
 			if err := cell.Encode(&buf); err != nil {
 				return nil, err
 			}
 		}
-
-		if err := buf.WriteByte(rowSeparator); err != nil {
-			return nil, err
-		}
 	}
 
 	return buf.Bytes(), nil
 }
 
 // decodeMatrix decode a byte stream into a matrix of serializable objects
 func decodeMatrix[T readable](factory func() T, data io.Reader) ([][]T, error) {
-	matrix := make([][]T, 0)
-	row := make([]T, 0)
-
-	for {
-		separator := make([]byte, 1)
+	var rowCount uint32
 
-		if _, err := data.Read(separator); err != nil {
-			if err == io.EOF {
-				break
-			}
-			return nil, err
-		}
+	// Read number of rows
+	if err := binary.Read(data, binary.LittleEndian, &rowCount); err != nil {
+		return nil, err
+	}
 
-		b := separator[0]
+	// Initialize matrix
+	matrix := make([][]T, rowCount)
 
-		if b == rowSeparator {
-			matrix = append(matrix, row)
-			row = make([]T, 0)
-			continue
+	// For each row, read its length and then its elements
+	for i := uint32(0); i < rowCount; i++ {
+		var colCount uint32
+		// Read row length
+		if err := binary.Read(data, binary.LittleEndian, &colCount); err != nil {
+			return nil, err
 		}
 
-		cell := factory()
+		// Initialize row
+		row := make([]T, colCount)
 
-		if err := cell.Decode(data); err != nil {
-			if err == io.EOF {
-				break
+		// Read row data
+		for j := uint32(0); j < colCount; j++ {
+			cell := factory()
+			if err := cell.Decode(data); err != nil {
+				return nil, err
 			}
-			return nil, err
+			row[j] = cell
 		}
-		row = append(row, cell)
+
+		matrix[i] = row
 	}
 
 	return matrix, nil

common/bitcointree/musig2_test.go

@@ -1,14 +1,14 @@
 package bitcointree_test
 
 import (
-	"bytes"
+	"encoding/hex"
 	"encoding/json"
-	"fmt"
+	"github.com/ark-network/ark/common/tree"
 	"os"
 	"testing"
 
 	"github.com/ark-network/ark/common/bitcointree"
-	"github.com/ark-network/ark/common/tree"
+	"github.com/btcsuite/btcd/btcec/v2"
 	"github.com/btcsuite/btcd/chaincfg/chainhash"
 	"github.com/btcsuite/btcd/txscript"
 	"github.com/btcsuite/btcd/wire"
@@ -27,36 +27,34 @@ var testTxid, _ = chainhash.NewHashFromStr("49f8664acc899be91902f8ade781b7eeb9cb
 func TestRoundTripSignTree(t *testing.T) {
 	fixtures := parseFixtures(t)
 	for _, f := range fixtures.Valid {
-		alice, err := secp256k1.GeneratePrivateKey()
-		require.NoError(t, err)
-
-		bob, err := secp256k1.GeneratePrivateKey()
-		require.NoError(t, err)
+		// Generate 20 cosigners
+		cosigners := make([]*secp256k1.PrivateKey, 20)
+		cosignerPubKeys := make([]*btcec.PublicKey, 20)
+		for i := 0; i < 20; i++ {
+			privKey, err := secp256k1.GeneratePrivateKey()
+			require.NoError(t, err)
+			cosigners[i] = privKey
+			cosignerPubKeys[i] = privKey.PubKey()
+		}
 
 		asp, err := secp256k1.GeneratePrivateKey()
 		require.NoError(t, err)
 
-		cosigners := make([]*secp256k1.PublicKey, 0)
-		cosigners = append(cosigners, alice.PubKey())
-		cosigners = append(cosigners, bob.PubKey())
-		cosigners = append(cosigners, asp.PubKey())
-
 		_, sharedOutputAmount, err := bitcointree.CraftSharedOutput(
-			cosigners,
+			cosignerPubKeys,
 			asp.PubKey(),
 			castReceivers(f.Receivers),
 			minRelayFee,
 			lifetime,
 		)
 		require.NoError(t, err)
 
-		// Create a new tree
 		tree, err := bitcointree.CraftCongestionTree(
 			&wire.OutPoint{
 				Hash:  *testTxid,
 				Index: 0,
 			},
-			cosigners,
+			cosignerPubKeys,
 			asp.PubKey(),
 			castReceivers(f.Receivers),
 			minRelayFee,
@@ -79,132 +77,48 @@ func TestRoundTripSignTree(t *testing.T) {
 			sharedOutputAmount,
 			tree,
 			root.CloneBytes(),
-			[]*secp256k1.PublicKey{alice.PubKey(), bob.PubKey(), asp.PubKey()},
+			cosignerPubKeys,
 		)
 		require.NoError(t, err)
 
-		aliceSession := bitcointree.NewTreeSignerSession(alice, sharedOutputAmount, tree, root.CloneBytes())
-		bobSession := bitcointree.NewTreeSignerSession(bob, sharedOutputAmount, tree, root.CloneBytes())
-		aspSession := bitcointree.NewTreeSignerSession(asp, sharedOutputAmount, tree, root.CloneBytes())
-
-		aliceNonces, err := aliceSession.GetNonces()
-		require.NoError(t, err)
-
-		bobNonces, err := bobSession.GetNonces()
-		require.NoError(t, err)
-
-		aspNonces, err := aspSession.GetNonces()
-		require.NoError(t, err)
-
-		s := bytes.NewBuffer(nil)
-
-		err = aspNonces[0][0].Encode(s)
-		require.NoError(t, err)
-
-		bitcointreeNonce := new(bitcointree.Musig2Nonce)
-		err = bitcointreeNonce.Decode(s)
-		require.NoError(t, err)
-
-		require.Equal(t, aspNonces[0][0], bitcointreeNonce)
-
-		var serializedNonces bytes.Buffer
-
-		err = aspNonces.Encode(&serializedNonces)
-		require.NoError(t, err)
-
-		decodedNonces, err := bitcointree.DecodeNonces(&serializedNonces)
-		require.NoError(t, err)
-
-		for i, nonces := range aspNonces {
-			for j, nonce := range nonces {
-				require.Equal(t, nonce.PubNonce, decodedNonces[i][j].PubNonce, fmt.Sprintf("matrix nonce not equal at index i: %d, j: %d", i, j))
-			}
+		// Create signer sessions for all cosigners
+		signerSessions := make([]bitcointree.SignerSession, 20)
+		for i, cosigner := range cosigners {
+			signerSessions[i] = bitcointree.NewTreeSignerSession(cosigner, sharedOutputAmount, tree, root.CloneBytes())
 		}
 
-		err = aspCoordinator.AddNonce(alice.PubKey(), aliceNonces)
-		require.NoError(t, err)
-
-		err = aspCoordinator.AddNonce(bob.PubKey(), bobNonces)
-		require.NoError(t, err)
-
-		err = aspCoordinator.AddNonce(asp.PubKey(), aspNonces)
-		require.NoError(t, err)
+		// Get nonces from all signers
+		for i, session := range signerSessions {
+			nonces, err := session.GetNonces()
+			require.NoError(t, err)
+			err = aspCoordinator.AddNonce(cosignerPubKeys[i], nonces)
+			require.NoError(t, err)
+		}
 
 		aggregatedNonce, err := aspCoordinator.AggregateNonces()
 		require.NoError(t, err)
 
-		// coordinator sends the combined nonce to all signers
-
-		err = aliceSession.SetKeys(
-			cosigners,
-		)
-		require.NoError(t, err)
-
-		err = aliceSession.SetAggregatedNonces(
-			aggregatedNonce,
-		)
-		require.NoError(t, err)
-
-		err = bobSession.SetKeys(
-			cosigners,
-		)
-		require.NoError(t, err)
-
-		err = bobSession.SetAggregatedNonces(
-			aggregatedNonce,
-		)
-		require.NoError(t, err)
-
-		err = aspSession.SetKeys(
-			cosigners,
-		)
-		require.NoError(t, err)
-
-		err = aspSession.SetAggregatedNonces(
-			aggregatedNonce,
-		)
-		require.NoError(t, err)
-
-		aliceSig, err := aliceSession.Sign()
-		require.NoError(t, err)
-
-		bobSig, err := bobSession.Sign()
-		require.NoError(t, err)
-
-		aspSig, err := aspSession.Sign()
-		require.NoError(t, err)
-
-		// check that the sigs are serializable
-
-		serializedSigs := bytes.NewBuffer(nil)
-
-		err = aspSig.Encode(serializedSigs)
-		require.NoError(t, err)
-
-		decodedSigs, err := bitcointree.DecodeSignatures(serializedSigs)
-		require.NoError(t, err)
-
-		for i, sigs := range aspSig {
-			for j, sig := range sigs {
-				require.Equal(t, sig.S, decodedSigs[i][j].S, fmt.Sprintf("matrix sig not equal at index i: %d, j: %d", i, j))
-			}
+		// Set keys and aggregated nonces for all signers
+		for _, session := range signerSessions {
+			err = session.SetKeys(cosignerPubKeys)
+			require.NoError(t, err)
+			err = session.SetAggregatedNonces(aggregatedNonce)
+			require.NoError(t, err)
 		}
 
-		// coordinator receives the signatures and combines them
-		err = aspCoordinator.AddSig(alice.PubKey(), aliceSig)
-		require.NoError(t, err)
-
-		err = aspCoordinator.AddSig(bob.PubKey(), bobSig)
-		require.NoError(t, err)
-
-		err = aspCoordinator.AddSig(asp.PubKey(), aspSig)
-		require.NoError(t, err)
+		// Get signatures from all signers
+		for i, session := range signerSessions {
+			sig, err := session.Sign()
+			require.NoError(t, err)
+			err = aspCoordinator.AddSig(cosignerPubKeys[i], sig)
+			require.NoError(t, err)
+		}
 
 		signedTree, err := aspCoordinator.SignTree()
 		require.NoError(t, err)
 
 		// verify the tree
-		aggregatedKey, err := bitcointree.AggregateKeys(cosigners, root.CloneBytes())
+		aggregatedKey, err := bitcointree.AggregateKeys(cosignerPubKeys, root.CloneBytes())
 		require.NoError(t, err)
 
 		err = bitcointree.ValidateTreeSigs(
@@ -222,6 +136,24 @@ type receiverFixture struct {
 	Pubkey string `json:"pubkey"`
 }
 
+func (r receiverFixture) toVtxoScript(asp *secp256k1.PublicKey) bitcointree.VtxoScript {
+	bytesKey, err := hex.DecodeString(r.Pubkey)
+	if err != nil {
+		panic(err)
+	}
+
+	pubkey, err := secp256k1.ParsePubKey(bytesKey)
+	if err != nil {
+		panic(err)
+	}
+
+	return &bitcointree.DefaultVtxoScript{
+		Owner:     pubkey,
+		Asp:       asp,
+		ExitDelay: exitDelay,
+	}
+}
+
 func castReceivers(receivers []receiverFixture) []tree.VtxoLeaf {
 	receiversOut := make([]tree.VtxoLeaf, 0, len(receivers))
 	for _, r := range receivers {