Creating a more robust rent management scheme (#377)

* Adding debug prints and moving from asserts to checked_adds to manage rent exemption. * Changing to new deposit workflow. * Removing bad message. * Fixing deposit test. * Removing Debug prints. * Saving files is good.
metaplex-foundation · Apr 21, 2022 · d370efb · d370efb
1 parent abb255a
commit d370efb
Show file tree

Hide file tree

Showing 4 changed files with 40 additions and 25 deletions.
diff --git a/auction-house/program/src/bid/mod.rs b/auction-house/program/src/bid/mod.rs
@@ -210,12 +210,17 @@ pub fn bid_logic<'info>(
     if is_native {
         assert_keys_equal(wallet.key(), payment_account.key())?;
 
-        if escrow_payment_account.lamports() < buyer_price.checked_add(rent.minimum_balance(escrow_payment_account.data_len())).ok_or(ErrorCode::NumericalOverflow)? {
+        if escrow_payment_account.lamports()
+            < buyer_price
+                .checked_add(rent.minimum_balance(escrow_payment_account.data_len()))
+                .ok_or(ErrorCode::NumericalOverflow)?
+        {
             let diff = buyer_price
                 .checked_add(rent.minimum_balance(escrow_payment_account.data_len()))
                 .ok_or(ErrorCode::NumericalOverflow)?
                 .checked_sub(escrow_payment_account.lamports())
                 .ok_or(ErrorCode::NumericalOverflow)?;
+
             invoke(
                 &system_instruction::transfer(
                     &payment_account.key(),

diff --git a/auction-house/program/src/lib.rs b/auction-house/program/src/lib.rs
@@ -396,11 +396,13 @@ pub mod auction_house {
             )?;
         } else {
             assert_keys_equal(receipt_account.key(), wallet.key())?;
+            let checked_amount =
+                rent_checked_sub(escrow_payment_account.to_account_info(), amount)?;
             invoke_signed(
                 &system_instruction::transfer(
                     &escrow_payment_account.key(),
                     &receipt_account.key(),
-                    amount,
+                    checked_amount,
                 ),
                 &[
                     escrow_payment_account.to_account_info(),
@@ -409,9 +411,6 @@ pub mod auction_house {
                 ],
                 &[&escrow_signer_seeds],
             )?;
-
-            // Verify that escrow is still rent exempt.
-            assert_escrow_rent_exempt(escrow_payment_account.to_account_info())?;
         }
 
         Ok(())
@@ -494,21 +493,22 @@ pub mod auction_house {
             )?;
         } else {
             assert_keys_equal(payment_account.key(), wallet.key())?;
+            // Reach rental exemption and then
+            let checked_amount = rent_checked_add(escrow_payment_account.to_account_info(), 0)?
+                .checked_add(amount)
+                .ok_or(ErrorCode::NumericalOverflow)?;
             invoke(
                 &system_instruction::transfer(
                     &payment_account.key(),
                     &escrow_payment_account.key(),
-                    amount,
+                    checked_amount,
                 ),
                 &[
                     escrow_payment_account.to_account_info(),
                     payment_account.to_account_info(),
                     system_program.to_account_info(),
                 ],
             )?;
-
-            // Verify enough exists in the escrow account to keep it rent exempt.
-            assert_escrow_rent_exempt(escrow_payment_account.to_account_info())?;
         }
 
         Ok(())
@@ -817,10 +817,6 @@ pub mod auction_house {
                 ],
                 &[&escrow_signer_seeds],
             )?;
-
-            // The buy instruction should handle accounting for minimum rent exemption for the
-            // escrow account, but double check here just in case.
-            assert_escrow_rent_exempt(escrow_payment_account.to_account_info())?;
         }
 
         if buyer_receipt_token_account.data_is_empty() {

diff --git a/auction-house/program/src/utils.rs b/auction-house/program/src/utils.rs
@@ -550,24 +550,38 @@ pub fn assert_valid_trade_state<'a>(
             &token_size_bytes,
         ],
     );
-    msg!(
-        "{:?}, {:?}, {:?}",
-        canonical_public_bump,
-        canonical_bump,
-        ts_bump
-    );
+
     match (canonical_public_bump, canonical_bump) {
         (Ok(public), Err(_)) if public == ts_bump => Ok(public),
         (Err(_), Ok(bump)) if bump == ts_bump => Ok(bump),
         _ => Err(ErrorCode::DerivedKeyInvalid.into()),
     }
 }
 
-pub fn assert_escrow_rent_exempt(escrow_account: AccountInfo) -> Result<()>{
-    if escrow_account.lamports() < (Rent::get()?).minimum_balance(escrow_account.data_len()) {
-        return err!(ErrorCode::EscrowUnderRentExemption);
+pub fn rent_checked_sub(escrow_account: AccountInfo, diff: u64) -> Result<u64> {
+    let rent_minimum: u64 = (Rent::get()?).minimum_balance(escrow_account.data_len());
+    let account_lamports: u64 = escrow_account
+        .lamports()
+        .checked_sub(diff)
+        .ok_or(ErrorCode::NumericalOverflow)?;
+
+    if account_lamports < rent_minimum {
+        Ok(escrow_account.lamports() - rent_minimum)
+    } else {
+        Ok(diff)
     }
-    else {
-        Ok(())
+}
+
+pub fn rent_checked_add(escrow_account: AccountInfo, diff: u64) -> Result<u64> {
+    let rent_minimum: u64 = (Rent::get()?).minimum_balance(escrow_account.data_len());
+    let account_lamports: u64 = escrow_account
+        .lamports()
+        .checked_add(diff)
+        .ok_or(ErrorCode::NumericalOverflow)?;
+
+    if account_lamports < rent_minimum {
+        Ok(rent_minimum - account_lamports)
+    } else {
+        Ok(diff)
     }
 }
diff --git a/auction-house/program/tests/deposit.rs b/auction-house/program/tests/deposit.rs
@@ -52,5 +52,5 @@ async fn deposit_success() {
         .await
         .expect("Error Getting Escrow")
         .expect("Trade State Escrow");
-    assert_eq!(escrow.lamports, 1000000000);
+    assert_eq!(escrow.lamports, 1000890880);
 }