From 9c3cdf0c3606545a2bcad78300795f3896ce1ff0 Mon Sep 17 00:00:00 2001
From: Jack May <jack@solana.com>
Date: Fri, 27 Aug 2021 10:53:07 -0700
Subject: [PATCH] Allow programs to realloc their accounts within limits

---
 program-runtime/src/instruction_processor.rs | 28 +++++++++++++-------
 programs/bpf_loader/src/serialization.rs     | 10 +++----
 2 files changed, 23 insertions(+), 15 deletions(-)

diff --git a/program-runtime/src/instruction_processor.rs b/program-runtime/src/instruction_processor.rs
index 7b5ed32b57242c..e1cc7ecb275bb7 100644
--- a/program-runtime/src/instruction_processor.rs
+++ b/program-runtime/src/instruction_processor.rs
@@ -11,7 +11,7 @@ use solana_sdk::{
     process_instruction::{Executor, InvokeContext, Logger, ProcessInstructionWithContext},
     pubkey::Pubkey,
     rent::Rent,
-    system_program,
+    system_instruction::MAX_PERMITTED_DATA_LENGTH,
 };
 use std::{
     cell::{Ref, RefCell},
@@ -151,13 +151,14 @@ impl PreAccount {
             }
         }
 
-        // Only the system program can change the size of the data
-        //  and only if the system program owns the account
+        // Account data size cannot exceed a maxumum length
+        if post.data().len() > MAX_PERMITTED_DATA_LENGTH as usize {
+            return Err(InstructionError::InvalidRealloc);
+        }
+
+        // The owner of the account can change the size of the data
         let data_len_changed = pre.data().len() != post.data().len();
-        if data_len_changed
-            && (!system_program::check_id(program_id) // line coverage used to get branch coverage
-                || !system_program::check_id(pre.owner()))
-        {
+        if data_len_changed && program_id != pre.owner() {
             return Err(InstructionError::AccountDataSizeChanged);
         }
 
@@ -774,7 +775,7 @@ impl InstructionProcessor {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use solana_sdk::{account::Account, instruction::InstructionError};
+    use solana_sdk::{account::Account, instruction::InstructionError, system_program};
 
     #[test]
     fn test_is_zeroed() {
@@ -1171,11 +1172,18 @@ mod tests {
             "system program should not be able to change another program's account data size"
         );
         assert_eq!(
-            Change::new(&alice_program_id, &alice_program_id)
+            Change::new(&alice_program_id, &solana_sdk::pubkey::new_rand())
                 .data(vec![0], vec![0, 0])
                 .verify(),
             Err(InstructionError::AccountDataSizeChanged),
-            "non-system programs cannot change their data size"
+            "another program should not be able to change another program's account data size"
+        );
+        assert_eq!(
+            Change::new(&alice_program_id, &alice_program_id)
+                .data(vec![0], vec![0, 0])
+                .verify(),
+            Ok(()),
+            "Programs can change their own data size"
         );
         assert_eq!(
             Change::new(&system_program::id(), &system_program::id())
diff --git a/programs/bpf_loader/src/serialization.rs b/programs/bpf_loader/src/serialization.rs
index 4f842abebb36d7..db8ba0383ef6f7 100644
--- a/programs/bpf_loader/src/serialization.rs
+++ b/programs/bpf_loader/src/serialization.rs
@@ -7,6 +7,7 @@ use solana_sdk::{
     instruction::InstructionError,
     keyed_account::KeyedAccount,
     pubkey::Pubkey,
+    system_instruction::MAX_PERMITTED_DATA_LENGTH,
 };
 use std::{
     io::prelude::*,
@@ -268,13 +269,12 @@ pub fn deserialize_parameters_aligned(
             let pre_len = account.data().len();
             let post_len = LittleEndian::read_u64(&buffer[start..]) as usize;
             start += size_of::<u64>(); // data length
-            let mut data_end = start + pre_len;
-            if post_len != pre_len
-                && (post_len.saturating_sub(pre_len)) <= MAX_PERMITTED_DATA_INCREASE
+            if post_len.saturating_sub(pre_len) > MAX_PERMITTED_DATA_INCREASE
+                || post_len > MAX_PERMITTED_DATA_LENGTH as usize
             {
-                data_end = start + post_len;
+                return Err(InstructionError::InvalidRealloc);
             }
-
+            let data_end = start + post_len;
             account.set_data_from_slice(&buffer[start..data_end]);
             start += pre_len + MAX_PERMITTED_DATA_INCREASE; // data
             start += (start as *const u8).align_offset(align_of::<u128>());