From 7ad782febf8215ec30e435a4957615bdc99e56e3 Mon Sep 17 00:00:00 2001
From: Andrew Schonfeld <andrew.schonfeld1@gmail.com>
Date: Tue, 16 Aug 2022 18:22:38 -0400
Subject: [PATCH] VoteProgram.safeWithdraw function to safeguard against
 accidental vote account closures (#26586)

feat: safe withdraw function

Co-authored-by: aschonfeld <andrew@proofofalpha.io>
---
 web3.js/src/programs/vote.ts            | 21 +++++++++++++++++++++
 web3.js/test/program-tests/vote.test.ts | 15 +++++++++++++++
 2 files changed, 36 insertions(+)

diff --git a/web3.js/src/programs/vote.ts b/web3.js/src/programs/vote.ts
index 6cd16671a977cd..db1a111da919a4 100644
--- a/web3.js/src/programs/vote.ts
+++ b/web3.js/src/programs/vote.ts
@@ -410,4 +410,25 @@ export class VoteProgram {
       data,
     });
   }
+
+  /**
+   * Generate a transaction to withdraw safely from a Vote account.
+   *
+   * This function was created as a safeguard for vote accounts running validators, `safeWithdraw`
+   * checks that the withdraw amount will not exceed the specified balance while leaving enough left
+   * to cover rent. If you wish to close the vote account by withdrawing the full amount, call the
+   * `withdraw` method directly.
+   */
+  static safeWithdraw(
+    params: WithdrawFromVoteAccountParams,
+    currentVoteAccountBalance: number,
+    rentExemptMinimum: number,
+  ): Transaction {
+    if (params.lamports > currentVoteAccountBalance - rentExemptMinimum) {
+      throw new Error(
+        'Withdraw will leave vote account with insuffcient funds.',
+      );
+    }
+    return VoteProgram.withdraw(params);
+  }
 }
diff --git a/web3.js/test/program-tests/vote.test.ts b/web3.js/test/program-tests/vote.test.ts
index 596e6e401b4fb7..6cd349a0c3a772 100644
--- a/web3.js/test/program-tests/vote.test.ts
+++ b/web3.js/test/program-tests/vote.test.ts
@@ -167,6 +167,21 @@ describe('VoteProgram', () => {
 
       // Withdraw from Vote account
       let recipient = Keypair.generate();
+      const voteBalance = await connection.getBalance(newVoteAccount.publicKey);
+
+      expect(() =>
+        VoteProgram.safeWithdraw(
+          {
+            votePubkey: newVoteAccount.publicKey,
+            authorizedWithdrawerPubkey: authorized.publicKey,
+            lamports: voteBalance - minimumAmount + 1,
+            toPubkey: recipient.publicKey,
+          },
+          voteBalance,
+          minimumAmount,
+        ),
+      ).to.throw('Withdraw will leave vote account with insuffcient funds.');
+
       let withdraw = VoteProgram.withdraw({
         votePubkey: newVoteAccount.publicKey,
         authorizedWithdrawerPubkey: authorized.publicKey,