From 2e449276bed905e2d235c5907734918a4775b28c Mon Sep 17 00:00:00 2001
From: "mergify[bot]" <37929162+mergify[bot]@users.noreply.github.com>
Date: Fri, 16 Oct 2020 19:14:26 +0000
Subject: [PATCH] Check payer balance for program account rent as needed
 (#12952) (#12955)

(cherry picked from commit b6bfed64cb159ee67bb6bdbaefc7f833bbed3563)

Co-authored-by: Tyera Eulberg <teulberg@gmail.com>
---
 cli/src/checks.rs | 36 +++++++++++++++++++++++++++++++++---
 cli/src/cli.rs    | 27 +++++++++++++++++----------
 2 files changed, 50 insertions(+), 13 deletions(-)

diff --git a/cli/src/checks.rs b/cli/src/checks.rs
index f1c6514878..7a2d01ae03 100644
--- a/cli/src/checks.rs
+++ b/cli/src/checks.rs
@@ -54,12 +54,42 @@ pub fn check_account_for_multiple_fees_with_commitment(
     fee_calculator: &FeeCalculator,
     messages: &[&Message],
     commitment: CommitmentConfig,
+) -> Result<(), CliError> {
+    check_account_for_spend_multiple_fees_with_commitment(
+        rpc_client,
+        account_pubkey,
+        0,
+        fee_calculator,
+        messages,
+        commitment,
+    )
+}
+
+pub fn check_account_for_spend_multiple_fees_with_commitment(
+    rpc_client: &RpcClient,
+    account_pubkey: &Pubkey,
+    balance: u64,
+    fee_calculator: &FeeCalculator,
+    messages: &[&Message],
+    commitment: CommitmentConfig,
 ) -> Result<(), CliError> {
     let fee = calculate_fee(fee_calculator, messages);
-    if !check_account_for_balance_with_commitment(rpc_client, account_pubkey, fee, commitment)
-        .map_err(Into::<ClientError>::into)?
+    if !check_account_for_balance_with_commitment(
+        rpc_client,
+        account_pubkey,
+        balance + fee,
+        commitment,
+    )
+    .map_err(Into::<ClientError>::into)?
     {
-        return Err(CliError::InsufficientFundsForFee(lamports_to_sol(fee)));
+        if balance > 0 {
+            return Err(CliError::InsufficientFundsForSpendAndFee(
+                lamports_to_sol(balance),
+                lamports_to_sol(fee),
+            ));
+        } else {
+            return Err(CliError::InsufficientFundsForFee(lamports_to_sol(fee)));
+        }
     }
     Ok(())
 }
diff --git a/cli/src/cli.rs b/cli/src/cli.rs
index b0b80ac15d..b4f84563f8 100644
--- a/cli/src/cli.rs
+++ b/cli/src/cli.rs
@@ -1166,11 +1166,12 @@ fn process_deploy(
     let signers = [config.signers[0], program_id];
 
     // Check program account to see if partial initialization has occurred
-    let initial_instructions = if let Some(account) = rpc_client
+    let (initial_instructions, balance_needed) = if let Some(account) = rpc_client
         .get_account_with_commitment(&program_id.pubkey(), config.commitment)?
         .value
     {
         let mut instructions: Vec<Instruction> = vec![];
+        let mut balance_needed = 0;
         if account.executable {
             return Err(CliError::DynamicProgramError(
                 "Program account is already executable".to_string(),
@@ -1194,21 +1195,26 @@ fn process_deploy(
             }
         }
         if account.lamports < minimum_balance {
+            let balance = minimum_balance - account.lamports;
             instructions.push(system_instruction::transfer(
                 &config.signers[0].pubkey(),
                 &program_id.pubkey(),
-                minimum_balance - account.lamports,
+                balance,
             ));
+            balance_needed = balance;
         }
-        instructions
+        (instructions, balance_needed)
     } else {
-        vec![system_instruction::create_account(
-            &config.signers[0].pubkey(),
-            &program_id.pubkey(),
+        (
+            vec![system_instruction::create_account(
+                &config.signers[0].pubkey(),
+                &program_id.pubkey(),
+                minimum_balance,
+                program_data.len() as u64,
+                &loader_id,
+            )],
             minimum_balance,
-            program_data.len() as u64,
-            &loader_id,
-        )]
+        )
     };
     let initial_message = if !initial_instructions.is_empty() {
         Some(Message::new(
@@ -1251,9 +1257,10 @@ fn process_deploy(
         .get_recent_blockhash_with_commitment(config.commitment)?
         .value;
 
-    check_account_for_multiple_fees_with_commitment(
+    check_account_for_spend_multiple_fees_with_commitment(
         rpc_client,
         &config.signers[0].pubkey(),
+        balance_needed,
         &fee_calculator,
         &messages,
         config.commitment,