Restore ability for programs to upgrade themselves (backport #20265) (#20295)

* Restore ability for programs to upgrade themselves (#20265)

* Make helper associated fn

* Add feature definition

* Add handling to preserve program-id write lock when upgradeable loader is present; restore bpf upgrade-self test

* Use single feature

(cherry picked from commit 2cd9dc99b6)

# Conflicts:
#	runtime/src/accounts.rs
#	sdk/program/src/message.rs
#	sdk/program/src/message/mapped.rs
#	sdk/program/src/message/sanitized.rs
#	sdk/src/feature_set.rs

* Fix conflicts

Co-authored-by: Tyera Eulberg <teulberg@gmail.com>
Co-authored-by: Tyera Eulberg <tyera@solana.com>

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)
     );
 }
 
@@ -2105,6 +2103,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(&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() {

runtime/src/accounts.rs

@@ -217,7 +217,6 @@ impl Accounts {
             let rent_for_sysvars = feature_set.is_active(&feature_set::rent_for_sysvars::id());
             let demote_program_write_locks =
                 feature_set.is_active(&feature_set::demote_program_write_locks::id());
-            let is_upgradeable_loader_present = is_upgradeable_loader_present(message);
 
             for (i, key) in message.account_keys.iter().enumerate() {
                 let account = if key_check.is_non_loader_key(key, i) {
@@ -250,7 +249,7 @@ impl Accounts {
                         if bpf_loader_upgradeable::check_id(account.owner()) {
                             if demote_program_write_locks
                                 && message.is_writable(i, demote_program_write_locks)
-                                && !is_upgradeable_loader_present
+                                && !message.is_upgradeable_loader_present()
                             {
                                 error_counters.invalid_writable_account += 1;
                                 return Err(TransactionError::InvalidWritableAccount);
@@ -1126,13 +1125,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>,

sdk/program/src/message.rs

@@ -402,6 +402,9 @@ impl Message {
     }
 
     pub fn is_writable(&self, i: usize, demote_program_write_locks: bool) -> bool {
+        let demote_program_id = demote_program_write_locks
+            && self.is_key_called_as_program(i)
+            && !self.is_upgradeable_loader_present();
         (i < (self.header.num_required_signatures - self.header.num_readonly_signed_accounts)
             as usize
             || (i >= self.header.num_required_signatures as usize
@@ -411,7 +414,7 @@ impl Message {
                 let key = self.account_keys[i];
                 sysvar::is_sysvar_id(&key) || BUILTIN_PROGRAMS_KEYS.contains(&key)
             }
-            && !(demote_program_write_locks && self.is_key_called_as_program(i))
+            && !demote_program_id
     }
 
     pub fn is_signer(&self, i: usize) -> bool {
@@ -537,6 +540,13 @@ impl Message {
             .min(self.header.num_required_signatures as usize);
         self.account_keys[..last_key].iter().collect()
     }
+
+    /// Returns true if any account is the bpf upgradeable loader
+    pub fn is_upgradeable_loader_present(&self) -> bool {
+        self.account_keys
+            .iter()
+            .any(|&key| key == bpf_loader_upgradeable::id())
+    }
 }
 
 #[cfg(test)]

sdk/program/src/message/mapped.rs

@@ -0,0 +1,300 @@
+use {
+    crate::{
+        bpf_loader_upgradeable,
+        message::{legacy::BUILTIN_PROGRAMS_KEYS, v0},
+        pubkey::Pubkey,
+        sysvar,
+    },
+    std::{collections::HashSet, convert::TryFrom},
+};
+
+/// Combination of a version #0 message and its mapped addresses
+#[derive(Debug, Clone)]
+pub struct MappedMessage {
+    /// Message which loaded a collection of mapped addresses
+    pub message: v0::Message,
+    /// Collection of mapped addresses loaded by this message
+    pub mapped_addresses: MappedAddresses,
+}
+
+/// Collection of mapped addresses loaded succinctly by a transaction using
+/// on-chain address map accounts.
+#[derive(Clone, Default, Debug, PartialEq, Serialize, Deserialize)]
+pub struct MappedAddresses {
+    /// List of addresses for writable loaded accounts
+    pub writable: Vec<Pubkey>,
+    /// List of addresses for read-only loaded accounts
+    pub readonly: Vec<Pubkey>,
+}
+
+impl MappedMessage {
+    /// Returns an iterator of account key segments. The ordering of segments
+    /// affects how account indexes from compiled instructions are resolved and
+    /// so should not be changed.
+    fn account_keys_segment_iter(&self) -> impl Iterator<Item = &Vec<Pubkey>> {
+        vec![
+            &self.message.account_keys,
+            &self.mapped_addresses.writable,
+            &self.mapped_addresses.readonly,
+        ]
+        .into_iter()
+    }
+
+    /// Returns the total length of loaded accounts for this message
+    pub fn account_keys_len(&self) -> usize {
+        let mut len = 0usize;
+        for key_segment in self.account_keys_segment_iter() {
+            len = len.saturating_add(key_segment.len());
+        }
+        len
+    }
+
+    /// Iterator for the addresses of the loaded accounts for this message
+    pub fn account_keys_iter(&self) -> impl Iterator<Item = &Pubkey> {
+        self.account_keys_segment_iter().flatten()
+    }
+
+    /// Returns true if any account keys are duplicates
+    pub fn has_duplicates(&self) -> bool {
+        let mut uniq = HashSet::new();
+        self.account_keys_iter().any(|x| !uniq.insert(x))
+    }
+
+    /// Returns the address of the account at the specified index of the list of
+    /// message account keys constructed from unmapped keys, followed by mapped
+    /// writable addresses, and lastly the list of mapped readonly addresses.
+    pub fn get_account_key(&self, mut index: usize) -> Option<&Pubkey> {
+        for key_segment in self.account_keys_segment_iter() {
+            if index < key_segment.len() {
+                return Some(&key_segment[index]);
+            }
+            index = index.saturating_sub(key_segment.len());
+        }
+
+        None
+    }
+
+    /// Returns true if the account at the specified index was requested to be
+    /// writable.  This method should not be used directly.
+    fn is_writable_index(&self, key_index: usize) -> bool {
+        let header = &self.message.header;
+        let num_account_keys = self.message.account_keys.len();
+        let num_signed_accounts = usize::from(header.num_required_signatures);
+        if key_index >= num_account_keys {
+            let mapped_addresses_index = key_index.saturating_sub(num_account_keys);
+            mapped_addresses_index < self.mapped_addresses.writable.len()
+        } else if key_index >= num_signed_accounts {
+            let num_unsigned_accounts = num_account_keys.saturating_sub(num_signed_accounts);
+            let num_writable_unsigned_accounts = num_unsigned_accounts
+                .saturating_sub(usize::from(header.num_readonly_unsigned_accounts));
+            let unsigned_account_index = key_index.saturating_sub(num_signed_accounts);
+            unsigned_account_index < num_writable_unsigned_accounts
+        } else {
+            let num_writable_signed_accounts = num_signed_accounts
+                .saturating_sub(usize::from(header.num_readonly_signed_accounts));
+            key_index < num_writable_signed_accounts
+        }
+    }
+
+    /// Returns true if the account at the specified index was loaded as writable
+    pub fn is_writable(&self, key_index: usize, demote_program_write_locks: bool) -> bool {
+        if self.is_writable_index(key_index) {
+            if let Some(key) = self.get_account_key(key_index) {
+                let demote_program_id = demote_program_write_locks
+                    && self.is_key_called_as_program(key_index)
+                    && !self.is_upgradeable_loader_present();
+                return !(sysvar::is_sysvar_id(key)
+                    || BUILTIN_PROGRAMS_KEYS.contains(key)
+                    || demote_program_id);
+            }
+        }
+        false
+    }
+
+    /// Returns true if the account at the specified index is called as a program by an instruction
+    pub fn is_key_called_as_program(&self, key_index: usize) -> bool {
+        if let Ok(key_index) = u8::try_from(key_index) {
+            self.message.instructions
+                .iter()
+                .any(|ix| ix.program_id_index == key_index)
+        } else {
+            false
+        }
+    }
+
+    /// Returns true if any account is the bpf upgradeable loader
+    pub fn is_upgradeable_loader_present(&self) -> bool {
+        self.account_keys_iter()
+            .any(|&key| key == bpf_loader_upgradeable::id())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::{instruction::CompiledInstruction, message::MessageHeader, system_program, sysvar};
+    use itertools::Itertools;
+
+    fn create_test_mapped_message() -> (MappedMessage, [Pubkey; 6]) {
+        let key0 = Pubkey::new_unique();
+        let key1 = Pubkey::new_unique();
+        let key2 = Pubkey::new_unique();
+        let key3 = Pubkey::new_unique();
+        let key4 = Pubkey::new_unique();
+        let key5 = Pubkey::new_unique();
+
+        let message = MappedMessage {
+            message: v0::Message {
+                header: MessageHeader {
+                    num_required_signatures: 2,
+                    num_readonly_signed_accounts: 1,
+                    num_readonly_unsigned_accounts: 1,
+                },
+                account_keys: vec![key0, key1, key2, key3],
+                ..v0::Message::default()
+            },
+            mapped_addresses: MappedAddresses {
+                writable: vec![key4],
+                readonly: vec![key5],
+            },
+        };
+
+        (message, [key0, key1, key2, key3, key4, key5])
+    }
+
+    #[test]
+    fn test_account_keys_segment_iter() {
+        let (message, keys) = create_test_mapped_message();
+
+        let expected_segments = vec![
+            vec![keys[0], keys[1], keys[2], keys[3]],
+            vec![keys[4]],
+            vec![keys[5]],
+        ];
+
+        let mut iter = message.account_keys_segment_iter();
+        for expected_segment in expected_segments {
+            assert_eq!(iter.next(), Some(&expected_segment));
+        }
+    }
+
+    #[test]
+    fn test_account_keys_len() {
+        let (message, keys) = create_test_mapped_message();
+
+        assert_eq!(message.account_keys_len(), keys.len());
+    }
+
+    #[test]
+    fn test_account_keys_iter() {
+        let (message, keys) = create_test_mapped_message();
+
+        let mut iter = message.account_keys_iter();
+        for expected_key in keys {
+            assert_eq!(iter.next(), Some(&expected_key));
+        }
+    }
+
+    #[test]
+    fn test_has_duplicates() {
+        let message = create_test_mapped_message().0;
+
+        assert!(!message.has_duplicates());
+    }
+
+    #[test]
+    fn test_has_duplicates_with_dupe_keys() {
+        let create_message_with_dupe_keys = |mut keys: Vec<Pubkey>| MappedMessage {
+            message: v0::Message {
+                account_keys: keys.split_off(2),
+                ..v0::Message::default()
+            },
+            mapped_addresses: MappedAddresses {
+                writable: keys.split_off(2),
+                readonly: keys,
+            },
+        };
+
+        let key0 = Pubkey::new_unique();
+        let key1 = Pubkey::new_unique();
+        let key2 = Pubkey::new_unique();
+        let key3 = Pubkey::new_unique();
+        let dupe_key = Pubkey::new_unique();
+
+        let keys = vec![key0, key1, key2, key3, dupe_key, dupe_key];
+        let keys_len = keys.len();
+        for keys in keys.into_iter().permutations(keys_len).unique() {
+            let message = create_message_with_dupe_keys(keys);
+            assert!(message.has_duplicates());
+        }
+    }
+
+    #[test]
+    fn test_get_account_key() {
+        let (message, keys) = create_test_mapped_message();
+
+        assert_eq!(message.get_account_key(0), Some(&keys[0]));
+        assert_eq!(message.get_account_key(1), Some(&keys[1]));
+        assert_eq!(message.get_account_key(2), Some(&keys[2]));
+        assert_eq!(message.get_account_key(3), Some(&keys[3]));
+        assert_eq!(message.get_account_key(4), Some(&keys[4]));
+        assert_eq!(message.get_account_key(5), Some(&keys[5]));
+    }
+
+    #[test]
+    fn test_is_writable_index() {
+        let message = create_test_mapped_message().0;
+
+        assert!(message.is_writable_index(0));
+        assert!(!message.is_writable_index(1));
+        assert!(message.is_writable_index(2));
+        assert!(!message.is_writable_index(3));
+        assert!(message.is_writable_index(4));
+        assert!(!message.is_writable_index(5));
+    }
+
+    #[test]
+    fn test_is_writable() {
+        let mut mapped_msg = create_test_mapped_message().0;
+
+        mapped_msg.message.account_keys[0] = sysvar::clock::id();
+        assert!(mapped_msg.is_writable_index(0));
+        assert!(!mapped_msg.is_writable(0, /*demote_program_write_locks=*/ true));
+
+        mapped_msg.message.account_keys[0] = system_program::id();
+        assert!(mapped_msg.is_writable_index(0));
+        assert!(!mapped_msg.is_writable(0, /*demote_program_write_locks=*/ true));
+    }
+
+    #[test]
+    fn test_demote_writable_program() {
+        let key0 = Pubkey::new_unique();
+        let key1 = Pubkey::new_unique();
+        let key2 = Pubkey::new_unique();
+        let mapped_msg = MappedMessage {
+            message: v0::Message {
+                header: MessageHeader {
+                    num_required_signatures: 1,
+                    num_readonly_signed_accounts: 0,
+                    num_readonly_unsigned_accounts: 0,
+                },
+                account_keys: vec![key0],
+                instructions: vec![
+                    CompiledInstruction {
+                        program_id_index: 2,
+                        accounts: vec![1],
+                        data: vec![],
+                    }
+                ],
+                ..v0::Message::default()
+            },
+            mapped_addresses: MappedAddresses {
+                writable: vec![key1, key2],
+                readonly: vec![],
+            },
+        };
+
+        assert!(mapped_msg.is_writable_index(2));
+        assert!(!mapped_msg.is_writable(2, /*demote_program_write_locks=*/ true));
+    }
+}

sdk/program/src/message/sanitized.rs

@@ -0,0 +1,605 @@
+use {
+    crate::{
+        fee_calculator::FeeCalculator,
+        hash::Hash,
+        instruction::{CompiledInstruction, Instruction},
+        message::{MappedAddresses, MappedMessage, Message, MessageHeader},
+        pubkey::Pubkey,
+        sanitize::{Sanitize, SanitizeError},
+        secp256k1_program,
+        serialize_utils::{append_slice, append_u16, append_u8},
+    },
+    bitflags::bitflags,
+    std::convert::TryFrom,
+    thiserror::Error,
+};
+
+/// Sanitized message of a transaction which includes a set of atomic
+/// instructions to be executed on-chain
+#[derive(Debug, Clone)]
+pub enum SanitizedMessage {
+    /// Sanitized legacy message
+    Legacy(Message),
+    /// Sanitized version #0 message with mapped addresses
+    V0(MappedMessage),
+}
+
+#[derive(PartialEq, Debug, Error, Eq, Clone)]
+pub enum SanitizeMessageError {
+    #[error("index out of bounds")]
+    IndexOutOfBounds,
+    #[error("value out of bounds")]
+    ValueOutOfBounds,
+    #[error("invalid value")]
+    InvalidValue,
+    #[error("duplicate account key")]
+    DuplicateAccountKey,
+}
+
+impl From<SanitizeError> for SanitizeMessageError {
+    fn from(err: SanitizeError) -> Self {
+        match err {
+            SanitizeError::IndexOutOfBounds => Self::IndexOutOfBounds,
+            SanitizeError::ValueOutOfBounds => Self::ValueOutOfBounds,
+            SanitizeError::InvalidValue => Self::InvalidValue,
+        }
+    }
+}
+
+impl TryFrom<Message> for SanitizedMessage {
+    type Error = SanitizeMessageError;
+    fn try_from(message: Message) -> Result<Self, Self::Error> {
+        message.sanitize()?;
+
+        let sanitized_msg = Self::Legacy(message);
+        if sanitized_msg.has_duplicates() {
+            return Err(SanitizeMessageError::DuplicateAccountKey);
+        }
+
+        Ok(sanitized_msg)
+    }
+}
+
+bitflags! {
+    struct InstructionsSysvarAccountMeta: u8 {
+        const NONE = 0b00000000;
+        const IS_SIGNER = 0b00000001;
+        const IS_WRITABLE = 0b00000010;
+    }
+}
+
+impl SanitizedMessage {
+    /// Return true if this message contains duplicate account keys
+    pub fn has_duplicates(&self) -> bool {
+        match self {
+            SanitizedMessage::Legacy(message) => message.has_duplicates(),
+            SanitizedMessage::V0(message) => message.has_duplicates(),
+        }
+    }
+
+    /// Message header which identifies the number of signer and writable or
+    /// readonly accounts
+    pub fn header(&self) -> &MessageHeader {
+        match self {
+            Self::Legacy(message) => &message.header,
+            Self::V0(mapped_msg) => &mapped_msg.message.header,
+        }
+    }
+
+    /// Returns a legacy message if this sanitized message wraps one
+    pub fn legacy_message(&self) -> Option<&Message> {
+        if let Self::Legacy(message) = &self {
+            Some(message)
+        } else {
+            None
+        }
+    }
+
+    /// Returns the fee payer for the transaction
+    pub fn fee_payer(&self) -> &Pubkey {
+        self.get_account_key(0)
+            .expect("sanitized message always has non-program fee payer at index 0")
+    }
+
+    /// The hash of a recent block, used for timing out a transaction
+    pub fn recent_blockhash(&self) -> &Hash {
+        match self {
+            Self::Legacy(message) => &message.recent_blockhash,
+            Self::V0(mapped_msg) => &mapped_msg.message.recent_blockhash,
+        }
+    }
+
+    /// Program instructions that will be executed in sequence and committed in
+    /// one atomic transaction if all succeed.
+    pub fn instructions(&self) -> &[CompiledInstruction] {
+        match self {
+            Self::Legacy(message) => &message.instructions,
+            Self::V0(mapped_msg) => &mapped_msg.message.instructions,
+        }
+    }
+
+    /// Program instructions iterator which includes each instruction's program
+    /// id.
+    pub fn program_instructions_iter(
+        &self,
+    ) -> impl Iterator<Item = (&Pubkey, &CompiledInstruction)> {
+        match self {
+            Self::Legacy(message) => message.instructions.iter(),
+            Self::V0(mapped_msg) => mapped_msg.message.instructions.iter(),
+        }
+        .map(move |ix| {
+            (
+                self.get_account_key(usize::from(ix.program_id_index))
+                    .expect("program id index is sanitized"),
+                ix,
+            )
+        })
+    }
+
+    /// Iterator of all account keys referenced in this message, included mapped keys.
+    pub fn account_keys_iter(&self) -> Box<dyn Iterator<Item = &Pubkey> + '_> {
+        match self {
+            Self::Legacy(message) => Box::new(message.account_keys.iter()),
+            Self::V0(mapped_msg) => Box::new(mapped_msg.account_keys_iter()),
+        }
+    }
+
+    /// Length of all account keys referenced in this message, included mapped keys.
+    pub fn account_keys_len(&self) -> usize {
+        match self {
+            Self::Legacy(message) => message.account_keys.len(),
+            Self::V0(mapped_msg) => mapped_msg.account_keys_len(),
+        }
+    }
+
+    /// Returns the address of the account at the specified index.
+    pub fn get_account_key(&self, index: usize) -> Option<&Pubkey> {
+        match self {
+            Self::Legacy(message) => message.account_keys.get(index),
+            Self::V0(message) => message.get_account_key(index),
+        }
+    }
+
+    /// Returns true if the account at the specified index is an input to some
+    /// program instruction in this message.
+    fn is_key_passed_to_program(&self, key_index: usize) -> bool {
+        if let Ok(key_index) = u8::try_from(key_index) {
+            self.instructions()
+                .iter()
+                .any(|ix| ix.accounts.contains(&key_index))
+        } else {
+            false
+        }
+    }
+
+    /// Returns true if the account at the specified index is invoked as a
+    /// program in this message.
+    pub fn is_invoked(&self, key_index: usize) -> bool {
+        match self {
+            Self::Legacy(message) => message.is_key_called_as_program(key_index),
+            Self::V0(message) => message.is_key_called_as_program(key_index),
+        }
+    }
+
+    /// Returns true if the account at the specified index is not invoked as a
+    /// program or, if invoked, is passed to a program.
+    pub fn is_non_loader_key(&self, key_index: usize) -> bool {
+        !self.is_invoked(key_index) || self.is_key_passed_to_program(key_index)
+    }
+
+    /// Returns true if the account at the specified index is writable by the
+    /// instructions in this message.
+    pub fn is_writable(&self, index: usize, demote_program_write_locks: bool) -> bool {
+        match self {
+            Self::Legacy(message) => message.is_writable(index, demote_program_write_locks),
+            Self::V0(message) => message.is_writable(index, demote_program_write_locks),
+        }
+    }
+
+    /// Returns true if the account at the specified index signed this
+    /// message.
+    pub fn is_signer(&self, index: usize) -> bool {
+        index < usize::from(self.header().num_required_signatures)
+    }
+
+    // First encode the number of instructions:
+    // [0..2 - num_instructions
+    //
+    // Then a table of offsets of where to find them in the data
+    //  3..2 * num_instructions table of instruction offsets
+    //
+    // Each instruction is then encoded as:
+    //   0..2 - num_accounts
+    //   2 - meta_byte -> (bit 0 signer, bit 1 is_writable)
+    //   3..35 - pubkey - 32 bytes
+    //   35..67 - program_id
+    //   67..69 - data len - u16
+    //   69..data_len - data
+    #[allow(clippy::integer_arithmetic)]
+    pub fn serialize_instructions(&self, demote_program_write_locks: 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);
+        for _ in 0..self.instructions().len() {
+            append_u16(&mut data, 0);
+        }
+        for (i, (program_id, instruction)) in self.program_instructions_iter().enumerate() {
+            let start_instruction_offset = data.len() as u16;
+            let start = 2 + (2 * i);
+            data[start..start + 2].copy_from_slice(&start_instruction_offset.to_le_bytes());
+            append_u16(&mut data, instruction.accounts.len() as u16);
+            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, demote_program_write_locks);
+                let mut account_meta = InstructionsSysvarAccountMeta::NONE;
+                if is_signer {
+                    account_meta |= InstructionsSysvarAccountMeta::IS_SIGNER;
+                }
+                if is_writable {
+                    account_meta |= InstructionsSysvarAccountMeta::IS_WRITABLE;
+                }
+                append_u8(&mut data, account_meta.bits());
+                append_slice(
+                    &mut data,
+                    self.get_account_key(account_index).unwrap().as_ref(),
+                );
+            }
+
+            append_slice(&mut data, program_id.as_ref());
+            append_u16(&mut data, instruction.data.len() as u16);
+            append_slice(&mut data, &instruction.data);
+        }
+        data
+    }
+
+    /// Return the mapped addresses for this message if it has any.
+    fn mapped_addresses(&self) -> Option<&MappedAddresses> {
+        match &self {
+            SanitizedMessage::V0(message) => Some(&message.mapped_addresses),
+            _ => None,
+        }
+    }
+
+    /// Return the number of readonly accounts loaded by this message.
+    pub fn num_readonly_accounts(&self) -> usize {
+        let mapped_readonly_addresses = self
+            .mapped_addresses()
+            .map(|keys| keys.readonly.len())
+            .unwrap_or_default();
+        mapped_readonly_addresses
+            .saturating_add(usize::from(self.header().num_readonly_signed_accounts))
+            .saturating_add(usize::from(self.header().num_readonly_unsigned_accounts))
+    }
+
+    fn try_position(&self, key: &Pubkey) -> Option<u8> {
+        u8::try_from(self.account_keys_iter().position(|k| k == key)?).ok()
+    }
+
+    /// Try to compile an instruction using the account keys in this message.
+    pub fn try_compile_instruction(&self, ix: &Instruction) -> Option<CompiledInstruction> {
+        let accounts: Vec<_> = ix
+            .accounts
+            .iter()
+            .map(|account_meta| self.try_position(&account_meta.pubkey))
+            .collect::<Option<_>>()?;
+
+        Some(CompiledInstruction {
+            program_id_index: self.try_position(&ix.program_id)?,
+            data: ix.data.clone(),
+            accounts,
+        })
+    }
+
+    /// Calculate the total fees for a transaction given a fee calculator
+    pub fn calculate_fee(&self, fee_calculator: &FeeCalculator) -> u64 {
+        let mut num_secp256k1_signatures: u64 = 0;
+        for (program_id, instruction) in self.program_instructions_iter() {
+            if secp256k1_program::check_id(program_id) {
+                if let Some(num_signatures) = instruction.data.get(0) {
+                    num_secp256k1_signatures =
+                        num_secp256k1_signatures.saturating_add(u64::from(*num_signatures));
+                }
+            }
+        }
+
+        fee_calculator.lamports_per_signature.saturating_mul(
+            u64::from(self.header().num_required_signatures)
+                .saturating_add(num_secp256k1_signatures),
+        )
+    }
+
+    /// Inspect all message keys for the bpf upgradeable loader
+    pub fn is_upgradeable_loader_present(&self) -> bool {
+        match self {
+            Self::Legacy(message) => message.is_upgradeable_loader_present(),
+            Self::V0(message) => message.is_upgradeable_loader_present(),
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::{
+        instruction::{AccountMeta, Instruction},
+        message::v0,
+        secp256k1_program, system_instruction,
+    };
+
+    #[test]
+    fn test_try_from_message() {
+        let dupe_key = Pubkey::new_unique();
+        let legacy_message_with_dupes = Message {
+            header: MessageHeader {
+                num_required_signatures: 1,
+                ..MessageHeader::default()
+            },
+            account_keys: vec![dupe_key, dupe_key],
+            ..Message::default()
+        };
+
+        assert_eq!(
+            SanitizedMessage::try_from(legacy_message_with_dupes).err(),
+            Some(SanitizeMessageError::DuplicateAccountKey),
+        );
+
+        let legacy_message_with_no_signers = Message {
+            account_keys: vec![Pubkey::new_unique()],
+            ..Message::default()
+        };
+
+        assert_eq!(
+            SanitizedMessage::try_from(legacy_message_with_no_signers).err(),
+            Some(SanitizeMessageError::IndexOutOfBounds),
+        );
+    }
+
+    #[test]
+    fn test_is_non_loader_key() {
+        let key0 = Pubkey::new_unique();
+        let key1 = Pubkey::new_unique();
+        let loader_key = Pubkey::new_unique();
+        let instructions = vec![
+            CompiledInstruction::new(1, &(), vec![0]),
+            CompiledInstruction::new(2, &(), vec![0, 1]),
+        ];
+
+        let message = SanitizedMessage::try_from(Message::new_with_compiled_instructions(
+            1,
+            0,
+            2,
+            vec![key0, key1, loader_key],
+            Hash::default(),
+            instructions,
+        ))
+        .unwrap();
+
+        assert!(message.is_non_loader_key(0));
+        assert!(message.is_non_loader_key(1));
+        assert!(!message.is_non_loader_key(2));
+    }
+
+    #[test]
+    fn test_num_readonly_accounts() {
+        let key0 = Pubkey::new_unique();
+        let key1 = Pubkey::new_unique();
+        let key2 = Pubkey::new_unique();
+        let key3 = Pubkey::new_unique();
+        let key4 = Pubkey::new_unique();
+        let key5 = Pubkey::new_unique();
+
+        let legacy_message = SanitizedMessage::try_from(Message {
+            header: MessageHeader {
+                num_required_signatures: 2,
+                num_readonly_signed_accounts: 1,
+                num_readonly_unsigned_accounts: 1,
+            },
+            account_keys: vec![key0, key1, key2, key3],
+            ..Message::default()
+        })
+        .unwrap();
+
+        assert_eq!(legacy_message.num_readonly_accounts(), 2);
+
+        let mapped_message = SanitizedMessage::V0(MappedMessage {
+            message: v0::Message {
+                header: MessageHeader {
+                    num_required_signatures: 2,
+                    num_readonly_signed_accounts: 1,
+                    num_readonly_unsigned_accounts: 1,
+                },
+                account_keys: vec![key0, key1, key2, key3],
+                ..v0::Message::default()
+            },
+            mapped_addresses: MappedAddresses {
+                writable: vec![key4],
+                readonly: vec![key5],
+            },
+        });
+
+        assert_eq!(mapped_message.num_readonly_accounts(), 3);
+    }
+
+    #[test]
+    #[allow(deprecated)]
+    fn test_serialize_instructions() {
+        let program_id0 = Pubkey::new_unique();
+        let program_id1 = Pubkey::new_unique();
+        let id0 = Pubkey::new_unique();
+        let id1 = Pubkey::new_unique();
+        let id2 = Pubkey::new_unique();
+        let id3 = Pubkey::new_unique();
+        let instructions = vec![
+            Instruction::new_with_bincode(program_id0, &0, vec![AccountMeta::new(id0, false)]),
+            Instruction::new_with_bincode(program_id0, &0, vec![AccountMeta::new(id1, true)]),
+            Instruction::new_with_bincode(
+                program_id1,
+                &0,
+                vec![AccountMeta::new_readonly(id2, false)],
+            ),
+            Instruction::new_with_bincode(
+                program_id1,
+                &0,
+                vec![AccountMeta::new_readonly(id3, true)],
+            ),
+        ];
+
+        let demote_program_write_locks = true;
+        let message = Message::new(&instructions, Some(&id1));
+        let sanitized_message = SanitizedMessage::try_from(message.clone()).unwrap();
+        let serialized = sanitized_message.serialize_instructions(demote_program_write_locks);
+
+        // assert that SanitizedMessage::serialize_instructions has the same behavior as the
+        // deprecated Message::serialize_instructions method
+        assert_eq!(serialized, message.serialize_instructions());
+
+        // assert that Message::deserialize_instruction is compatible with SanitizedMessage::serialize_instructions
+        for (i, instruction) in instructions.iter().enumerate() {
+            assert_eq!(
+                Message::deserialize_instruction(i, &serialized).unwrap(),
+                *instruction
+            );
+        }
+    }
+
+    #[test]
+    fn test_calculate_fee() {
+        // Default: no fee.
+        let message =
+            SanitizedMessage::try_from(Message::new(&[], Some(&Pubkey::new_unique()))).unwrap();
+        assert_eq!(message.calculate_fee(&FeeCalculator::default()), 0);
+
+        // One signature, a fee.
+        assert_eq!(message.calculate_fee(&FeeCalculator::new(1)), 1);
+
+        // Two signatures, double the fee.
+        let key0 = Pubkey::new_unique();
+        let key1 = Pubkey::new_unique();
+        let ix0 = system_instruction::transfer(&key0, &key1, 1);
+        let ix1 = system_instruction::transfer(&key1, &key0, 1);
+        let message = SanitizedMessage::try_from(Message::new(&[ix0, ix1], Some(&key0))).unwrap();
+        assert_eq!(message.calculate_fee(&FeeCalculator::new(2)), 4);
+    }
+
+    #[test]
+    fn test_try_compile_instruction() {
+        let key0 = Pubkey::new_unique();
+        let key1 = Pubkey::new_unique();
+        let key2 = Pubkey::new_unique();
+        let program_id = Pubkey::new_unique();
+
+        let valid_instruction = Instruction {
+            program_id,
+            accounts: vec![
+                AccountMeta::new_readonly(key0, false),
+                AccountMeta::new_readonly(key1, false),
+                AccountMeta::new_readonly(key2, false),
+            ],
+            data: vec![],
+        };
+
+        let invalid_program_id_instruction = Instruction {
+            program_id: Pubkey::new_unique(),
+            accounts: vec![
+                AccountMeta::new_readonly(key0, false),
+                AccountMeta::new_readonly(key1, false),
+                AccountMeta::new_readonly(key2, false),
+            ],
+            data: vec![],
+        };
+
+        let invalid_account_key_instruction = Instruction {
+            program_id: Pubkey::new_unique(),
+            accounts: vec![
+                AccountMeta::new_readonly(key0, false),
+                AccountMeta::new_readonly(key1, false),
+                AccountMeta::new_readonly(Pubkey::new_unique(), false),
+            ],
+            data: vec![],
+        };
+
+        let legacy_message = SanitizedMessage::try_from(Message {
+            header: MessageHeader {
+                num_required_signatures: 1,
+                num_readonly_signed_accounts: 0,
+                num_readonly_unsigned_accounts: 0,
+            },
+            account_keys: vec![key0, key1, key2, program_id],
+            ..Message::default()
+        })
+        .unwrap();
+
+        let mapped_message = SanitizedMessage::V0(MappedMessage {
+            message: v0::Message {
+                header: MessageHeader {
+                    num_required_signatures: 1,
+                    num_readonly_signed_accounts: 0,
+                    num_readonly_unsigned_accounts: 0,
+                },
+                account_keys: vec![key0, key1],
+                ..v0::Message::default()
+            },
+            mapped_addresses: MappedAddresses {
+                writable: vec![key2],
+                readonly: vec![program_id],
+            },
+        });
+
+        for message in vec![legacy_message, mapped_message] {
+            assert_eq!(
+                message.try_compile_instruction(&valid_instruction),
+                Some(CompiledInstruction {
+                    program_id_index: 3,
+                    accounts: vec![0, 1, 2],
+                    data: vec![],
+                })
+            );
+
+            assert!(message
+                .try_compile_instruction(&invalid_program_id_instruction)
+                .is_none());
+            assert!(message
+                .try_compile_instruction(&invalid_account_key_instruction)
+                .is_none());
+        }
+    }
+
+    #[test]
+    fn test_calculate_fee_secp256k1() {
+        let key0 = Pubkey::new_unique();
+        let key1 = Pubkey::new_unique();
+        let ix0 = system_instruction::transfer(&key0, &key1, 1);
+
+        let mut secp_instruction1 = Instruction {
+            program_id: secp256k1_program::id(),
+            accounts: vec![],
+            data: vec![],
+        };
+        let mut secp_instruction2 = Instruction {
+            program_id: secp256k1_program::id(),
+            accounts: vec![],
+            data: vec![1],
+        };
+
+        let message = SanitizedMessage::try_from(Message::new(
+            &[
+                ix0.clone(),
+                secp_instruction1.clone(),
+                secp_instruction2.clone(),
+            ],
+            Some(&key0),
+        ))
+        .unwrap();
+        assert_eq!(message.calculate_fee(&FeeCalculator::new(1)), 2);
+
+        secp_instruction1.data = vec![0];
+        secp_instruction2.data = vec![10];
+        let message = SanitizedMessage::try_from(Message::new(
+            &[ix0, secp_instruction1, secp_instruction2],
+            Some(&key0),
+        ))
+        .unwrap();
+        assert_eq!(message.calculate_fee(&FeeCalculator::new(1)), 11);
+    }
+}

sdk/src/feature_set.rs

@@ -263,7 +263,7 @@ lazy_static! {
         (mem_overlap_fix::id(), "Memory overlap fix"),
         (close_upgradeable_program_accounts::id(), "enable closing upgradeable program accounts"),
         (stake_program_advance_activating_credits_observed::id(), "Enable advancing credits observed for activation epoch #19309"),
-        (demote_program_write_locks::id(), "demote program write locks to readonly #19593"),
+        (demote_program_write_locks::id(), "demote program write locks to readonly, except when upgradeable loader present #19593 #20265"),
         (allow_native_ids::id(), "allow native program ids in program derived addresses"),
         (check_seed_length::id(), "Check program address seed lengths"),
         (fix_write_privs::id(), "fix native invoke write privileges"),