From 019c86382114a89fa7f56f8c21e284eada8fff18 Mon Sep 17 00:00:00 2001
From: Vincent <5889406+VincentDebug@users.noreply.github.com>
Date: Wed, 26 Jun 2024 22:39:28 +0900
Subject: [PATCH] Fix order of signatures in PartialSign function (#225)

Co-authored-by: Vincent <VincentDebug@users.noreply.github.com>
---
 transaction.go      | 16 ++++++++++++----
 transaction_test.go | 39 +++++++++++++++++++++++++++++++--------
 2 files changed, 43 insertions(+), 12 deletions(-)

diff --git a/transaction.go b/transaction.go
index 53a6f83e..dfa05955 100644
--- a/transaction.go
+++ b/transaction.go
@@ -536,18 +536,26 @@ func (tx *Transaction) PartialSign(getter privateKeyGetter) (out []Signature, er
 	}
 	signerKeys := tx.Message.signerKeys()
 
-	signedSignatures := []Signature{}
-	for _, key := range signerKeys {
+	// Ensure that the transaction has the correct number of signatures initialized
+	if len(tx.Signatures) == 0 {
+		// Initialize the Signatures slice to the correct length if it's empty
+		tx.Signatures = make([]Signature, len(signerKeys))
+	} else if len(tx.Signatures) != len(signerKeys) {
+		// Return an error if the current length of the Signatures slice doesn't match the expected number
+		return nil, fmt.Errorf("invalid signatures length, expected %d, actual %d", len(signerKeys), len(tx.Signatures))
+	}
+
+	for i, key := range signerKeys {
 		privateKey := getter(key)
 		if privateKey != nil {
 			s, err := privateKey.Sign(messageContent)
 			if err != nil {
 				return nil, fmt.Errorf("failed to signed with key %q: %w", key.String(), err)
 			}
-			signedSignatures = append(signedSignatures, s)
+			// Directly assign the signature to the corresponding position in the transaction's signature slice
+			tx.Signatures[i] = s
 		}
 	}
-	tx.Signatures = append(tx.Signatures, signedSignatures...)
 	return tx.Signatures, nil
 }
 
diff --git a/transaction_test.go b/transaction_test.go
index 3047134f..25384b6b 100644
--- a/transaction_test.go
+++ b/transaction_test.go
@@ -111,12 +111,14 @@ func TestPartialSignTransaction(t *testing.T) {
 	signers := []PrivateKey{
 		NewWallet().PrivateKey,
 		NewWallet().PrivateKey,
+		NewWallet().PrivateKey,
 	}
 	instructions := []Instruction{
 		&testTransactionInstructions{
 			accounts: []*AccountMeta{
 				{PublicKey: signers[0].PublicKey(), IsSigner: true, IsWritable: false},
 				{PublicKey: signers[1].PublicKey(), IsSigner: true, IsWritable: true},
+				{PublicKey: signers[2].PublicKey(), IsSigner: true, IsWritable: false},
 			},
 			data:      []byte{0xaa, 0xbb},
 			programID: MustPublicKeyFromBase58("11111111111111111111111111111111"),
@@ -129,16 +131,37 @@ func TestPartialSignTransaction(t *testing.T) {
 	trx, err := NewTransaction(instructions, blockhash)
 	require.NoError(t, err)
 
-	assert.Equal(t, trx.Message.Header.NumRequiredSignatures, uint8(2))
+	assert.Equal(t, trx.Message.Header.NumRequiredSignatures, uint8(3))
+
+	// Test various signing orders
+	signingOrders := [][]int{
+		{0, 1, 2}, // ABC
+		{0, 2, 1}, // ACB
+		{1, 0, 2}, // BAC
+		{1, 2, 0}, // BCA
+		{2, 0, 1}, // CAB
+		{2, 1, 0}, // CBA
+	}
+
+	for _, order := range signingOrders {
+		// Reset the transaction signatures before each test
+		trx.Signatures = make([]Signature, len(signers))
 
-	signatures, err := trx.PartialSign(func(key PublicKey) *PrivateKey {
-		if key.Equals(signers[0].PublicKey()) {
-			return &signers[0]
+		// Sign the transaction in the specified order
+		for _, idx := range order {
+			signer := signers[idx]
+			signatures, err := trx.PartialSign(func(key PublicKey) *PrivateKey {
+				if key.Equals(signer.PublicKey()) {
+					return &signer
+				}
+				return nil
+			})
+			require.NoError(t, err)
+			assert.Equal(t, len(signatures), 3)
 		}
-		return nil
-	})
-	require.NoError(t, err)
-	assert.Equal(t, len(signatures), 1)
+		// Verify Signatures
+		require.NoError(t, trx.VerifySignatures())
+	}
 }
 
 func TestSignTransaction(t *testing.T) {