Add handling to preserve program-id write lock when upgradeable loade…

…r is present; restore bpf upgrade-self test

cli-output/src/display.rs

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

core/src/banking_stage.rs

@@ -1091,7 +1091,7 @@ impl BankingStage {
         feature_set: &Arc<feature_set::FeatureSet>,
         cost_tracker: &Arc<RwLock<CostTracker>>,
         banking_stage_stats: &BankingStageStats,
-        demote_program_write_locks: bool,
+        demote_program_write_lock_features: bool,
         votes_only: bool,
     ) -> (Vec<SanitizedTransaction>, Vec<usize>, Vec<usize>) {
         let mut retryable_transaction_packet_indexes: Vec<usize> = vec![];
@@ -1126,7 +1126,7 @@ impl BankingStage {
                 .into_iter()
                 .filter_map(|(tx, tx_index)| {
                     let result = cost_tracker_readonly
-                        .would_transaction_fit(&tx, demote_program_write_locks);
+                        .would_transaction_fit(&tx, demote_program_write_lock_features);
                     if result.is_err() {
                         debug!("transaction {:?} would exceed limit: {:?}", tx, result);
                         retryable_transaction_packet_indexes.push(tx_index);
@@ -1208,7 +1208,7 @@ impl BankingStage {
                 &bank.feature_set,
                 cost_tracker,
                 banking_stage_stats,
-                bank.demote_program_write_locks(),
+                bank.demote_program_write_lock_features(),
                 bank.vote_only_bank(),
             );
         packet_conversion_time.stop();
@@ -1249,7 +1249,7 @@ impl BankingStage {
                 cost_tracker
                     .write()
                     .unwrap()
-                    .add_transaction_cost(tx, bank.demote_program_write_locks());
+                    .add_transaction_cost(tx, bank.demote_program_write_lock_features());
             }
         });
         cost_tracking_time.stop();
@@ -1315,7 +1315,7 @@ impl BankingStage {
                 &bank.feature_set,
                 cost_tracker,
                 banking_stage_stats,
-                bank.demote_program_write_locks(),
+                bank.demote_program_write_lock_features(),
                 bank.vote_only_bank(),
             );
         unprocessed_packet_conversion_time.stop();

core/src/cost_model.rs

@@ -116,7 +116,7 @@ impl CostModel {
     pub fn calculate_cost(
         &mut self,
         transaction: &SanitizedTransaction,
-        demote_program_write_locks: bool,
+        demote_program_write_lock_features: bool,
     ) -> &TransactionCost {
         self.transaction_cost.reset();
 
@@ -126,7 +126,7 @@ impl CostModel {
         // calculate account access cost
         let message = transaction.message();
         message.account_keys_iter().enumerate().for_each(|(i, k)| {
-            let is_writable = message.is_writable(i, demote_program_write_locks);
+            let is_writable = message.is_writable(i, demote_program_write_lock_features);
 
             if is_writable {
                 self.transaction_cost.writable_accounts.push(*k);
@@ -357,7 +357,8 @@ mod tests {
         .unwrap();
 
         let mut cost_model = CostModel::default();
-        let tx_cost = cost_model.calculate_cost(&tx, /*demote_program_write_locks=*/ true);
+        let tx_cost =
+            cost_model.calculate_cost(&tx, /*demote_program_write_lock_features=*/ true);
         assert_eq!(2 + 2, tx_cost.writable_accounts.len());
         assert_eq!(signer1.pubkey(), tx_cost.writable_accounts[0]);
         assert_eq!(signer2.pubkey(), tx_cost.writable_accounts[1]);
@@ -399,7 +400,8 @@ mod tests {
         cost_model
             .upsert_instruction_cost(&system_program::id(), expected_execution_cost)
             .unwrap();
-        let tx_cost = cost_model.calculate_cost(&tx, /*demote_program_write_locks=*/ true);
+        let tx_cost =
+            cost_model.calculate_cost(&tx, /*demote_program_write_lock_features=*/ true);
         assert_eq!(expected_account_cost, tx_cost.account_access_cost);
         assert_eq!(expected_execution_cost, tx_cost.execution_cost);
         assert_eq!(2, tx_cost.writable_accounts.len());
@@ -470,7 +472,7 @@ mod tests {
                     thread::spawn(move || {
                         let mut cost_model = cost_model.write().unwrap();
                         let tx_cost = cost_model
-                            .calculate_cost(&tx, /*demote_program_write_locks=*/ true);
+                            .calculate_cost(&tx, /*demote_program_write_lock_features=*/ true);
                         assert_eq!(3, tx_cost.writable_accounts.len());
                         assert_eq!(expected_account_cost, tx_cost.account_access_cost);
                     })

core/src/cost_tracker.rs

@@ -46,10 +46,10 @@ impl CostTracker {
     pub fn would_transaction_fit(
         &self,
         transaction: &SanitizedTransaction,
-        demote_program_write_locks: bool,
+        demote_program_write_lock_features: bool,
     ) -> Result<(), CostModelError> {
         let mut cost_model = self.cost_model.write().unwrap();
-        let tx_cost = cost_model.calculate_cost(transaction, demote_program_write_locks);
+        let tx_cost = cost_model.calculate_cost(transaction, demote_program_write_lock_features);
         self.would_fit(
             &tx_cost.writable_accounts,
             &(tx_cost.account_access_cost + tx_cost.execution_cost),
@@ -59,10 +59,10 @@ impl CostTracker {
     pub fn add_transaction_cost(
         &mut self,
         transaction: &SanitizedTransaction,
-        demote_program_write_locks: bool,
+        demote_program_write_lock_features: bool,
     ) {
         let mut cost_model = self.cost_model.write().unwrap();
-        let tx_cost = cost_model.calculate_cost(transaction, demote_program_write_locks);
+        let tx_cost = cost_model.calculate_cost(transaction, demote_program_write_lock_features);
         let cost = tx_cost.account_access_cost + tx_cost.execution_cost;
         for account_key in tx_cost.writable_accounts.iter() {
             *self

ledger-tool/src/main.rs

@@ -794,7 +794,7 @@ fn compute_slot_cost(blockstore: &Blockstore, slot: Slot) -> Result<(), String>
 
                 let tx_cost = cost_model.calculate_cost(
                     &transaction,
-                    true, // demote_program_write_locks
+                    true, // demote_program_write_lock_features
                 );
                 if cost_tracker.try_add(tx_cost).is_err() {
                     println!(

program-runtime/src/instruction_processor.rs

@@ -4,7 +4,9 @@ use solana_sdk::{
     account::{AccountSharedData, ReadableAccount, WritableAccount},
     account_utils::StateMut,
     bpf_loader_upgradeable::{self, UpgradeableLoaderState},
-    feature_set::{demote_program_write_locks, fix_write_privs},
+    feature_set::{
+        demote_program_write_locks, fix_write_privs, restore_write_lock_when_upgradeable,
+    },
     ic_msg,
     instruction::{Instruction, InstructionError},
     message::Message,
@@ -599,10 +601,11 @@ impl InstructionProcessor {
         );
         if result.is_ok() {
             // Verify the called program has not misbehaved
-            let demote_program_write_locks =
-                invoke_context.is_feature_active(&demote_program_write_locks::id());
+            let demote_program_write_lock_features = invoke_context
+                .is_feature_active(&demote_program_write_locks::id())
+                && invoke_context.is_feature_active(&restore_write_lock_when_upgradeable::id());
             let write_privileges: Vec<bool> = (0..message.account_keys.len())
-                .map(|i| message.is_writable(i, demote_program_write_locks))
+                .map(|i| message.is_writable(i, demote_program_write_lock_features))
                 .collect();
             result =
                 invoke_context.verify_and_update(instruction, account_indices, &write_privileges);

program-test/src/lib.rs

@@ -24,7 +24,7 @@ use {
         compute_budget::ComputeBudget,
         entrypoint::{ProgramResult, SUCCESS},
         epoch_schedule::EpochSchedule,
-        feature_set::demote_program_write_locks,
+        feature_set::{demote_program_write_locks, restore_write_lock_when_upgradeable},
         fee_calculator::{FeeCalculator, FeeRateGovernor},
         genesis_config::{ClusterType, GenesisConfig},
         hash::Hash,
@@ -262,14 +262,15 @@ impl solana_sdk::program_stubs::SyscallStubs for SyscallStubs {
         let message = Message::new(&[instruction.clone()], None);
         let program_id_index = message.instructions[0].program_id_index as usize;
         let program_id = message.account_keys[program_id_index];
-        let demote_program_write_locks =
-            invoke_context.is_feature_active(&demote_program_write_locks::id());
+        let demote_program_write_lock_features = invoke_context
+            .is_feature_active(&demote_program_write_locks::id())
+            && 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, demote_program_write_locks))
+            .map(|(i, _)| message.is_writable(i, demote_program_write_lock_features))
             .collect::<Vec<bool>>();
 
         stable_log::program_invoke(&logger, &program_id, invoke_context.invoke_depth());
@@ -295,7 +296,7 @@ impl solana_sdk::program_stubs::SyscallStubs for SyscallStubs {
                 account.set_executable(account_info.executable);
                 account.set_rent_epoch(account_info.rent_epoch);
             }
-            let account_info = if message.is_writable(i, demote_program_write_locks) {
+            let account_info = if message.is_writable(i, demote_program_write_lock_features) {
                 Some(account_info)
             } else {
                 None

programs/bpf/tests/programs.rs

@@ -1862,7 +1862,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(),
@@ -1881,12 +1881,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)
     );
 }
 
@@ -2182,6 +2180,96 @@ fn test_program_bpf_upgrade_via_cpi() {
     assert_ne!(programdata, original_programdata);
 }
 
+#[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_for_tests(&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),
+        ],
+    );
+
+    // 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

@@ -21,8 +21,9 @@ use solana_sdk::{
     feature_set::{
         allow_native_ids, blake3_syscall_enabled, check_seed_length,
         close_upgradeable_program_accounts, demote_program_write_locks, disable_fees_sysvar,
-        libsecp256k1_0_5_upgrade_enabled, mem_overlap_fix, return_data_syscall_enabled,
-        secp256k1_recover_syscall_enabled, sol_log_data_syscall_enabled,
+        libsecp256k1_0_5_upgrade_enabled, mem_overlap_fix, restore_write_lock_when_upgradeable,
+        return_data_syscall_enabled, secp256k1_recover_syscall_enabled,
+        sol_log_data_syscall_enabled,
     },
     hash::{Hasher, HASH_BYTES},
     ic_msg,
@@ -2070,8 +2071,9 @@ fn get_translated_accounts<'a, T, F>(
 where
     F: Fn(&T, &mut dyn InvokeContext) -> Result<AccountReferences<'a>, EbpfError<BpfError>>,
 {
-    let demote_program_write_locks =
-        invoke_context.is_feature_active(&demote_program_write_locks::id());
+    let demote_program_write_lock_features = invoke_context
+        .is_feature_active(&demote_program_write_locks::id())
+        && invoke_context.is_feature_active(&restore_write_lock_when_upgradeable::id());
     let mut account_indices = Vec::with_capacity(message.account_keys.len());
     let mut accounts = Vec::with_capacity(message.account_keys.len());
     for (i, account_key) in message.account_keys.iter().enumerate() {
@@ -2095,7 +2097,7 @@ where
                     account.set_executable(account_ref.executable);
                     account.set_rent_epoch(account_ref.rent_epoch);
                 }
-                let account_ref = if message.is_writable(i, demote_program_write_locks) {
+                let account_ref = if message.is_writable(i, demote_program_write_lock_features) {
                     Some(account_ref)
                 } else {
                     None

rpc/src/transaction_status_service.rs

@@ -113,8 +113,8 @@ impl TransactionStatusService {
                             })
                             .expect("FeeCalculator must exist");
                         let fee = transaction.message().calculate_fee(&fee_calculator);
-                        let tx_account_locks =
-                            transaction.get_account_locks(bank.demote_program_write_locks());
+                        let tx_account_locks = transaction
+                            .get_account_locks(bank.demote_program_write_lock_features());
 
                         let inner_instructions = inner_instructions.map(|inner_instructions| {
                             inner_instructions