Consume from AccountsDataMeter (#21994)

program-runtime/src/accounts_data_meter.rs

@@ -64,10 +64,21 @@ impl AccountsDataMeter {
             let amount = amount.abs() as u64;
             self.current = self.current.saturating_sub(amount);
         }
+
         Ok(())
     }
 }
 
+#[cfg(test)]
+impl AccountsDataMeter {
+    pub fn set_maximum(&mut self, maximum: u64) {
+        self.maximum = maximum;
+    }
+    pub fn set_current(&mut self, current: u64) {
+        self.current = current;
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

program-runtime/src/invoke_context.rs

@@ -9,8 +9,9 @@ use {
         bpf_loader_upgradeable::{self, UpgradeableLoaderState},
         compute_budget::ComputeBudget,
         feature_set::{
-            do_support_realloc, neon_evm_compute_budget, reject_empty_instruction_without_program,
-            remove_native_loader, requestable_heap_size, tx_wide_compute_cap, FeatureSet,
+            cap_accounts_data_len, do_support_realloc, neon_evm_compute_budget,
+            reject_empty_instruction_without_program, remove_native_loader, requestable_heap_size,
+            tx_wide_compute_cap, FeatureSet,
         },
         hash::Hash,
         instruction::{AccountMeta, CompiledInstruction, Instruction, InstructionError},
@@ -367,6 +368,7 @@ impl<'a> InvokeContext<'a> {
         program_indices: &[usize],
     ) -> Result<(), InstructionError> {
         let do_support_realloc = self.feature_set.is_active(&do_support_realloc::id());
+        let cap_accounts_data_len = self.feature_set.is_active(&cap_accounts_data_len::id());
         let program_id = self
             .transaction_context
             .get_program_key()
@@ -423,6 +425,14 @@ impl<'a> InvokeContext<'a> {
             post_sum = post_sum
                 .checked_add(u128::from(account.lamports()))
                 .ok_or(InstructionError::UnbalancedInstruction)?;
+
+            if cap_accounts_data_len {
+                let pre_data_len = pre_account.data().len() as i64;
+                let post_data_len = account.data().len() as i64;
+                let data_len_delta = post_data_len.saturating_sub(pre_data_len);
+                self.accounts_data_meter.consume(data_len_delta)?;
+            }
+
             Ok(())
         };
         visit_each_account_once(instruction_accounts, &mut work)?;
@@ -441,6 +451,7 @@ impl<'a> InvokeContext<'a> {
         before_instruction_context_push: bool,
     ) -> Result<(), InstructionError> {
         let do_support_realloc = self.feature_set.is_active(&do_support_realloc::id());
+        let cap_accounts_data_len = self.feature_set.is_active(&cap_accounts_data_len::id());
         let transaction_context = &self.transaction_context;
         let instruction_context = transaction_context.get_current_instruction_context()?;
         let program_id = transaction_context
@@ -505,6 +516,14 @@ impl<'a> InvokeContext<'a> {
                         if is_writable && !pre_account.executable() {
                             pre_account.update(account.clone());
                         }
+
+                        if cap_accounts_data_len {
+                            let pre_data_len = pre_account.data().len() as i64;
+                            let post_data_len = account.data().len() as i64;
+                            let data_len_delta = post_data_len.saturating_sub(pre_data_len);
+                            self.accounts_data_meter.consume(data_len_delta)?;
+                        }
+
                         return Ok(());
                     }
                 }
@@ -1062,6 +1081,9 @@ mod tests {
             compute_units_to_consume: u64,
             desired_result: Result<(), InstructionError>,
         },
+        Resize {
+            new_len: usize,
+        },
     }
 
     #[test]
@@ -1191,6 +1213,12 @@ mod tests {
                         .unwrap();
                     return desired_result;
                 }
+                MockInstruction::Resize { new_len } => {
+                    keyed_account_at_index(keyed_accounts, first_instruction_account)?
+                        .try_account_ref_mut()?
+                        .data_mut()
+                        .resize_with(new_len, Default::default)
+                }
             }
         } else {
             return Err(InstructionError::InvalidInstructionData);
@@ -1507,4 +1535,111 @@ mod tests {
         );
         invoke_context.pop().unwrap();
     }
+
+    #[test]
+    fn test_process_instruction_accounts_data_meter() {
+        solana_logger::setup();
+
+        let program_key = Pubkey::new_unique();
+        let user_account_data_len = 123;
+        let user_account = AccountSharedData::new(100, user_account_data_len, &program_key);
+        let dummy_account = AccountSharedData::new(10, 0, &program_key);
+        let mut program_account = AccountSharedData::new(500, 500, &native_loader::id());
+        program_account.set_executable(true);
+        let accounts = vec![
+            (Pubkey::new_unique(), user_account),
+            (Pubkey::new_unique(), dummy_account),
+            (program_key, program_account),
+        ];
+
+        let builtin_programs = [BuiltinProgram {
+            program_id: program_key,
+            process_instruction: mock_process_instruction,
+        }];
+
+        let mut transaction_context = TransactionContext::new(accounts, 1);
+        let mut invoke_context =
+            InvokeContext::new_mock(&mut transaction_context, &builtin_programs);
+
+        invoke_context
+            .accounts_data_meter
+            .set_current(user_account_data_len as u64);
+        invoke_context
+            .accounts_data_meter
+            .set_maximum(user_account_data_len as u64 * 3);
+        let remaining_account_data_len = invoke_context.accounts_data_meter.remaining() as usize;
+
+        let instruction_accounts = [
+            InstructionAccount {
+                index_in_transaction: 0,
+                index_in_caller: 1,
+                is_signer: false,
+                is_writable: true,
+            },
+            InstructionAccount {
+                index_in_transaction: 1,
+                index_in_caller: 2,
+                is_signer: false,
+                is_writable: false,
+            },
+        ];
+
+        // Test 1: Resize the account to use up all the space; this must succeed
+        {
+            let new_len = user_account_data_len + remaining_account_data_len;
+            dbg!(new_len);
+            let instruction_data =
+                bincode::serialize(&MockInstruction::Resize { new_len }).unwrap();
+
+            let result = invoke_context.process_instruction(
+                &instruction_data,
+                &instruction_accounts,
+                &[2],
+                &mut 0,
+            );
+
+            assert!(result.is_ok());
+            assert_eq!(invoke_context.accounts_data_meter.remaining(), 0);
+        }
+
+        // Test 2: Resize the account to *the same size*, so not consuming any additional size; this must succeed
+        {
+            let new_len = user_account_data_len + remaining_account_data_len;
+            dbg!(new_len);
+            let instruction_data =
+                bincode::serialize(&MockInstruction::Resize { new_len }).unwrap();
+
+            let result = invoke_context.process_instruction(
+                &instruction_data,
+                &instruction_accounts,
+                &[2],
+                &mut 0,
+            );
+
+            assert!(result.is_ok());
+            assert_eq!(invoke_context.accounts_data_meter.remaining(), 0);
+        }
+
+        // Test 3: Resize the account to exceed the budget; this must fail
+        {
+            let new_len = user_account_data_len + remaining_account_data_len + 1;
+            dbg!(new_len);
+            let instruction_data =
+                bincode::serialize(&MockInstruction::Resize { new_len }).unwrap();
+
+            let result = invoke_context.process_instruction(
+                &instruction_data,
+                &instruction_accounts,
+                &[2],
+                &mut 0,
+            );
+
+            assert!(result.is_err());
+            assert!(matches!(
+                result,
+                Err(solana_sdk::instruction::InstructionError::AccountsDataBudgetExceeded)
+            ));
+            assert_eq!(invoke_context.accounts_data_meter.remaining(), 0);
+        }
+    }
 }

