Add runtime support for address table lookups (#22223)

* Add support for address table lookups in runtime

* feedback

* feedback

accountsdb-plugin-postgres/scripts/create_schema.sql

@@ -49,7 +49,11 @@ Create TYPE "TransactionErrorCode" AS ENUM (
     'UnsupportedVersion',
     'InvalidWritableAccount',
     'WouldExceedMaxAccountDataCostLimit',
-    'TooManyAccountLocks'
+    'TooManyAccountLocks',
+    'AddressLookupTableNotFound',
+    'InvalidAddressLookupTableOwner',
+    'InvalidAddressLookupTableData',
+    'InvalidAddressLookupTableIndex'
 );
 
 CREATE TYPE "TransactionError" AS (

accountsdb-plugin-postgres/src/postgres_client/postgres_client_transaction.rs

@@ -332,6 +332,10 @@ pub enum DbTransactionErrorCode {
     InvalidWritableAccount,
     WouldExceedMaxAccountDataCostLimit,
     TooManyAccountLocks,
+    AddressLookupTableNotFound,
+    InvalidAddressLookupTableOwner,
+    InvalidAddressLookupTableData,
+    InvalidAddressLookupTableIndex,
 }
 
 impl From<&TransactionError> for DbTransactionErrorCode {
@@ -364,6 +368,14 @@ impl From<&TransactionError> for DbTransactionErrorCode {
                 Self::WouldExceedMaxAccountDataCostLimit
             }
             TransactionError::TooManyAccountLocks => Self::TooManyAccountLocks,
+            TransactionError::AddressLookupTableNotFound => Self::AddressLookupTableNotFound,
+            TransactionError::InvalidAddressLookupTableOwner => {
+                Self::InvalidAddressLookupTableOwner
+            }
+            TransactionError::InvalidAddressLookupTableData => Self::InvalidAddressLookupTableData,
+            TransactionError::InvalidAddressLookupTableIndex => {
+                Self::InvalidAddressLookupTableIndex
+            }
         }
     }
 }

core/src/banking_stage.rs

@@ -33,7 +33,10 @@ use {
             MAX_TRANSACTION_FORWARDING_DELAY_GPU,
         },
         feature_set,
-        message::Message,
+        message::{
+            v0::{LoadedAddresses, MessageAddressTableLookup},
+            Message,
+        },
         pubkey::Pubkey,
         short_vec::decode_shortu16_len,
         signature::Signature,
@@ -1114,6 +1117,7 @@ impl BankingStage {
         transaction_indexes: &[usize],
         feature_set: &Arc<feature_set::FeatureSet>,
         votes_only: bool,
+        address_loader: impl Fn(&[MessageAddressTableLookup]) -> transaction::Result<LoadedAddresses>,
     ) -> (Vec<SanitizedTransaction>, Vec<usize>) {
         transaction_indexes
             .iter()
@@ -1130,7 +1134,7 @@ impl BankingStage {
                     tx,
                     message_hash,
                     Some(p.meta.is_simple_vote_tx()),
-                    |_| Err(TransactionError::UnsupportedVersion),
+                    &address_loader,
                 )
                 .ok()?;
                 tx.verify_precompiles(feature_set).ok()?;
@@ -1196,6 +1200,7 @@ impl BankingStage {
             &packet_indexes,
             &bank.feature_set,
             bank.vote_only_bank(),
+            |lookup| bank.load_lookup_table_addresses(lookup),
         );
         packet_conversion_time.stop();
         inc_new_counter_info!("banking_stage-packet_conversion", 1);
@@ -1270,6 +1275,7 @@ impl BankingStage {
             transaction_indexes,
             &bank.feature_set,
             bank.vote_only_bank(),
+            |lookup| bank.load_lookup_table_addresses(lookup),
         );
         unprocessed_packet_conversion_time.stop();
 
@@ -3109,6 +3115,7 @@ mod tests {
                 &packet_indexes,
                 &Arc::new(FeatureSet::default()),
                 votes_only,
+                |_| Err(TransactionError::UnsupportedVersion),
             );
             assert_eq!(2, txs.len());
             assert_eq!(vec![0, 1], tx_packet_index);
@@ -3119,6 +3126,7 @@ mod tests {
                 &packet_indexes,
                 &Arc::new(FeatureSet::default()),
                 votes_only,
+                |_| Err(TransactionError::UnsupportedVersion),
             );
             assert_eq!(0, txs.len());
             assert_eq!(0, tx_packet_index.len());
@@ -3138,6 +3146,7 @@ mod tests {
                 &packet_indexes,
                 &Arc::new(FeatureSet::default()),
                 votes_only,
+                |_| Err(TransactionError::UnsupportedVersion),
             );
             assert_eq!(3, txs.len());
             assert_eq!(vec![0, 1, 2], tx_packet_index);
@@ -3148,6 +3157,7 @@ mod tests {
                 &packet_indexes,
                 &Arc::new(FeatureSet::default()),
                 votes_only,
+                |_| Err(TransactionError::UnsupportedVersion),
             );
             assert_eq!(2, txs.len());
             assert_eq!(vec![0, 2], tx_packet_index);
@@ -3167,6 +3177,7 @@ mod tests {
                 &packet_indexes,
                 &Arc::new(FeatureSet::default()),
                 votes_only,
+                |_| Err(TransactionError::UnsupportedVersion),
             );
             assert_eq!(3, txs.len());
             assert_eq!(vec![0, 1, 2], tx_packet_index);
@@ -3177,6 +3188,7 @@ mod tests {
                 &packet_indexes,
                 &Arc::new(FeatureSet::default()),
                 votes_only,
+                |_| Err(TransactionError::UnsupportedVersion),
             );
             assert_eq!(3, txs.len());
             assert_eq!(vec![0, 1, 2], tx_packet_index);

ledger/src/blockstore/blockstore_purge.rs

@@ -335,8 +335,8 @@ impl Blockstore {
                 if let Some(&signature) = transaction.signatures.get(0) {
                     batch.delete::<cf::TransactionStatus>((0, signature, slot))?;
                     batch.delete::<cf::TransactionStatus>((1, signature, slot))?;
-                    // TODO: support purging mapped addresses from versioned transactions
-                    for pubkey in transaction.message.unmapped_keys() {
+                    // TODO: support purging dynamically loaded addresses from versioned transactions
+                    for pubkey in transaction.message.into_static_account_keys() {
                         batch.delete::<cf::AddressSignatures>((0, pubkey, slot, signature))?;
                         batch.delete::<cf::AddressSignatures>((1, pubkey, slot, signature))?;
                     }

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),
+        );
+    }
 }