From e15a5c010e929082357c5792b6cf1bca538f0d8c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Varga-Somogyi=20=C3=81kos?=
 <vsakos@users.noreply.github.com>
Date: Thu, 22 Sep 2022 02:22:40 +0200
Subject: [PATCH] feat: implement addSignature in VersionedTransaction (#27945)

* Implement addSignature in VersionedTransaction

* Update asserts

* VersionedTransaction.addSignature test

* Update VersionedTransaction tests
---
 web3.js/src/transaction/versioned.ts |  17 ++++
 web3.js/test/transaction.test.ts     | 115 ++++++++++++++++++++++-----
 2 files changed, 112 insertions(+), 20 deletions(-)

diff --git a/web3.js/src/transaction/versioned.ts b/web3.js/src/transaction/versioned.ts
index 0eac6dd0989a06..53e78ade1fff09 100644
--- a/web3.js/src/transaction/versioned.ts
+++ b/web3.js/src/transaction/versioned.ts
@@ -7,6 +7,7 @@ import {SIGNATURE_LENGTH_IN_BYTES} from './constants';
 import * as shortvec from '../utils/shortvec-encoding';
 import * as Layout from '../layout';
 import {sign} from '../utils/ed25519';
+import {PublicKey} from '../publickey';
 
 export type TransactionVersion = 'legacy' | 0;
 
@@ -106,4 +107,20 @@ export class VersionedTransaction {
       this.signatures[signerIndex] = sign(messageData, signer.secretKey);
     }
   }
+
+  addSignature(publicKey: PublicKey, signature: Uint8Array) {
+    assert(signature.byteLength === 64, 'Signature must be 64 bytes long');
+    const signerPubkeys = this.message.staticAccountKeys.slice(
+      0,
+      this.message.header.numRequiredSignatures,
+    );
+    const signerIndex = signerPubkeys.findIndex(pubkey =>
+      pubkey.equals(publicKey),
+    );
+    assert(
+      signerIndex >= 0,
+      `Can not add signature; \`${publicKey.toBase58()}\` is not required to sign this transaction`,
+    );
+    this.signatures[signerIndex] = signature;
+  }
 }
diff --git a/web3.js/test/transaction.test.ts b/web3.js/test/transaction.test.ts
index 9786384a2c2783..6fec2cd8e14af8 100644
--- a/web3.js/test/transaction.test.ts
+++ b/web3.js/test/transaction.test.ts
@@ -8,6 +8,7 @@ import {PublicKey} from '../src/publickey';
 import {
   Transaction,
   TransactionInstruction,
+  TransactionMessage,
   VersionedTransaction,
 } from '../src/transaction';
 import {StakeProgram, SystemProgram} from '../src/programs';
@@ -860,26 +861,6 @@ describe('Transaction', () => {
     expect(tx.verifySignatures()).to.be.true;
   });
 
-  it('deserializes versioned transactions', () => {
-    const serializedVersionedTx = Buffer.from(
-      'AdTIDASR42TgVuXKkd7mJKk373J3LPVp85eyKMVcrboo9KTY8/vm6N/Cv0NiHqk2I8iYw6VX5ZaBKG8z' +
-        '9l1XjwiAAQACA+6qNbqfjaIENwt9GzEK/ENiB/ijGwluzBUmQ9xlTAMcCaS0ctnyxTcXXlJr7u2qtnaM' +
-        'gIAO2/c7RBD0ipHWUcEDBkZv5SEXMv/srbpyw5vnvIzlu8X3EmssQ5s6QAAAAJbI7VNs6MzREUlnzRaJ' +
-        'pBKP8QQoDn2dWQvD0KIgHFDiAwIACQAgoQcAAAAAAAIABQEAAAQAATYPBwAKBDIBAyQWIw0oCxIdCA4i' +
-        'JzQRKwUZHxceHCohMBUJJiwpMxAaGC0TLhQxGyAMBiU2NS8VDgAAAADuAgAAAAAAAAIAAAAAAAAAAdGCT' +
-        'Qiq5yw3+3m1sPoRNj0GtUNNs0FIMocxzt3zuoSZHQABAwQFBwgLDA8RFBcYGhwdHh8iIyUnKiwtLi8yF' +
-        'wIGCQoNDhASExUWGRsgISQmKCkrMDEz',
-      'base64',
-    );
-
-    expect(() => Transaction.from(serializedVersionedTx)).to.throw(
-      'Versioned messages must be deserialized with VersionedMessage.deserialize()',
-    );
-
-    const versionedTx = VersionedTransaction.deserialize(serializedVersionedTx);
-    expect(versionedTx.message.version).to.eq(0);
-  });
-
   it('can serialize, deserialize, and reserialize with a partial signer', () => {
     const signer = Keypair.generate();
     const acc0Writable = Keypair.generate();
@@ -963,3 +944,97 @@ describe('Transaction', () => {
     t1.serialize();
   });
 });