runtime/src/bank.rs

@@ -15190,4 +15190,48 @@ pub(crate) mod tests {
             Some(Vec::<TransactionLogInfo>::new()),
         );
     }
+
+    /// Test exceeding the accounts data budget by creating accounts in a loop
+    #[test]
+    fn test_accounts_data_budget_exceeded() {
+        use solana_program_runtime::accounts_data_meter::MAX_ACCOUNTS_DATA_LEN;
+        use solana_sdk::system_instruction::MAX_PERMITTED_DATA_LENGTH;
+
+        solana_logger::setup();
+        let (genesis_config, mint_keypair) = create_genesis_config(1_000_000_000_000);
+        let mut bank = Bank::new_for_tests(&genesis_config);
+        bank.activate_feature(&solana_sdk::feature_set::cap_accounts_data_len::id());
+
+        let mut i = 0;
+        let result = loop {
+            let txn = system_transaction::create_account(
+                &mint_keypair,
+                &Keypair::new(),
+                bank.last_blockhash(),
+                1,
+                MAX_PERMITTED_DATA_LENGTH,
+                &solana_sdk::system_program::id(),
+            );
+
+            let result = bank.process_transaction(&txn);
+            assert!(bank.load_accounts_data_len() <= MAX_ACCOUNTS_DATA_LEN);
+            if result.is_err() {
+                break result;
+            }
+
+            assert!(
+                i < MAX_ACCOUNTS_DATA_LEN / MAX_PERMITTED_DATA_LENGTH,
+                "test must complete within bounded limits"
+            );
+            i += 1;
+        };
+
+        assert!(matches!(
+            result,
+            Err(TransactionError::InstructionError(
+                _,
+                solana_sdk::instruction::InstructionError::AccountsDataBudgetExceeded,
+            ))
+        ));
+    }
 }