Add runtime support for address table lookups (backport #22223) (#22354)

programs/address-lookup-table/src/state.rs

@@ -6,6 +6,7 @@ use {
         instruction::InstructionError,
         pubkey::Pubkey,
         slot_hashes::{SlotHashes, MAX_ENTRIES},
+        transaction::AddressLookupError,
     },
     std::borrow::Cow,
 };
@@ -133,6 +134,49 @@ impl<'a> AddressLookupTable<'a> {
         Ok(())
     }
 
+    /// Get the length of addresses that are active for lookups
+    pub fn get_active_addresses_len(
+        &self,
+        current_slot: Slot,
+        slot_hashes: &SlotHashes,
+    ) -> Result<usize, AddressLookupError> {
+        if !self.meta.is_active(current_slot, slot_hashes) {
+            // Once a lookup table is no longer active, it can be closed
+            // at any point, so returning a specific error for deactivated
+            // lookup tables could result in a race condition.
+            return Err(AddressLookupError::LookupTableAccountNotFound);
+        }
+
+        // If the address table was extended in the same slot in which it is used
+        // to lookup addresses for another transaction, the recently extended
+        // addresses are not considered active and won't be accessible.
+        let active_addresses_len = if current_slot > self.meta.last_extended_slot {
+            self.addresses.len()
+        } else {
+            self.meta.last_extended_slot_start_index as usize
+        };
+
+        Ok(active_addresses_len)
+    }
+
+    /// Lookup addresses for provided table indexes. Since lookups are performed on
+    /// tables which are not read-locked, this implementation needs to be careful
+    /// about resolving addresses consistently.
+    pub fn lookup(
+        &self,
+        current_slot: Slot,
+        indexes: &[u8],
+        slot_hashes: &SlotHashes,
+    ) -> Result<Vec<Pubkey>, AddressLookupError> {
+        let active_addresses_len = self.get_active_addresses_len(current_slot, slot_hashes)?;
+        let active_addresses = &self.addresses[0..active_addresses_len];
+        indexes
+            .iter()
+            .map(|idx| active_addresses.get(*idx as usize).cloned())
+            .collect::<Option<_>>()
+            .ok_or(AddressLookupError::InvalidLookupIndex)
+    }
+
     /// Serialize an address table including its addresses
     pub fn serialize_for_tests(self, data: &mut Vec<u8>) -> Result<(), InstructionError> {
         data.resize(LOOKUP_TABLE_META_SIZE, 0);
@@ -322,4 +366,117 @@ mod tests {
             test_case(case);
         }
     }
+
+    #[test]
+    fn test_lookup_from_empty_table() {
+        let lookup_table = AddressLookupTable {
+            meta: LookupTableMeta::default(),
+            addresses: Cow::Owned(vec![]),
+        };
+
+        assert_eq!(
+            lookup_table.lookup(0, &[], &SlotHashes::default()),
+            Ok(vec![])
+        );
+        assert_eq!(
+            lookup_table.lookup(0, &[0], &SlotHashes::default()),
+            Err(AddressLookupError::InvalidLookupIndex)
+        );
+    }
+
+    #[test]
+    fn test_lookup_from_deactivating_table() {
+        let current_slot = 1;
+        let slot_hashes = SlotHashes::default();
+        let addresses = vec![Pubkey::new_unique()];
+        let lookup_table = AddressLookupTable {
+            meta: LookupTableMeta {
+                deactivation_slot: current_slot,
+                last_extended_slot: current_slot - 1,
+                ..LookupTableMeta::default()
+            },
+            addresses: Cow::Owned(addresses.clone()),
+        };
+
+        assert_eq!(
+            lookup_table.meta.status(current_slot, &slot_hashes),
+            LookupTableStatus::Deactivating {
+                remaining_blocks: MAX_ENTRIES + 1
+            }
+        );
+
+        assert_eq!(
+            lookup_table.lookup(current_slot, &[0], &slot_hashes),
+            Ok(vec![addresses[0]]),
+        );
+    }
+
+    #[test]
+    fn test_lookup_from_deactivated_table() {
+        let current_slot = 1;
+        let slot_hashes = SlotHashes::default();
+        let lookup_table = AddressLookupTable {
+            meta: LookupTableMeta {
+                deactivation_slot: current_slot - 1,
+                last_extended_slot: current_slot - 1,
+                ..LookupTableMeta::default()
+            },
+            addresses: Cow::Owned(vec![]),
+        };
+
+        assert_eq!(
+            lookup_table.meta.status(current_slot, &slot_hashes),
+            LookupTableStatus::Deactivated
+        );
+        assert_eq!(
+            lookup_table.lookup(current_slot, &[0], &slot_hashes),
+            Err(AddressLookupError::LookupTableAccountNotFound)
+        );
+    }
+
+    #[test]
+    fn test_lookup_from_table_extended_in_current_slot() {
+        let current_slot = 0;
+        let addresses: Vec<_> = (0..2).map(|_| Pubkey::new_unique()).collect();
+        let lookup_table = AddressLookupTable {
+            meta: LookupTableMeta {
+                last_extended_slot: current_slot,
+                last_extended_slot_start_index: 1,
+                ..LookupTableMeta::default()
+            },
+            addresses: Cow::Owned(addresses.clone()),
+        };
+
+        assert_eq!(
+            lookup_table.lookup(current_slot, &[0], &SlotHashes::default()),
+            Ok(vec![addresses[0]])
+        );
+        assert_eq!(
+            lookup_table.lookup(current_slot, &[1], &SlotHashes::default()),
+            Err(AddressLookupError::InvalidLookupIndex),
+        );
+    }
+
+    #[test]
+    fn test_lookup_from_table_extended_in_previous_slot() {
+        let current_slot = 1;
+        let addresses: Vec<_> = (0..10).map(|_| Pubkey::new_unique()).collect();
+        let lookup_table = AddressLookupTable {
+            meta: LookupTableMeta {
+                last_extended_slot: current_slot - 1,
+                last_extended_slot_start_index: 1,
+                ..LookupTableMeta::default()
+            },
+            addresses: Cow::Owned(addresses.clone()),
+        };
+
+        assert_eq!(
+            lookup_table.lookup(current_slot, &[0, 3, 1, 5], &SlotHashes::default()),
+            Ok(vec![addresses[0], addresses[3], addresses[1], addresses[5]])
+        );
+        assert_eq!(
+            lookup_table.lookup(current_slot, &[10], &SlotHashes::default()),
+            Err(AddressLookupError::InvalidLookupIndex),
+        );
+    }
 }