Fix CPI recursion depth (#17659)

* Fix CPI recursion depth

programs/bpf/c/src/invoke/invoke.c

@@ -18,6 +18,7 @@ static const uint8_t TEST_RETURN_ERROR = 11;
 static const uint8_t TEST_PRIVILEGE_DEESCALATION_ESCALATION_SIGNER = 12;
 static const uint8_t TEST_PRIVILEGE_DEESCALATION_ESCALATION_WRITABLE = 13;
 static const uint8_t TEST_WRITABLE_DEESCALATION_WRITABLE = 14;
+static const uint8_t TEST_NESTED_INVOKE_TOO_DEEP = 15;
 
 static const int MINT_INDEX = 0;
 static const int ARGUMENT_INDEX = 1;
@@ -31,6 +32,35 @@ static const int DERIVED_KEY3_INDEX = 8;
 static const int SYSTEM_PROGRAM_INDEX = 9;
 static const int FROM_INDEX = 10;
 
+uint64_t do_nested_invokes(uint64_t num_nested_invokes,
+                           SolAccountInfo *accounts, uint64_t num_accounts) {
+  sol_assert(accounts[ARGUMENT_INDEX].is_signer);
+
+  *accounts[ARGUMENT_INDEX].lamports -= 5;
+  *accounts[INVOKED_ARGUMENT_INDEX].lamports += 5;
+
+  SolAccountMeta arguments[] = {
+      {accounts[INVOKED_ARGUMENT_INDEX].key, true, true},
+      {accounts[ARGUMENT_INDEX].key, true, true},
+      {accounts[INVOKED_PROGRAM_INDEX].key, false, false}};
+  uint8_t data[] = {NESTED_INVOKE, num_nested_invokes};
+  const SolInstruction instruction = {accounts[INVOKED_PROGRAM_INDEX].key,
+                                      arguments, SOL_ARRAY_SIZE(arguments),
+                                      data, SOL_ARRAY_SIZE(data)};
+
+  sol_log("First invoke");
+  sol_assert(SUCCESS == sol_invoke(&instruction, accounts, num_accounts));
+  sol_log("2nd invoke from first program");
+  sol_assert(SUCCESS == sol_invoke(&instruction, accounts, num_accounts));
+
+  sol_assert(*accounts[ARGUMENT_INDEX].lamports ==
+             42 - 5 + (2 * num_nested_invokes));
+  sol_assert(*accounts[INVOKED_ARGUMENT_INDEX].lamports ==
+             10 + 5 - (2 * num_nested_invokes));
+
+  return SUCCESS;
+}
+
 extern uint64_t entrypoint(const uint8_t *input) {
   sol_log("Invoke C program");
 
@@ -203,32 +233,9 @@ extern uint64_t entrypoint(const uint8_t *input) {
                  sol_invoke(&instruction, accounts, SOL_ARRAY_SIZE(accounts)));
     }
 
-    sol_log("Test invoke");
+    sol_log("Test nested invoke");
     {
-      sol_assert(accounts[ARGUMENT_INDEX].is_signer);
+      sol_assert(SUCCESS == do_nested_invokes(4, accounts, params.ka_num));
-
-      *accounts[ARGUMENT_INDEX].lamports -= 5;
-      *accounts[INVOKED_ARGUMENT_INDEX].lamports += 5;
-
-      SolAccountMeta arguments[] = {
-          {accounts[INVOKED_ARGUMENT_INDEX].key, true, true},
-          {accounts[ARGUMENT_INDEX].key, true, true},
-          {accounts[INVOKED_PROGRAM_DUP_INDEX].key, false, false}};
-      uint8_t data[] = {NESTED_INVOKE};
-      const SolInstruction instruction = {accounts[INVOKED_PROGRAM_INDEX].key,
-                                          arguments, SOL_ARRAY_SIZE(arguments),
-                                          data, SOL_ARRAY_SIZE(data)};
-
-      sol_log("First invoke");
-      sol_assert(SUCCESS ==
-                 sol_invoke(&instruction, accounts, SOL_ARRAY_SIZE(accounts)));
-      sol_log("2nd invoke from first program");
-      sol_assert(SUCCESS ==
-                 sol_invoke(&instruction, accounts, SOL_ARRAY_SIZE(accounts)));
-
-      sol_assert(*accounts[ARGUMENT_INDEX].lamports == 42 - 5 + 1 + 1 + 1 + 1);
-      sol_assert(*accounts[INVOKED_ARGUMENT_INDEX].lamports ==
-                 10 + 5 - 1 - 1 - 1 - 1);
     }
 
     sol_log("Test privilege deescalation");
@@ -505,24 +512,29 @@ extern uint64_t entrypoint(const uint8_t *input) {
     break;
   }
   case TEST_WRITABLE_DEESCALATION_WRITABLE: {
-  sol_log("Test writable deescalation");
+    sol_log("Test writable deescalation");
-      uint8_t buffer[10];
+    uint8_t buffer[10];
-      for (int i = 0; i < 10; i++) {
+    for (int i = 0; i < 10; i++) {
-        buffer[i] = accounts[INVOKED_ARGUMENT_INDEX].data[i];
+      buffer[i] = accounts[INVOKED_ARGUMENT_INDEX].data[i];
-      }
-      SolAccountMeta arguments[] = {
-          {accounts[INVOKED_ARGUMENT_INDEX].key, false, false}};
-      uint8_t data[] = {WRITE_ACCOUNT, 10};
-      const SolInstruction instruction = {accounts[INVOKED_PROGRAM_INDEX].key,
-                                          arguments, SOL_ARRAY_SIZE(arguments),
-                                          data, SOL_ARRAY_SIZE(data)};
-      sol_invoke(&instruction, accounts, SOL_ARRAY_SIZE(accounts));
-
-      for (int i = 0; i < 10; i++) {
-        sol_assert(buffer[i] == accounts[INVOKED_ARGUMENT_INDEX].data[i]);
-      }
-      break;
     }
+    SolAccountMeta arguments[] = {
+        {accounts[INVOKED_ARGUMENT_INDEX].key, false, false}};
+    uint8_t data[] = {WRITE_ACCOUNT, 10};
+    const SolInstruction instruction = {accounts[INVOKED_PROGRAM_INDEX].key,
+                                        arguments, SOL_ARRAY_SIZE(arguments),
+                                        data, SOL_ARRAY_SIZE(data)};
+    sol_invoke(&instruction, accounts, SOL_ARRAY_SIZE(accounts));
+
+    for (int i = 0; i < 10; i++) {
+      sol_assert(buffer[i] == accounts[INVOKED_ARGUMENT_INDEX].data[i]);
+    }
+    break;
+  }
+  case TEST_NESTED_INVOKE_TOO_DEEP: {
+    do_nested_invokes(5, accounts, params.ka_num);
+    break;
+  }
+
   default:
     sol_panic();
   }

