From b76784774a45e2b82e8a4d7f4ae25a729e6d4906 Mon Sep 17 00:00:00 2001
From: "mergify[bot]" <37929162+mergify[bot]@users.noreply.github.com>
Date: Thu, 16 Nov 2023 06:29:46 +1100
Subject: [PATCH] v1.17: cpi: fix capacity check in update_caller_account
 (backport of #34064) (#34081)

cpi: fix capacity check in update_caller_account (#34064)

reserve(additional) reserves additional bytes on top of the current _length_
not capacity. Before this fix we could potentially reserve less capacity than
required.

(cherry picked from commit d009d7304a7c82bddc9809684ef27183b766de71)

Co-authored-by: Alessandro Decina <alessandro.d@gmail.com>
---
 programs/bpf_loader/src/syscalls/cpi.rs   |  3 ++-
 programs/sbf/rust/invoke/src/processor.rs | 27 +++++++++++++++++++++++
 programs/sbf/tests/programs.rs            | 16 +++++++++-----
 3 files changed, 40 insertions(+), 6 deletions(-)

diff --git a/programs/bpf_loader/src/syscalls/cpi.rs b/programs/bpf_loader/src/syscalls/cpi.rs
index 3f8190d26df143..c91e3705423151 100644
--- a/programs/bpf_loader/src/syscalls/cpi.rs
+++ b/programs/bpf_loader/src/syscalls/cpi.rs
@@ -1470,7 +1470,8 @@ fn update_caller_account(
             // invalid address.
             let min_capacity = caller_account.original_data_len;
             if callee_account.capacity() < min_capacity {
-                callee_account.reserve(min_capacity.saturating_sub(callee_account.capacity()))?;
+                callee_account
+                    .reserve(min_capacity.saturating_sub(callee_account.get_data().len()))?;
                 zero_all_mapped_spare_capacity = true;
             }
 
diff --git a/programs/sbf/rust/invoke/src/processor.rs b/programs/sbf/rust/invoke/src/processor.rs
index 7c689cfcf860ae..a14ee76c2e301b 100644
--- a/programs/sbf/rust/invoke/src/processor.rs
+++ b/programs/sbf/rust/invoke/src/processor.rs
@@ -1286,6 +1286,33 @@ fn process_instruction(
                 },
                 &vec![0; original_data_len - new_len]
             );
+
+            // Realloc to [0xFC; 2]
+            invoke(
+                &create_instruction(
+                    *callee_program_id,
+                    &[
+                        (accounts[ARGUMENT_INDEX].key, true, false),
+                        (callee_program_id, false, false),
+                    ],
+                    vec![0xFC; 2],
+                ),
+                accounts,
+            )
+            .unwrap();
+
+            // Check that [2..20] is zeroed
+            let new_len = account.data_len();
+            assert_eq!(&*account.data.borrow(), &[0xFC; 2]);
+            assert_eq!(
+                unsafe {
+                    slice::from_raw_parts(
+                        account.data.borrow().as_ptr().add(new_len),
+                        original_data_len - new_len,
+                    )
+                },
+                &vec![0; original_data_len - new_len]
+            );
         }
         TEST_WRITE_ACCOUNT => {
             msg!("TEST_WRITE_ACCOUNT");
diff --git a/programs/sbf/tests/programs.rs b/programs/sbf/tests/programs.rs
index 8a7cfb693afdd9..83c1ba00a77a22 100644
--- a/programs/sbf/tests/programs.rs
+++ b/programs/sbf/tests/programs.rs
@@ -4413,12 +4413,18 @@ fn test_cpi_change_account_data_memory_allocation() {
 
         // Test changing the account data both in place and by changing the
         // underlying vector. CPI will have to detect the vector change and
-        // update the corresponding memory region. In both cases CPI will have
+        // update the corresponding memory region. In all cases CPI will have
         // to zero the spare bytes correctly.
-        if instruction_data[0] == 0xFE {
-            account.set_data(instruction_data.to_vec());
-        } else {
-            account.set_data_from_slice(instruction_data);
+        match instruction_data[0] {
+            0xFE => account.set_data(instruction_data.to_vec()),
+            0xFD => account.set_data_from_slice(instruction_data),
+            0xFC => {
+                // Exercise the update_caller_account capacity check where account len != capacity.
+                let mut data = instruction_data.to_vec();
+                data.reserve_exact(1);
+                account.set_data(data)
+            }
+            _ => panic!(),
         }
 
         Ok(())