v1.6: Restore ability for programs to upgrade themselves (#20263)

* Make helper associated fn

* Add feature definition

* Restore program-id write lock when upgradeable loader is present; restore bpf upgrade-self test

* Remove spurious comment

cli-output/src/display.rs

@@ -140,7 +140,7 @@ fn format_account_mode(message: &Message, index: usize) -> String {
         } else {
             "-"
         },
-        if message.is_writable(index) {
+        if message.is_writable(index, /* restore_write_lock_when_upgradeable=*/ true) {
             "w" // comment for consistent rust fmt (no joking; lol)
         } else {
             "-"

core/src/transaction_status_service.rs

@@ -110,7 +110,9 @@ impl TransactionStatusService {
                             .expect("FeeCalculator must exist");
                         let fee = fee_calculator.calculate_fee(transaction.message());
                         let (writable_keys, readonly_keys) =
-                            transaction.message.get_account_keys_by_lock_type();
+                            transaction.message.get_account_keys_by_lock_type(
+                                bank.restore_write_lock_when_upgradeable(),
+                            );
 
                         let inner_instructions = inner_instructions.map(|inner_instructions| {
                             inner_instructions

program-test/src/lib.rs

@@ -19,6 +19,7 @@ use {
         clock::{Clock, Slot},
         entrypoint::{ProgramResult, SUCCESS},
         epoch_schedule::EpochSchedule,
+        feature_set::restore_write_lock_when_upgradeable,
         fee_calculator::{FeeCalculator, FeeRateGovernor},
         genesis_config::{ClusterType, GenesisConfig},
         hash::Hash,
@@ -255,12 +256,14 @@ impl solana_sdk::program_stubs::SyscallStubs for SyscallStubs {
             }
             panic!("Program id {} wasn't found in account_infos", program_id);
         };
+        let restore_write_lock_when_upgradeable =
+            invoke_context.is_feature_active(&restore_write_lock_when_upgradeable::id());
         // TODO don't have the caller's keyed_accounts so can't validate writer or signer escalation or deescalation yet
         let caller_privileges = message
             .account_keys
             .iter()
             .enumerate()
-            .map(|(i, _)| message.is_writable(i))
+            .map(|(i, _)| message.is_writable(i, restore_write_lock_when_upgradeable))
             .collect::<Vec<bool>>();
 
         stable_log::program_invoke(&logger, &program_id, invoke_context.invoke_depth());
@@ -331,7 +334,7 @@ impl solana_sdk::program_stubs::SyscallStubs for SyscallStubs {
 
         // Copy writeable account modifications back into the caller's AccountInfos
         for (i, account_pubkey) in message.account_keys.iter().enumerate() {
-            if !message.is_writable(i) {
+            if !message.is_writable(i, restore_write_lock_when_upgradeable) {
                 continue;
             }
 

programs/bpf/tests/programs.rs

@@ -1782,7 +1782,7 @@ fn test_program_bpf_upgrade_and_invoke_in_same_tx() {
         "solana_bpf_rust_panic",
     );
 
-    // Attempt to invoke, then upgrade the program in same tx
+    // Invoke, then upgrade the program, and then invoke again in same tx
     let message = Message::new(
         &[
             invoke_instruction.clone(),
@@ -1801,12 +1801,10 @@ fn test_program_bpf_upgrade_and_invoke_in_same_tx() {
         message.clone(),
         bank.last_blockhash(),
     );
-    // program_id is automatically demoted to readonly, preventing the upgrade, which requires
-    // writeability
     let (result, _) = process_transaction_and_record_inner(&bank, tx);
     assert_eq!(
         result.unwrap_err(),
-        TransactionError::InstructionError(1, InstructionError::InvalidArgument)
+        TransactionError::InstructionError(2, InstructionError::ProgramFailedToComplete)
     );
 }
 
@@ -2087,6 +2085,97 @@ fn test_program_bpf_upgrade_via_cpi() {
     );
 }
 
+#[cfg(feature = "bpf_rust")]
+#[test]
+fn test_program_bpf_upgrade_self_via_cpi() {
+    solana_logger::setup();
+
+    let GenesisConfigInfo {
+        genesis_config,
+        mint_keypair,
+        ..
+    } = create_genesis_config(50);
+    let mut bank = Bank::new(&genesis_config);
+    let (name, id, entrypoint) = solana_bpf_loader_program!();
+    bank.add_builtin(&name, id, entrypoint);
+    let (name, id, entrypoint) = solana_bpf_loader_upgradeable_program!();
+    bank.add_builtin(&name, id, entrypoint);
+    let bank = Arc::new(bank);
+    let bank_client = BankClient::new_shared(&bank);
+    let noop_program_id = load_bpf_program(
+        &bank_client,
+        &bpf_loader::id(),
+        &mint_keypair,
+        "solana_bpf_rust_noop",
+    );
+
+    // Deploy upgradeable program
+    let buffer_keypair = Keypair::new();
+    let program_keypair = Keypair::new();
+    let program_id = program_keypair.pubkey();
+    let authority_keypair = Keypair::new();
+    load_upgradeable_bpf_program(
+        &bank_client,
+        &mint_keypair,
+        &buffer_keypair,
+        &program_keypair,
+        &authority_keypair,
+        "solana_bpf_rust_invoke_and_return",
+    );
+
+    let mut invoke_instruction = Instruction::new_with_bytes(
+        program_id,
+        &[0],
+        vec![
+            AccountMeta::new_readonly(noop_program_id, false),
+            AccountMeta::new_readonly(noop_program_id, false),
+            AccountMeta::new_readonly(clock::id(), false),
+            AccountMeta::new_readonly(fees::id(), false),
+        ],
+    );
+
+    // Call the upgraded program
+    invoke_instruction.data[0] += 1;
+    let result =
+        bank_client.send_and_confirm_instruction(&mint_keypair, invoke_instruction.clone());
+    assert!(result.is_ok());
+
+    // Prepare for upgrade
+    let buffer_keypair = Keypair::new();
+    load_upgradeable_buffer(
+        &bank_client,
+        &mint_keypair,
+        &buffer_keypair,
+        &authority_keypair,
+        "solana_bpf_rust_panic",
+    );
+
+    // Invoke, then upgrade the program, and then invoke again in same tx
+    let message = Message::new(
+        &[
+            invoke_instruction.clone(),
+            bpf_loader_upgradeable::upgrade(
+                &program_id,
+                &buffer_keypair.pubkey(),
+                &authority_keypair.pubkey(),
+                &mint_keypair.pubkey(),
+            ),
+            invoke_instruction,
+        ],
+        Some(&mint_keypair.pubkey()),
+    );
+    let tx = Transaction::new(
+        &[&mint_keypair, &authority_keypair],
+        message.clone(),
+        bank.last_blockhash(),
+    );
+    let (result, _) = process_transaction_and_record_inner(&bank, tx);
+    assert_eq!(
+        result.unwrap_err(),
+        TransactionError::InstructionError(2, InstructionError::ProgramFailedToComplete)
+    );
+}
+
 #[cfg(feature = "bpf_rust")]
 #[test]
 fn test_program_bpf_set_upgrade_authority_via_cpi() {

programs/bpf_loader/src/syscalls.rs

@@ -20,8 +20,8 @@ use solana_sdk::{
     epoch_schedule::EpochSchedule,
     feature_set::{
         cpi_data_cost, cpi_share_ro_and_exec_accounts, enforce_aligned_host_addrs,
-        keccak256_syscall_enabled, memory_ops_syscalls, set_upgrade_authority_via_cpi_enabled,
+        keccak256_syscall_enabled, memory_ops_syscalls, restore_write_lock_when_upgradeable,
-        sysvar_via_syscall, update_data_on_realloc,
+        set_upgrade_authority_via_cpi_enabled, sysvar_via_syscall, update_data_on_realloc,
     },
     hash::{Hasher, HASH_BYTES},
     ic_msg,
@@ -2140,7 +2140,14 @@ fn call<'a>(
     signers_seeds_len: u64,
     memory_mapping: &MemoryMapping,
 ) -> Result<u64, EbpfError<BpfError>> {
-    let (message, executables, accounts, account_refs, caller_write_privileges) = {
+    let (
+        message,
+        executables,
+        accounts,
+        account_refs,
+        caller_write_privileges,
+        restore_write_lock_when_upgradeable,
+    ) = {
         let invoke_context = syscall.get_context()?;
 
         invoke_context
@@ -2230,6 +2237,7 @@ fn call<'a>(
             accounts,
             account_refs,
             caller_write_privileges,
+            invoke_context.is_feature_active(&restore_write_lock_when_upgradeable::id()),
         )
     };
 
@@ -2255,7 +2263,9 @@ fn call<'a>(
         for (i, (account, account_ref)) in accounts.iter().zip(account_refs).enumerate() {
             let account = account.borrow();
             if let Some(mut account_ref) = account_ref {
-                if message.is_writable(i) && !account.executable() {
+                if message.is_writable(i, restore_write_lock_when_upgradeable)
+                    && !account.executable()
+                {
                     *account_ref.lamports = account.lamports();
                     *account_ref.owner = *account.owner();
                     if account_ref.data.len() != account.data().len() {

runtime/src/accounts.rs

@@ -170,8 +170,11 @@ impl Accounts {
         false
     }
 
-    fn construct_instructions_account(message: &Message) -> AccountSharedData {
+    fn construct_instructions_account(
-        let mut data = message.serialize_instructions();
+        message: &Message,
+        restore_write_lock_when_upgradeable: bool,
+    ) -> AccountSharedData {
+        let mut data = message.serialize_instructions(restore_write_lock_when_upgradeable);
         // add room for current instruction index.
         data.resize(data.len() + 2, 0);
         AccountSharedData::from(Account {
@@ -201,7 +204,8 @@ impl Accounts {
             let mut accounts = Vec::with_capacity(message.account_keys.len());
             let mut account_deps = Vec::with_capacity(message.account_keys.len());
             let mut rent_debits = RentDebits::default();
-            let is_upgradeable_loader_present = is_upgradeable_loader_present(message);
+            let restore_write_lock_when_upgradeable =
+                feature_set.is_active(&feature_set::restore_write_lock_when_upgradeable::id());
 
             for (i, key) in message.account_keys.iter().enumerate() {
                 let account = if message.is_non_loader_key(key, i) {
@@ -212,13 +216,16 @@ impl Accounts {
                     if solana_sdk::sysvar::instructions::check_id(key)
                         && feature_set.is_active(&feature_set::instructions_sysvar_enabled::id())
                     {
-                        Self::construct_instructions_account(message)
+                        Self::construct_instructions_account(
+                            message,
+                            restore_write_lock_when_upgradeable,
+                        )
                     } else {
                         let (account, rent) = self
                             .accounts_db
                             .load(ancestors, key)
                             .map(|(mut account, _)| {
-                                if message.is_writable(i) {
+                                if message.is_writable(i, restore_write_lock_when_upgradeable) {
                                     let rent_due = rent_collector
                                         .collect_from_existing_account(&key, &mut account);
                                     (account, rent_due)
@@ -229,7 +236,9 @@ impl Accounts {
                             .unwrap_or_default();
 
                         if bpf_loader_upgradeable::check_id(&account.owner) {
-                            if message.is_writable(i) && !is_upgradeable_loader_present {
+                            if message.is_writable(i, restore_write_lock_when_upgradeable)
+                                && !message.is_upgradeable_loader_present()
+                            {
                                 error_counters.invalid_writable_account += 1;
                                 return Err(TransactionError::InvalidWritableAccount);
                             }
@@ -255,7 +264,9 @@ impl Accounts {
                                     return Err(TransactionError::InvalidProgramForExecution);
                                 }
                             }
-                        } else if account.executable && message.is_writable(i) {
+                        } else if account.executable
+                            && message.is_writable(i, restore_write_lock_when_upgradeable)
+                        {
                             error_counters.invalid_writable_account += 1;
                             return Err(TransactionError::InvalidWritableAccount);
                         }
@@ -774,13 +785,21 @@ impl Accounts {
         Ok(())
     }
 
-    fn unlock_account(&self, tx: &Transaction, result: &Result<()>, locks: &mut AccountLocks) {
+    fn unlock_account(
+        &self,
+        tx: &Transaction,
+        result: &Result<()>,
+        locks: &mut AccountLocks,
+        restore_write_lock_when_upgradeable: bool,
+    ) {
         match result {
             Err(TransactionError::AccountInUse) => (),
             Err(TransactionError::SanitizeFailure) => (),
             Err(TransactionError::AccountLoadedTwice) => (),
             _ => {
-                let (writable_keys, readonly_keys) = &tx.message().get_account_keys_by_lock_type();
+                let (writable_keys, readonly_keys) = &tx
+                    .message()
+                    .get_account_keys_by_lock_type(restore_write_lock_when_upgradeable);
                 for k in writable_keys {
                     locks.unlock_write(k);
                 }
@@ -809,7 +828,11 @@ impl Accounts {
     /// This function will prevent multiple threads from modifying the same account state at the
     /// same time
     #[must_use]
-    pub fn lock_accounts<'a>(&self, txs: impl Iterator<Item = &'a Transaction>) -> Vec<Result<()>> {
+    pub fn lock_accounts<'a>(
+        &self,
+        txs: impl Iterator<Item = &'a Transaction>,
+        restore_write_lock_when_upgradeable: bool,
+    ) -> Vec<Result<()>> {
         use solana_sdk::sanitize::Sanitize;
         let keys: Vec<Result<_>> = txs
             .map(|tx| {
@@ -819,7 +842,9 @@ impl Accounts {
                     return Err(TransactionError::AccountLoadedTwice);
                 }
 
-                Ok(tx.message().get_account_keys_by_lock_type())
+                Ok(tx
+                    .message()
+                    .get_account_keys_by_lock_type(restore_write_lock_when_upgradeable))
             })
             .collect();
         let mut account_locks = &mut self.account_locks.lock().unwrap();
@@ -838,11 +863,17 @@ impl Accounts {
         &self,
         txs: impl Iterator<Item = &'a Transaction>,
         results: &[Result<()>],
+        restore_write_lock_when_upgradeable: bool,
     ) {
         let mut account_locks = self.account_locks.lock().unwrap();
         debug!("bank unlock accounts");
         for (tx, lock_result) in txs.zip(results) {
-            self.unlock_account(tx, lock_result, &mut account_locks);
+            self.unlock_account(
+                tx,
+                lock_result,
+                &mut account_locks,
+                restore_write_lock_when_upgradeable,
+            );
         }
     }
 
@@ -859,6 +890,7 @@ impl Accounts {
         last_blockhash_with_fee_calculator: &(Hash, FeeCalculator),
         fix_recent_blockhashes_sysvar_delay: bool,
         merge_nonce_error_into_system_error: bool,
+        restore_write_lock_when_upgradeable: bool,
     ) {
         let accounts_to_store = self.collect_accounts_to_store(
             txs,
@@ -868,6 +900,7 @@ impl Accounts {
             last_blockhash_with_fee_calculator,
             fix_recent_blockhashes_sysvar_delay,
             merge_nonce_error_into_system_error,
+            restore_write_lock_when_upgradeable,
         );
         self.accounts_db.store_cached(slot, &accounts_to_store);
     }
@@ -893,6 +926,7 @@ impl Accounts {
         last_blockhash_with_fee_calculator: &(Hash, FeeCalculator),
         fix_recent_blockhashes_sysvar_delay: bool,
         merge_nonce_error_into_system_error: bool,
+        restore_write_lock_when_upgradeable: bool,
     ) -> Vec<(&'a Pubkey, &'a AccountSharedData)> {
         let mut accounts = Vec::with_capacity(loaded.len());
         for (i, ((raccs, _nonce_rollback), tx)) in loaded.iter_mut().zip(txs).enumerate() {
@@ -945,7 +979,7 @@ impl Accounts {
                     fee_payer_index = Some(i);
                 }
                 let is_fee_payer = Some(i) == fee_payer_index;
-                if message.is_writable(i)
+                if message.is_writable(i, restore_write_lock_when_upgradeable)
                     && (res.is_ok()
                         || (maybe_nonce_rollback.is_some() && (is_nonce_account || is_fee_payer)))
                 {
@@ -1029,13 +1063,6 @@ pub fn prepare_if_nonce_account(
     false
 }
 
-fn is_upgradeable_loader_present(message: &Message) -> bool {
-    message
-        .account_keys
-        .iter()
-        .any(|&key| key == bpf_loader_upgradeable::id())
-}
-
 pub fn create_test_accounts(
     accounts: &Accounts,
     pubkeys: &mut Vec<Pubkey>,
@@ -1988,6 +2015,8 @@ mod tests {
         accounts.store_slow_uncached(0, &keypair2.pubkey(), &account2);
         accounts.store_slow_uncached(0, &keypair3.pubkey(), &account3);
 
+        let restore_write_lock_when_upgradeable = true;
+
         let instructions = vec![CompiledInstruction::new(2, &(), vec![0, 1])];
         let message = Message::new_with_compiled_instructions(
             1,
@@ -1998,7 +2027,8 @@ mod tests {
             instructions,
         );
         let tx = Transaction::new(&[&keypair0], message, Hash::default());
-        let results0 = accounts.lock_accounts([tx.clone()].iter());
+        let results0 =
+            accounts.lock_accounts([tx.clone()].iter(), restore_write_lock_when_upgradeable);
 
         assert!(results0[0].is_ok());
         assert_eq!(
@@ -2033,7 +2063,7 @@ mod tests {
         );
         let tx1 = Transaction::new(&[&keypair1], message, Hash::default());
         let txs = vec![tx0, tx1];
-        let results1 = accounts.lock_accounts(txs.iter());
+        let results1 = accounts.lock_accounts(txs.iter(), restore_write_lock_when_upgradeable);
 
         assert!(results1[0].is_ok()); // Read-only account (keypair1) can be referenced multiple times
         assert!(results1[1].is_err()); // Read-only account (keypair1) cannot also be locked as writable
@@ -2048,8 +2078,8 @@ mod tests {
             2
         );
 
-        accounts.unlock_accounts([tx].iter(), &results0);
+        accounts.unlock_accounts([tx].iter(), &results0, restore_write_lock_when_upgradeable);
-        accounts.unlock_accounts(txs.iter(), &results1);
+        accounts.unlock_accounts(txs.iter(), &results1, restore_write_lock_when_upgradeable);
         let instructions = vec![CompiledInstruction::new(2, &(), vec![0, 1])];
         let message = Message::new_with_compiled_instructions(
             1,
@@ -2060,7 +2090,7 @@ mod tests {
             instructions,
         );
         let tx = Transaction::new(&[&keypair1], message, Hash::default());
-        let results2 = accounts.lock_accounts([tx].iter());
+        let results2 = accounts.lock_accounts([tx].iter(), restore_write_lock_when_upgradeable);
         assert!(results2[0].is_ok()); // Now keypair1 account can be locked as writable
 
         // Check that read-only lock with zero references is deleted
@@ -2096,6 +2126,8 @@ mod tests {
         accounts.store_slow_uncached(0, &keypair1.pubkey(), &account1);
         accounts.store_slow_uncached(0, &keypair2.pubkey(), &account2);
 
+        let restore_write_lock_when_upgradeable = true;
+
         let accounts_arc = Arc::new(accounts);
 
         let instructions = vec![CompiledInstruction::new(2, &(), vec![0, 1])];
@@ -2128,13 +2160,19 @@ mod tests {
             let exit_clone = exit_clone.clone();
             loop {
                 let txs = vec![writable_tx.clone()];
-                let results = accounts_clone.clone().lock_accounts(txs.iter());
+                let results = accounts_clone
+                    .clone()
+                    .lock_accounts(txs.iter(), restore_write_lock_when_upgradeable);
                 for result in results.iter() {
                     if result.is_ok() {
                         counter_clone.clone().fetch_add(1, Ordering::SeqCst);
                     }
                 }
-                accounts_clone.unlock_accounts(txs.iter(), &results);
+                accounts_clone.unlock_accounts(
+                    txs.iter(),
+                    &results,
+                    restore_write_lock_when_upgradeable,
+                );
                 if exit_clone.clone().load(Ordering::Relaxed) {
                     break;
                 }
@@ -2143,13 +2181,15 @@ mod tests {
         let counter_clone = counter;
         for _ in 0..5 {
             let txs = vec![readonly_tx.clone()];
-            let results = accounts_arc.clone().lock_accounts(txs.iter());
+            let results = accounts_arc
+                .clone()
+                .lock_accounts(txs.iter(), restore_write_lock_when_upgradeable);
             if results[0].is_ok() {
                 let counter_value = counter_clone.clone().load(Ordering::SeqCst);
                 thread::sleep(time::Duration::from_millis(50));
                 assert_eq!(counter_value, counter_clone.clone().load(Ordering::SeqCst));
             }
-            accounts_arc.unlock_accounts(txs.iter(), &results);
+            accounts_arc.unlock_accounts(txs.iter(), &results, restore_write_lock_when_upgradeable);
             thread::sleep(time::Duration::from_millis(50));
         }
         exit.store(true, Ordering::Relaxed);
@@ -2178,6 +2218,8 @@ mod tests {
         accounts.store_slow_uncached(0, &keypair2.pubkey(), &account2);
         accounts.store_slow_uncached(0, &keypair3.pubkey(), &account3);
 
+        let restore_write_lock_when_upgradeable = true;
+
         let instructions = vec![CompiledInstruction::new(2, &(), vec![0, 1])];
         let message = Message::new_with_compiled_instructions(
             1,
@@ -2188,7 +2230,7 @@ mod tests {
             instructions,
         );
         let tx = Transaction::new(&[&keypair0], message, Hash::default());
-        let results0 = accounts.lock_accounts([tx].iter());
+        let results0 = accounts.lock_accounts([tx].iter(), restore_write_lock_when_upgradeable);
 
         assert!(results0[0].is_ok());
         // Instruction program-id account demoted to readonly
@@ -2305,6 +2347,7 @@ mod tests {
             &(Hash::default(), FeeCalculator::default()),
             true,
             true, // merge_nonce_error_into_system_error
+            true, // restore_write_lock_when_upgradeable
         );
         assert_eq!(collected_accounts.len(), 2);
         assert!(collected_accounts
@@ -2690,6 +2733,7 @@ mod tests {
             &(next_blockhash, FeeCalculator::default()),
             true,
             true, // merge_nonce_error_into_system_error
+            true, // restore_write_lock_when_upgradeable
         );
         assert_eq!(collected_accounts.len(), 2);
         assert_eq!(
@@ -2806,6 +2850,7 @@ mod tests {
             &(next_blockhash, FeeCalculator::default()),
             true,
             true, // merge_nonce_error_into_system_error
+            true, // restore_write_lock_when_upgradeable
         );
         assert_eq!(collected_accounts.len(), 1);
         let collected_nonce_account = collected_accounts

runtime/src/bank.rs

@@ -2503,10 +2503,10 @@ impl Bank {
         txs: impl Iterator<Item = &'b Transaction>,
     ) -> TransactionBatch<'a, 'b> {
         let hashed_txs: Vec<HashedTransaction> = txs.map(HashedTransaction::from).collect();
-        let lock_results = self
+        let lock_results = self.rc.accounts.lock_accounts(
-            .rc
+            hashed_txs.as_transactions_iter(),
-            .accounts
+            self.restore_write_lock_when_upgradeable(),
-            .lock_accounts(hashed_txs.as_transactions_iter());
+        );
         TransactionBatch::new(lock_results, self, Cow::Owned(hashed_txs))
     }
 
@@ -2514,10 +2514,10 @@ impl Bank {
         &'a self,
         hashed_txs: &'b [HashedTransaction],
     ) -> TransactionBatch<'a, 'b> {
-        let lock_results = self
+        let lock_results = self.rc.accounts.lock_accounts(
-            .rc
+            hashed_txs.as_transactions_iter(),
-            .accounts
+            self.restore_write_lock_when_upgradeable(),
-            .lock_accounts(hashed_txs.as_transactions_iter());
+        );
         TransactionBatch::new(lock_results, self, Cow::Borrowed(hashed_txs))
     }
 
@@ -2585,9 +2585,11 @@ impl Bank {
     pub fn unlock_accounts(&self, batch: &mut TransactionBatch) {
         if batch.needs_unlock {
             batch.needs_unlock = false;
-            self.rc
+            self.rc.accounts.unlock_accounts(
-                .accounts
+                batch.transactions_iter(),
-                .unlock_accounts(batch.transactions_iter(), batch.lock_results())
+                batch.lock_results(),
+                self.restore_write_lock_when_upgradeable(),
+            )
         }
     }
 
@@ -3333,6 +3335,7 @@ impl Bank {
             &self.last_blockhash_with_fee_calculator(),
             self.fix_recent_blockhashes_sysvar_delay(),
             self.merge_nonce_error_into_system_error(),
+            self.restore_write_lock_when_upgradeable(),
         );
         let rent_debits = self.collect_rent(executed, loaded_accounts);
 
@@ -4888,6 +4891,11 @@ impl Bank {
             .is_active(&feature_set::stakes_remove_delegation_if_inactive::id())
     }
 
+    pub fn restore_write_lock_when_upgradeable(&self) -> bool {
+        self.feature_set
+            .is_active(&feature_set::restore_write_lock_when_upgradeable::id())
+    }
+
     // Check if the wallclock time from bank creation to now has exceeded the allotted
     // time for transaction processing
     pub fn should_bank_still_be_processing_txs(

runtime/src/message_processor.rs

@@ -9,8 +9,8 @@ use solana_sdk::{
     account_utils::StateMut,
     bpf_loader_upgradeable::{self, UpgradeableLoaderState},
     feature_set::{
-        cpi_share_ro_and_exec_accounts, instructions_sysvar_enabled, updated_verify_policy,
+        cpi_share_ro_and_exec_accounts, instructions_sysvar_enabled,
-        FeatureSet,
+        restore_write_lock_when_upgradeable, updated_verify_policy, FeatureSet,
     },
     ic_msg,
     instruction::{CompiledInstruction, Instruction, InstructionError},
@@ -360,6 +360,8 @@ impl<'a> InvokeContext for ThisInvokeContext<'a> {
                 caller_write_privileges,
                 &mut self.timings,
                 self.feature_set.is_active(&updated_verify_policy::id()),
+                self.feature_set
+                    .is_active(&restore_write_lock_when_upgradeable::id()),
             ),
             None => Err(InstructionError::GenericError), // Should never happen
         }
@@ -572,6 +574,7 @@ impl MessageProcessor {
         instruction: &'a CompiledInstruction,
         executable_accounts: &'a [(Pubkey, Rc<RefCell<AccountSharedData>>)],
         accounts: &'a [Rc<RefCell<AccountSharedData>>],
+        restore_write_lock_when_upgradeable: bool,
     ) -> Vec<KeyedAccount<'a>> {
         let mut keyed_accounts = create_keyed_readonly_accounts(&executable_accounts);
         let mut keyed_accounts2: Vec<_> = instruction
@@ -582,7 +585,7 @@ impl MessageProcessor {
                 let index = index as usize;
                 let key = &message.account_keys[index];
                 let account = &accounts[index];
-                if message.is_writable(index) {
+                if message.is_writable(index, restore_write_lock_when_upgradeable) {
                     KeyedAccount::new(key, is_signer, account)
                 } else {
                     KeyedAccount::new_readonly(key, is_signer, account)
@@ -727,7 +730,14 @@ impl MessageProcessor {
     ) -> Result<(), InstructionError> {
         let invoke_context = RefCell::new(invoke_context);
 
-        let (message, executables, accounts, account_refs, caller_write_privileges) = {
+        let (
+            message,
+            executables,
+            accounts,
+            account_refs,
+            caller_write_privileges,
+            restore_write_lock_when_upgradeable,
+        ) = {
             let invoke_context = invoke_context.borrow();
 
             let caller_program_id = invoke_context.get_caller()?;
@@ -819,6 +829,7 @@ impl MessageProcessor {
                 accounts,
                 account_refs,
                 caller_write_privileges,
+                invoke_context.is_feature_active(&restore_write_lock_when_upgradeable::id()),
             )
         };
 
@@ -837,7 +848,9 @@ impl MessageProcessor {
             let invoke_context = invoke_context.borrow();
             for (i, (account, account_ref)) in accounts.iter().zip(account_refs).enumerate() {
                 let account = account.borrow();
-                if message.is_writable(i) && !account.executable {
+                if message.is_writable(i, restore_write_lock_when_upgradeable)
+                    && !account.executable
+                {
                     account_ref.try_account_ref_mut()?.lamports = account.lamports;
                     account_ref.try_account_ref_mut()?.owner = account.owner;
                     if account_ref.data_len()? != account.data().len()
@@ -881,8 +894,15 @@ impl MessageProcessor {
                 Some(caller_write_privileges),
             )?;
             // Construct keyed accounts
-            let keyed_accounts =
+            let restore_write_lock_when_upgradeable =
-                Self::create_keyed_accounts(message, instruction, executable_accounts, accounts);
+                invoke_context.is_feature_active(&restore_write_lock_when_upgradeable::id());
+            let keyed_accounts = Self::create_keyed_accounts(
+                message,
+                instruction,
+                executable_accounts,
+                accounts,
+                restore_write_lock_when_upgradeable,
+            );
 
             // Invoke callee
             invoke_context.push(program_id)?;
@@ -954,6 +974,7 @@ impl MessageProcessor {
         rent: &Rent,
         timings: &mut ExecuteDetailsTimings,
         updated_verify_policy: bool,
+        restore_write_lock_when_upgradeable: bool,
     ) -> Result<(), InstructionError> {
         // Verify all executable accounts have zero outstanding refs
         Self::verify_account_references(executable_accounts)?;
@@ -972,7 +993,7 @@ impl MessageProcessor {
                 let account = accounts[account_index].borrow();
                 pre_accounts[unique_index].verify(
                     &program_id,
-                    message.is_writable(account_index),
+                    message.is_writable(account_index, restore_write_lock_when_upgradeable),
                     rent,
                     &account,
                     timings,
@@ -1004,6 +1025,7 @@ impl MessageProcessor {
         caller_write_privileges: Option<&[bool]>,
         timings: &mut ExecuteDetailsTimings,
         updated_verify_policy: bool,
+        restore_write_lock_when_upgradeable: bool,
     ) -> Result<(), InstructionError> {
         // Verify the per-account instruction results
         let (mut pre_sum, mut post_sum) = (0_u128, 0_u128);
@@ -1014,7 +1036,7 @@ impl MessageProcessor {
                 let is_writable = if let Some(caller_write_privileges) = caller_write_privileges {
                     caller_write_privileges[account_index]
                 } else {
-                    message.is_writable(account_index)
+                    message.is_writable(account_index, restore_write_lock_when_upgradeable)
                 };
                 // Find the matching PreAccount
                 for pre_account in pre_accounts.iter_mut() {
@@ -1110,8 +1132,15 @@ impl MessageProcessor {
             account_db,
             ancestors,
         );
-        let keyed_accounts =
+        let restore_write_lock_when_upgradeable =
-            Self::create_keyed_accounts(message, instruction, executable_accounts, accounts);
+            invoke_context.is_feature_active(&restore_write_lock_when_upgradeable::id());
+        let keyed_accounts = Self::create_keyed_accounts(
+            message,
+            instruction,
+            executable_accounts,
+            accounts,
+            restore_write_lock_when_upgradeable,
+        );
         self.process_instruction(
             program_id,
             &keyed_accounts,
@@ -1127,6 +1156,7 @@ impl MessageProcessor {
             &rent_collector.rent,
             timings,
             invoke_context.is_feature_active(&updated_verify_policy::id()),
+            restore_write_lock_when_upgradeable,
         )?;
 
         timings.accumulate(&invoke_context.timings);
@@ -2142,6 +2172,8 @@ mod tests {
         let programs: Vec<(_, ProcessInstructionWithContext)> =
             vec![(callee_program_id, mock_process_instruction)];
         let feature_set = FeatureSet::all_enabled();
+        let restore_write_lock_when_upgradeable =
+            feature_set.is_active(&restore_write_lock_when_upgradeable::id());
 
         let ancestors = Ancestors::default();
         let mut invoke_context = ThisInvokeContext::new(
@@ -2180,7 +2212,7 @@ mod tests {
             .account_keys
             .iter()
             .enumerate()
-            .map(|(i, _)| message.is_writable(i))
+            .map(|(i, _)| message.is_writable(i, restore_write_lock_when_upgradeable))
             .collect::<Vec<bool>>();
         assert_eq!(
             MessageProcessor::process_cross_program_instruction(
@@ -2215,7 +2247,7 @@ mod tests {
                 .account_keys
                 .iter()
                 .enumerate()
-                .map(|(i, _)| message.is_writable(i))
+                .map(|(i, _)| message.is_writable(i, restore_write_lock_when_upgradeable))
                 .collect::<Vec<bool>>();
             assert_eq!(
                 MessageProcessor::process_cross_program_instruction(

sdk/benches/serialize_instructions.rs

@@ -8,6 +8,8 @@ use solana_sdk::pubkey;
 use solana_sdk::sysvar::instructions;
 use test::Bencher;
 
+const RESTORE_WRITE_LOCK_WHEN_UPGRADEABLE: bool = true;
+
 fn make_instructions() -> Vec<Instruction> {
     let meta = AccountMeta::new(pubkey::new_rand(), false);
     let inst = Instruction::new_with_bincode(pubkey::new_rand(), &[0; 10], vec![meta; 4]);
@@ -27,7 +29,7 @@ fn bench_manual_instruction_serialize(b: &mut Bencher) {
     let instructions = make_instructions();
     let message = Message::new(&instructions, None);
     b.iter(|| {
-        test::black_box(message.serialize_instructions());
+        test::black_box(message.serialize_instructions(RESTORE_WRITE_LOCK_WHEN_UPGRADEABLE));
     });
 }
 
@@ -44,7 +46,7 @@ fn bench_bincode_instruction_deserialize(b: &mut Bencher) {
 fn bench_manual_instruction_deserialize(b: &mut Bencher) {
     let instructions = make_instructions();
     let message = Message::new(&instructions, None);
-    let serialized = message.serialize_instructions();
+    let serialized = message.serialize_instructions(RESTORE_WRITE_LOCK_WHEN_UPGRADEABLE);
     b.iter(|| {
         for i in 0..instructions.len() {
             test::black_box(instructions::load_instruction_at(i, &serialized).unwrap());
@@ -56,7 +58,7 @@ fn bench_manual_instruction_deserialize(b: &mut Bencher) {
 fn bench_manual_instruction_deserialize_single(b: &mut Bencher) {
     let instructions = make_instructions();
     let message = Message::new(&instructions, None);
-    let serialized = message.serialize_instructions();
+    let serialized = message.serialize_instructions(RESTORE_WRITE_LOCK_WHEN_UPGRADEABLE);
     b.iter(|| {
         test::black_box(instructions::load_instruction_at(3, &serialized).unwrap());
     });

sdk/program/src/message.rs

@@ -373,7 +373,7 @@ impl Message {
         self.program_position(i).is_some()
     }
 
-    pub fn is_writable(&self, i: usize) -> bool {
+    pub fn is_writable(&self, i: usize, restore_write_lock_when_upgradeable: bool) -> bool {
         (i < (self.header.num_required_signatures - self.header.num_readonly_signed_accounts)
             as usize
             || (i >= self.header.num_required_signatures as usize
@@ -383,18 +383,22 @@ impl Message {
                 let key = self.account_keys[i];
                 sysvar::is_sysvar_id(&key) || BUILTIN_PROGRAMS_KEYS.contains(&key)
             }
-            && !self.is_key_called_as_program(i)
+            && (!self.is_key_called_as_program(i)
+                || (restore_write_lock_when_upgradeable && self.is_upgradeable_loader_present()))
     }
 
     pub fn is_signer(&self, i: usize) -> bool {
         i < self.header.num_required_signatures as usize
     }
 
-    pub fn get_account_keys_by_lock_type(&self) -> (Vec<&Pubkey>, Vec<&Pubkey>) {
+    pub fn get_account_keys_by_lock_type(
+        &self,
+        restore_write_lock_when_upgradeable: bool,
+    ) -> (Vec<&Pubkey>, Vec<&Pubkey>) {
         let mut writable_keys = vec![];
         let mut readonly_keys = vec![];
         for (i, key) in self.account_keys.iter().enumerate() {
-            if self.is_writable(i) {
+            if self.is_writable(i, restore_write_lock_when_upgradeable) {
                 writable_keys.push(key);
             } else {
                 readonly_keys.push(key);
@@ -416,7 +420,7 @@ impl Message {
     //   35..67 - program_id
     //   67..69 - data len - u16
     //   69..data_len - data
-    pub fn serialize_instructions(&self) -> Vec<u8> {
+    pub fn serialize_instructions(&self, restore_write_lock_when_upgradeable: bool) -> Vec<u8> {
         // 64 bytes is a reasonable guess, calculating exactly is slower in benchmarks
         let mut data = Vec::with_capacity(self.instructions.len() * (32 * 2));
         append_u16(&mut data, self.instructions.len() as u16);
@@ -431,7 +435,8 @@ impl Message {
             for account_index in &instruction.accounts {
                 let account_index = *account_index as usize;
                 let is_signer = self.is_signer(account_index);
-                let is_writable = self.is_writable(account_index);
+                let is_writable =
+                    self.is_writable(account_index, restore_write_lock_when_upgradeable);
                 let mut meta_byte = 0;
                 if is_signer {
                     meta_byte |= 1 << Self::IS_SIGNER_BIT;
@@ -506,6 +511,12 @@ impl Message {
             .min(self.header.num_required_signatures as usize);
         self.account_keys[..last_key].iter().collect()
     }
+
+    pub fn is_upgradeable_loader_present(&self) -> bool {
+        self.account_keys
+            .iter()
+            .any(|&key| key == bpf_loader_upgradeable::id())
+    }
 }
 
 #[cfg(test)]
@@ -870,12 +881,13 @@ mod tests {
             recent_blockhash: Hash::default(),
             instructions: vec![],
         };
-        assert!(message.is_writable(0));
+        let restore_write_lock_when_upgradeable = true;
-        assert!(!message.is_writable(1));
+        assert!(message.is_writable(0, restore_write_lock_when_upgradeable));
-        assert!(!message.is_writable(2));
+        assert!(!message.is_writable(1, restore_write_lock_when_upgradeable));
-        assert!(message.is_writable(3));
+        assert!(!message.is_writable(2, restore_write_lock_when_upgradeable));
-        assert!(message.is_writable(4));
+        assert!(message.is_writable(3, restore_write_lock_when_upgradeable));
-        assert!(!message.is_writable(5));
+        assert!(message.is_writable(4, restore_write_lock_when_upgradeable));
+        assert!(!message.is_writable(5, restore_write_lock_when_upgradeable));
     }
 
     #[test]
@@ -903,7 +915,7 @@ mod tests {
             Some(&id1),
         );
         assert_eq!(
-            message.get_account_keys_by_lock_type(),
+            message.get_account_keys_by_lock_type(/*restore_write_lock_when_upgradeable=*/ true),
             (vec![&id1, &id0], vec![&id3, &id2, &program_id])
         );
     }
@@ -933,7 +945,8 @@ mod tests {
         ];
 
         let message = Message::new(&instructions, Some(&id1));
-        let serialized = message.serialize_instructions();
+        let serialized =
+            message.serialize_instructions(/*restore_write_lock_when_upgradeable=*/ true);
         for (i, instruction) in instructions.iter().enumerate() {
             assert_eq!(
                 Message::deserialize_instruction(i, &serialized).unwrap(),
@@ -954,7 +967,8 @@ mod tests {
         ];
 
         let message = Message::new(&instructions, Some(&id1));
-        let serialized = message.serialize_instructions();
+        let serialized =
+            message.serialize_instructions(/*restore_write_lock_when_upgradeable=*/ true);
         assert_eq!(
             Message::deserialize_instruction(instructions.len(), &serialized).unwrap_err(),
             SanitizeError::IndexOutOfBounds,

sdk/src/feature_set.rs

@@ -174,6 +174,10 @@ pub mod stakes_remove_delegation_if_inactive {
     solana_sdk::declare_id!("HFpdDDNQjvcXnXKec697HDDsyk6tFoWS2o8fkxuhQZpL");
 }
 
+pub mod restore_write_lock_when_upgradeable {
+    solana_sdk::declare_id!("3Tye2iVqQTxprFSJNpyz5W6SjKNQVfRUDR2s3oVYS6h6");
+}
+
 lazy_static! {
     /// Map of feature identifiers to user-visible description
     pub static ref FEATURE_NAMES: HashMap<Pubkey, &'static str> = [
@@ -216,6 +220,7 @@ lazy_static! {
         (spl_token_v2_set_authority_fix::id(), "spl-token set_authority fix"),
         (demote_program_write_locks::id(), "demote program write locks to readonly #19593"),
         (stakes_remove_delegation_if_inactive::id(), "remove delegations from stakes cache when inactive"),
+        (restore_write_lock_when_upgradeable::id(), "restore program-id write lock when upgradeable loader present"),
         /*************** ADD NEW FEATURES HERE ***************/
     ]
     .iter()

transaction-status/src/parse_accounts.rs

@@ -13,7 +13,7 @@ pub fn parse_accounts(message: &Message) -> Vec<ParsedAccount> {
     for (i, account_key) in message.account_keys.iter().enumerate() {
         accounts.push(ParsedAccount {
             pubkey: account_key.to_string(),
-            writable: message.is_writable(i),
+            writable: message.is_writable(i, /*restore_write_lock_when_upgradeable=*/ true),
             signer: message.is_signer(i),
         });
     }