Migrate from address maps to address lookup tables (#21634)

* Migrate from address maps to address lookup tables

* update sanitize error

* cargo fmt

* update abi

sdk/program/src/message/versions/v0/loaded.rs

@@ -0,0 +1,308 @@
+use {
+    crate::{
+        bpf_loader_upgradeable,
+        message::{legacy::BUILTIN_PROGRAMS_KEYS, v0},
+        pubkey::Pubkey,
+        sysvar,
+    },
+    std::{collections::HashSet, ops::Deref, convert::TryFrom},
+};
+
+/// Combination of a version #0 message and its loaded addresses
+#[derive(Debug, Clone)]
+pub struct LoadedMessage {
+    /// Message which loaded a collection of lookup table addresses
+    pub message: v0::Message,
+    /// Addresses loaded with on-chain address lookup tables
+    pub loaded_addresses: LoadedAddresses,
+}
+
+impl Deref for LoadedMessage {
+    type Target = v0::Message;
+    fn deref(&self) -> &Self::Target {
+        &self.message
+    }
+}
+
+/// Collection of addresses loaded from on-chain lookup tables, split
+/// by readonly and writable.
+#[derive(Clone, Default, Debug, PartialEq, Serialize, Deserialize)]
+pub struct LoadedAddresses {
+    /// List of addresses for writable loaded accounts
+    pub writable: Vec<Pubkey>,
+    /// List of addresses for read-only loaded accounts
+    pub readonly: Vec<Pubkey>,
+}
+
+impl LoadedMessage {
+    /// 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.loaded_addresses.writable,
+            &self.loaded_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.loaded_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::*,
+        crate::{instruction::CompiledInstruction, message::MessageHeader, system_program, sysvar},
+        itertools::Itertools,
+    };
+
+    fn check_test_loaded_message() -> (LoadedMessage, [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 = LoadedMessage {
+            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()
+            },
+            loaded_addresses: LoadedAddresses {
+                writable: vec![key4],
+                readonly: vec![key5],
+            },
+        };
+
+        (message, [key0, key1, key2, key3, key4, key5])
+    }
+
+    #[test]
+    fn test_account_keys_segment_iter() {
+        let (message, keys) = check_test_loaded_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) = check_test_loaded_message();
+
+        assert_eq!(message.account_keys_len(), keys.len());
+    }
+
+    #[test]
+    fn test_account_keys_iter() {
+        let (message, keys) = check_test_loaded_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 = check_test_loaded_message().0;
+
+        assert!(!message.has_duplicates());
+    }
+
+    #[test]
+    fn test_has_duplicates_with_dupe_keys() {
+        let create_message_with_dupe_keys = |mut keys: Vec<Pubkey>| LoadedMessage {
+            message: v0::Message {
+                account_keys: keys.split_off(2),
+                ..v0::Message::default()
+            },
+            loaded_addresses: LoadedAddresses {
+                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) = check_test_loaded_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 = check_test_loaded_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 message = check_test_loaded_message().0;
+
+        message.message.account_keys[0] = sysvar::clock::id();
+        assert!(message.is_writable_index(0));
+        assert!(!message.is_writable(0, /*demote_program_write_locks=*/ true));
+
+        message.message.account_keys[0] = system_program::id();
+        assert!(message.is_writable_index(0));
+        assert!(!message.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 message = LoadedMessage {
+            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()
+            },
+            loaded_addresses: LoadedAddresses {
+                writable: vec![key1, key2],
+                readonly: vec![],
+            },
+        };
+
+        assert!(message.is_writable_index(2));
+        assert!(!message.is_writable(2, /*demote_program_write_locks=*/ true));
+    }
+}