From 7046136320e1ca63c3c18df71216fbbf5d4dee27 Mon Sep 17 00:00:00 2001
From: Kyrylo Stepanov <k.stepanov@tutanota.com>
Date: Tue, 5 Mar 2024 15:56:25 +0200
Subject: [PATCH] fixed behaviour

---
 core/rust/utils/src/account.rs | 28 ++++++++++++++++++----------
 1 file changed, 18 insertions(+), 10 deletions(-)

diff --git a/core/rust/utils/src/account.rs b/core/rust/utils/src/account.rs
index 2f221177d..b4b5016ce 100644
--- a/core/rust/utils/src/account.rs
+++ b/core/rust/utils/src/account.rs
@@ -3,6 +3,7 @@ use solana_program::{
     entrypoint::ProgramResult,
     msg,
     program::{invoke, invoke_signed},
+    program_error::ProgramError,
     pubkey::Pubkey,
     rent::Rent,
     system_instruction,
@@ -65,31 +66,38 @@ pub fn resize_or_reallocate_account_raw<'a>(
 ) -> ProgramResult {
     let rent = Rent::get()?;
     let new_minimum_balance = rent.minimum_balance(new_size);
-    let current_ta_lamports = target_account.lamports();
+    let lamports_diff =
+        new_minimum_balance.abs_diff(target_account.lamports());
+    if new_size == target_account.data_len() {
+        return Ok(());
+    }
+
     let account_infos = &[
         funding_account.clone(),
         target_account.clone(),
         system_program.clone(),
     ];
 
-    // account will be shrunk
-    if target_account.data_len() > new_size {
-        let lamports_diff = new_minimum_balance.saturating_sub(current_ta_lamports);
+    if new_size > target_account.data_len() {
         invoke(
             &system_instruction::transfer(funding_account.key, target_account.key, lamports_diff),
             account_infos,
         )?;
-    } else {
-        // account will be extended
-        let excess_lamports = current_ta_lamports.saturating_sub(new_minimum_balance);
+    } else if target_account.owner == system_program.owner {
         invoke(
-            &system_instruction::transfer(target_account.key, funding_account.key, excess_lamports),
+            &system_instruction::transfer(target_account.key, funding_account.key, lamports_diff),
             account_infos,
         )?;
+    } else {
+        (**target_account.try_borrow_mut_lamports()?)
+            .checked_sub(lamports_diff)
+            .ok_or(ProgramError::InvalidRealloc)?;
+        (**funding_account.try_borrow_mut_lamports()?)
+            .checked_add(lamports_diff)
+            .ok_or(ProgramError::InvalidRealloc)?;
     }
-    target_account.realloc(new_size, false)?;
 
-    Ok(())
+    target_account.realloc(new_size, false)
 }
 
 /// Close src_account and transfer lamports to dst_account, lifted from Solana Cookbook