+
+describe('VersionedTransaction', () => {
+  it('deserializes versioned transactions', () => {
+    const serializedVersionedTx = Buffer.from(
+      'AdTIDASR42TgVuXKkd7mJKk373J3LPVp85eyKMVcrboo9KTY8/vm6N/Cv0NiHqk2I8iYw6VX5ZaBKG8z' +
+        '9l1XjwiAAQACA+6qNbqfjaIENwt9GzEK/ENiB/ijGwluzBUmQ9xlTAMcCaS0ctnyxTcXXlJr7u2qtnaM' +
+        'gIAO2/c7RBD0ipHWUcEDBkZv5SEXMv/srbpyw5vnvIzlu8X3EmssQ5s6QAAAAJbI7VNs6MzREUlnzRaJ' +
+        'pBKP8QQoDn2dWQvD0KIgHFDiAwIACQAgoQcAAAAAAAIABQEAAAQAATYPBwAKBDIBAyQWIw0oCxIdCA4i' +
+        'JzQRKwUZHxceHCohMBUJJiwpMxAaGC0TLhQxGyAMBiU2NS8VDgAAAADuAgAAAAAAAAIAAAAAAAAAAdGCT' +
+        'Qiq5yw3+3m1sPoRNj0GtUNNs0FIMocxzt3zuoSZHQABAwQFBwgLDA8RFBcYGhwdHh8iIyUnKiwtLi8yF' +
+        'wIGCQoNDhASExUWGRsgISQmKCkrMDEz',
+      'base64',
+    );
+
+    expect(() => Transaction.from(serializedVersionedTx)).to.throw(
+      'Versioned messages must be deserialized with VersionedMessage.deserialize()',
+    );
+
+    const versionedTx = VersionedTransaction.deserialize(serializedVersionedTx);
+    expect(versionedTx.message.version).to.eq(0);
+  });
+
+  describe('addSignature', () => {
+    const signer1 = Keypair.generate();
+    const signer2 = Keypair.generate();
+    const signer3 = Keypair.generate();
+
+    const recentBlockhash = new PublicKey(3).toBuffer();
+
+    const message = new TransactionMessage({
+      payerKey: signer1.publicKey,
+      instructions: [
+        new TransactionInstruction({
+          data: Buffer.from('Hello!'),
+          keys: [
+            {
+              pubkey: signer1.publicKey,
+              isSigner: true,
+              isWritable: true,
+            },
+            {
+              pubkey: signer2.publicKey,
+              isSigner: true,
+              isWritable: true,
+            },
+            {
+              pubkey: signer3.publicKey,
+              isSigner: false,
+              isWritable: false,
+            },
+          ],
+          programId: new PublicKey(
+            'MemoSq4gqABAXKb96qnH8TysNcWxMyWCqXgDLGmfcHr',
+          ),
+        }),
+      ],
+      recentBlockhash: bs58.encode(recentBlockhash),
+    });
+
+    const transaction = new VersionedTransaction(message.compileToV0Message());
+
+    it('appends externally generated signatures at correct indexes', () => {
+      const signature1 = sign(
+        transaction.message.serialize(),
+        signer1.secretKey,
+      );
+      const signature2 = sign(
+        transaction.message.serialize(),
+        signer2.secretKey,
+      );
+
+      transaction.addSignature(signer2.publicKey, signature2);
+      transaction.addSignature(signer1.publicKey, signature1);
+
+      expect(transaction.signatures).to.have.length(2);
+      expect(transaction.signatures[0]).to.eq(signature1);
+      expect(transaction.signatures[1]).to.eq(signature2);
+    });
+
+    it('fatals when the signature is the wrong length', () => {
+      expect(() => {
+        transaction.addSignature(signer1.publicKey, new Uint8Array(32));
+      }).to.throw('Signature must be 64 bytes long');
+    });
+
+    it('fatals when adding a signature for a public key that has not been marked as a signer', () => {
+      expect(() => {
+        transaction.addSignature(signer3.publicKey, new Uint8Array(64));
+      }).to.throw(
+        `Can not add signature; \`${signer3.publicKey.toBase58()}\` is not required to sign this transaction`,
+      );
+    });
+  });
+});