programs/bpf/c/src/invoked/invoked.c

@@ -228,16 +228,17 @@ extern uint64_t entrypoint(const uint8_t *input) {
     *accounts[INVOKED_ARGUMENT_INDEX].lamports -= 1;
     *accounts[ARGUMENT_INDEX].lamports += 1;
 
-    if (params.ka_num == 3) {
+    uint8_t remaining_invokes = params.data[1];
+    if (remaining_invokes > 1) {
+      sol_log("Invoke again");
       SolAccountMeta arguments[] = {
           {accounts[INVOKED_ARGUMENT_INDEX].key, true, true},
-          {accounts[ARGUMENT_INDEX].key, true, true}};
+          {accounts[ARGUMENT_INDEX].key, true, true},
-      uint8_t data[] = {NESTED_INVOKE};
+          {accounts[INVOKED_PROGRAM_INDEX].key, false, false}};
+      uint8_t data[] = {NESTED_INVOKE, remaining_invokes - 1};
       const SolInstruction instruction = {accounts[INVOKED_PROGRAM_INDEX].key,
                                           arguments, SOL_ARRAY_SIZE(arguments),
                                           data, SOL_ARRAY_SIZE(data)};
-
-      sol_log("Invoke again");
       sol_assert(SUCCESS == sol_invoke(&instruction, accounts, params.ka_num));
     } else {
       sol_log("Last invoked");

programs/bpf/rust/invoke/src/lib.rs

@@ -30,19 +30,53 @@ const TEST_RETURN_ERROR: u8 = 11;
 const TEST_PRIVILEGE_DEESCALATION_ESCALATION_SIGNER: u8 = 12;
 const TEST_PRIVILEGE_DEESCALATION_ESCALATION_WRITABLE: u8 = 13;
 const TEST_WRITABLE_DEESCALATION_WRITABLE: u8 = 14;
+const TEST_NESTED_INVOKE_TOO_DEEP: u8 = 15;
 
-// const MINT_INDEX: usize = 0;
+// const MINT_INDEX: usize = 0; // unused placeholder
 const ARGUMENT_INDEX: usize = 1;
 const INVOKED_PROGRAM_INDEX: usize = 2;
 const INVOKED_ARGUMENT_INDEX: usize = 3;
 const INVOKED_PROGRAM_DUP_INDEX: usize = 4;
-// const ARGUMENT_DUP_INDEX: usize = 5;
+// const ARGUMENT_DUP_INDEX: usize = 5; unused placeholder
 const DERIVED_KEY1_INDEX: usize = 6;
 const DERIVED_KEY2_INDEX: usize = 7;
 const DERIVED_KEY3_INDEX: usize = 8;
 const SYSTEM_PROGRAM_INDEX: usize = 9;
 const FROM_INDEX: usize = 10;
 
+fn do_nested_invokes(num_nested_invokes: u64, accounts: &[AccountInfo]) -> ProgramResult {
+    assert!(accounts[ARGUMENT_INDEX].is_signer);
+
+    let pre_argument_lamports = accounts[ARGUMENT_INDEX].lamports();
+    let pre_invoke_argument_lamports = accounts[INVOKED_ARGUMENT_INDEX].lamports();
+    **accounts[ARGUMENT_INDEX].lamports.borrow_mut() -= 5;
+    **accounts[INVOKED_ARGUMENT_INDEX].lamports.borrow_mut() += 5;
+
+    msg!("First invoke");
+    let instruction = create_instruction(
+        *accounts[INVOKED_PROGRAM_INDEX].key,
+        &[
+            (accounts[ARGUMENT_INDEX].key, true, true),
+            (accounts[INVOKED_ARGUMENT_INDEX].key, true, true),
+            (accounts[INVOKED_PROGRAM_INDEX].key, false, false),
+        ],
+        vec![NESTED_INVOKE, num_nested_invokes as u8],
+    );
+    invoke(&instruction, accounts)?;
+    msg!("2nd invoke from first program");
+    invoke(&instruction, accounts)?;
+
+    assert_eq!(
+        accounts[ARGUMENT_INDEX].lamports(),
+        pre_argument_lamports - 5 + (2 * num_nested_invokes)
+    );
+    assert_eq!(
+        accounts[INVOKED_ARGUMENT_INDEX].lamports(),
+        pre_invoke_argument_lamports + 5 - (2 * num_nested_invokes)
+    );
+    Ok(())
+}
+
 entrypoint!(process_instruction);
 fn process_instruction(
     program_id: &Pubkey,
@@ -282,31 +316,7 @@ fn process_instruction(
 
             msg!("Test nested invoke");
             {
-                assert!(accounts[ARGUMENT_INDEX].is_signer);
+                do_nested_invokes(4, accounts)?;
-
-                **accounts[ARGUMENT_INDEX].lamports.borrow_mut() -= 5;
-                **accounts[INVOKED_ARGUMENT_INDEX].lamports.borrow_mut() += 5;
-
-                msg!("First invoke");
-                let instruction = create_instruction(
-                    *accounts[INVOKED_PROGRAM_INDEX].key,
-                    &[
-                        (accounts[ARGUMENT_INDEX].key, true, true),
-                        (accounts[INVOKED_ARGUMENT_INDEX].key, true, true),
-                        (accounts[INVOKED_PROGRAM_DUP_INDEX].key, false, false),
-                        (accounts[INVOKED_PROGRAM_DUP_INDEX].key, false, false),
-                    ],
-                    vec![NESTED_INVOKE],
-                );
-                invoke(&instruction, accounts)?;
-                msg!("2nd invoke from first program");
-                invoke(&instruction, accounts)?;
-
-                assert_eq!(accounts[ARGUMENT_INDEX].lamports(), 42 - 5 + 1 + 1 + 1 + 1);
-                assert_eq!(
-                    accounts[INVOKED_ARGUMENT_INDEX].lamports(),
-                    10 + 5 - 1 - 1 - 1 - 1
-                );
             }
 
             msg!("Test privilege deescalation");
@@ -602,6 +612,9 @@ fn process_instruction(
                 accounts[INVOKED_ARGUMENT_INDEX].data.borrow_mut()[..NUM_BYTES]
             );
         }
+        TEST_NESTED_INVOKE_TOO_DEEP => {
+            let _ = do_nested_invokes(5, accounts);
+        }
         _ => panic!(),
     }
 

programs/bpf/rust/invoked/src/processor.rs

@@ -202,21 +202,24 @@ fn process_instruction(
             msg!("nested invoke");
             const ARGUMENT_INDEX: usize = 0;
             const INVOKED_ARGUMENT_INDEX: usize = 1;
-            const INVOKED_PROGRAM_INDEX: usize = 3;
+            const INVOKED_PROGRAM_INDEX: usize = 2;
 
             assert!(accounts[INVOKED_ARGUMENT_INDEX].is_signer);
+            assert!(instruction_data.len() > 1);
 
             **accounts[INVOKED_ARGUMENT_INDEX].lamports.borrow_mut() -= 1;
             **accounts[ARGUMENT_INDEX].lamports.borrow_mut() += 1;
-            if accounts.len() > 2 {
+            let remaining_invokes = instruction_data[1];
+            if remaining_invokes > 1 {
                 msg!("Invoke again");
                 let invoked_instruction = create_instruction(
                     *accounts[INVOKED_PROGRAM_INDEX].key,
                     &[
                         (accounts[ARGUMENT_INDEX].key, true, true),
                         (accounts[INVOKED_ARGUMENT_INDEX].key, true, true),
+                        (accounts[INVOKED_PROGRAM_INDEX].key, false, false),
                     ],
-                    vec![NESTED_INVOKE],
+                    vec![NESTED_INVOKE, remaining_invokes - 1],
                 );
                 invoke(&invoked_instruction, accounts)?;
             } else {

programs/bpf/tests/programs.rs

@@ -757,6 +757,7 @@ fn test_program_bpf_invoke_sanity() {
     const TEST_PRIVILEGE_DEESCALATION_ESCALATION_SIGNER: u8 = 12;
     const TEST_PRIVILEGE_DEESCALATION_ESCALATION_WRITABLE: u8 = 13;
     const TEST_WRITABLE_DEESCALATION_WRITABLE: u8 = 14;
+    const TEST_NESTED_INVOKE_TOO_DEEP: u8 = 15;
 
     #[allow(dead_code)]
     #[derive(Debug)]
@@ -874,6 +875,10 @@ fn test_program_bpf_invoke_sanity() {
                 invoked_program_id.clone(),
                 invoked_program_id.clone(),
                 invoked_program_id.clone(),
+                invoked_program_id.clone(),
+                invoked_program_id.clone(),
+                invoked_program_id.clone(),
+                invoked_program_id.clone(),
             ],
             Languages::Rust => vec![
                 solana_sdk::system_program::id(),
@@ -893,6 +898,10 @@ fn test_program_bpf_invoke_sanity() {
                 invoked_program_id.clone(),
                 invoked_program_id.clone(),
                 invoked_program_id.clone(),
+                invoked_program_id.clone(),
+                invoked_program_id.clone(),
+                invoked_program_id.clone(),
+                invoked_program_id.clone(),
                 solana_sdk::system_program::id(),
             ],
         };
@@ -1004,6 +1013,18 @@ fn test_program_bpf_invoke_sanity() {
             &[invoked_program_id.clone()],
         );
 
+        do_invoke_failure_test_local(
+            TEST_NESTED_INVOKE_TOO_DEEP,
+            TransactionError::InstructionError(0, InstructionError::CallDepth),
+            &[
+                invoked_program_id.clone(),
+                invoked_program_id.clone(),
+                invoked_program_id.clone(),
+                invoked_program_id.clone(),
+                invoked_program_id.clone(),
+            ],
+        );
+
         // Check resulting state
 
         assert_eq!(43, bank.get_balance(&derived_key1));

runtime/src/message_processor.rs

@@ -330,11 +330,18 @@ impl<'a> InvokeContext for ThisInvokeContext<'a> {
         if self.invoke_stack.len() > self.bpf_compute_budget.max_invoke_depth {
             return Err(InstructionError::CallDepth);
         }
-        let frame_index = self.invoke_stack.iter().position(|frame| frame.key == *key);
+
-        if frame_index != None && frame_index != Some(self.invoke_stack.len().saturating_sub(1)) {
+        let contains = self.invoke_stack.iter().any(|frame| frame.key == *key);
+        let is_last = if let Some(last_frame) = self.invoke_stack.last() {
+            last_frame.key == *key
+        } else {
+            false
+        };
+        if contains && !is_last {
             // Reentrancy not allowed unless caller is calling itself
             return Err(InstructionError::ReentrancyNotAllowed);
         }
+
         // Alias the keys and account references in the provided keyed_accounts
         // with the ones already existing in self, so that the lifetime 'a matches.
         fn transmute_lifetime<'a, 'b, T: Sized>(value: &'a T) -> &'b T {

sdk/src/process_instruction.rs

@@ -163,7 +163,7 @@ pub struct BpfComputeBudget {
     /// Number of compute units consumed by an invoke call (not including the cost incurred by
     /// the called program)
     pub invoke_units: u64,
-    /// Maximum cross-program invocation depth allowed including the original caller
+    /// Maximum cross-program invocation depth allowed
     pub max_invoke_depth: usize,
     /// Base number of compute units consumed to call SHA256
     pub sha256_base_cost: u64,