AccountLocks: validate_account_locks (#2448)

accounts-db/src/account_locks.rs

@@ -2,8 +2,12 @@
 use qualifier_attr::qualifiers;
 use {
     ahash::{AHashMap, AHashSet},
-    solana_sdk::{pubkey::Pubkey, transaction::TransactionError},
-    std::collections::hash_map,
+    solana_sdk::{
+        message::AccountKeys,
+        pubkey::Pubkey,
+        transaction::{TransactionError, MAX_TX_ACCOUNT_LOCKS},
+    },
+    std::{cell::RefCell, collections::hash_map},
 };
 
 #[derive(Debug, Default)]
@@ -110,9 +114,51 @@ impl AccountLocks {
     }
 }
 
+/// Validate account locks before locking.
+pub fn validate_account_locks(
+    account_keys: AccountKeys,
+    tx_account_lock_limit: usize,
+) -> Result<(), TransactionError> {
+    if account_keys.len() > tx_account_lock_limit {
+        Err(TransactionError::TooManyAccountLocks)
+    } else if has_duplicates(account_keys) {
+        Err(TransactionError::AccountLoadedTwice)
+    } else {
+        Ok(())
+    }
+}
+
+thread_local! {
+    static HAS_DUPLICATES_SET: RefCell<AHashSet<Pubkey>> = RefCell::new(AHashSet::with_capacity(MAX_TX_ACCOUNT_LOCKS));
+}
+
+/// Check for duplicate account keys.
+fn has_duplicates(account_keys: AccountKeys) -> bool {
+    // Benchmarking has shown that for sets of 32 or more keys, it is faster to
+    // use a HashSet to check for duplicates.
+    // For smaller sets a brute-force O(n^2) check seems to be faster.
+    const USE_ACCOUNT_LOCK_SET_SIZE: usize = 32;
+    if account_keys.len() >= USE_ACCOUNT_LOCK_SET_SIZE {
+        HAS_DUPLICATES_SET.with_borrow_mut(|set| {
+            let has_duplicates = account_keys.iter().any(|key| !set.insert(*key));
+            set.clear();
+            has_duplicates
+        })
+    } else {
+        for (idx, key) in account_keys.iter().enumerate() {
+            for jdx in idx + 1..account_keys.len() {
+                if key == &account_keys[jdx] {
+                    return true;
+                }
+            }
+        }
+        false
+    }
+}
+
 #[cfg(test)]
 mod tests {
-    use super::*;
+    use {super::*, solana_sdk::message::v0::LoadedAddresses};
 
     #[test]
     fn test_account_locks() {
@@ -152,4 +198,94 @@ mod tests {
         account_locks.unlock_accounts([(&key2, false)].into_iter());
         assert!(!account_locks.is_locked_readonly(&key2));
     }
+
+    #[test]
+    fn test_validate_account_locks_valid_no_dynamic() {
+        let static_keys = &[Pubkey::new_unique(), Pubkey::new_unique()];
+        let account_keys = AccountKeys::new(static_keys, None);
+        assert!(validate_account_locks(account_keys, MAX_TX_ACCOUNT_LOCKS).is_ok());
+    }
+
+    #[test]
+    fn test_validate_account_locks_too_many_no_dynamic() {
+        let static_keys = &[Pubkey::new_unique(), Pubkey::new_unique()];
+        let account_keys = AccountKeys::new(static_keys, None);
+        assert_eq!(
+            validate_account_locks(account_keys, 1),
+            Err(TransactionError::TooManyAccountLocks)
+        );
+    }
+
+    #[test]
+    fn test_validate_account_locks_duplicate_no_dynamic() {
+        let duplicate_key = Pubkey::new_unique();
+        let static_keys = &[duplicate_key, Pubkey::new_unique(), duplicate_key];
+        let account_keys = AccountKeys::new(static_keys, None);
+        assert_eq!(
+            validate_account_locks(account_keys, MAX_TX_ACCOUNT_LOCKS),
+            Err(TransactionError::AccountLoadedTwice)
+        );
+    }
+
+    #[test]
+    fn test_validate_account_locks_valid_dynamic() {
+        let static_keys = &[Pubkey::new_unique(), Pubkey::new_unique()];
+        let dynamic_keys = LoadedAddresses {
+            writable: vec![Pubkey::new_unique()],
+            readonly: vec![Pubkey::new_unique()],
+        };
+        let account_keys = AccountKeys::new(static_keys, Some(&dynamic_keys));
+        assert!(validate_account_locks(account_keys, MAX_TX_ACCOUNT_LOCKS).is_ok());
+    }
+
+    #[test]
+    fn test_validate_account_locks_too_many_dynamic() {
+        let static_keys = &[Pubkey::new_unique()];
+        let dynamic_keys = LoadedAddresses {
+            writable: vec![Pubkey::new_unique()],
+            readonly: vec![Pubkey::new_unique()],
+        };
+        let account_keys = AccountKeys::new(static_keys, Some(&dynamic_keys));
+        assert_eq!(
+            validate_account_locks(account_keys, 2),
+            Err(TransactionError::TooManyAccountLocks)
+        );
+    }
+
+    #[test]
+    fn test_validate_account_locks_duplicate_dynamic() {
+        let duplicate_key = Pubkey::new_unique();
+        let static_keys = &[duplicate_key];
+        let dynamic_keys = LoadedAddresses {
+            writable: vec![Pubkey::new_unique()],
+            readonly: vec![duplicate_key],
+        };
+        let account_keys = AccountKeys::new(static_keys, Some(&dynamic_keys));
+        assert_eq!(
+            validate_account_locks(account_keys, MAX_TX_ACCOUNT_LOCKS),
+            Err(TransactionError::AccountLoadedTwice)
+        );
+    }
+
+    #[test]
+    fn test_has_duplicates_small() {
+        let mut keys = (0..16).map(|_| Pubkey::new_unique()).collect::<Vec<_>>();
+        let account_keys = AccountKeys::new(&keys, None);
+        assert!(!has_duplicates(account_keys));
+
+        keys[14] = keys[3]; // Duplicate key
+        let account_keys = AccountKeys::new(&keys, None);
+        assert!(has_duplicates(account_keys));
+    }
+
+    #[test]
+    fn test_has_duplicates_large() {
+        let mut keys = (0..64).map(|_| Pubkey::new_unique()).collect::<Vec<_>>();
+        let account_keys = AccountKeys::new(&keys, None);
+        assert!(!has_duplicates(account_keys));
+
+        keys[47] = keys[3]; // Duplicate key
+        let account_keys = AccountKeys::new(&keys, None);
+        assert!(has_duplicates(account_keys));
+    }
 }

accounts-db/src/accounts.rs

@@ -1,6 +1,6 @@
 use {
     crate::{
-        account_locks::AccountLocks,
+        account_locks::{validate_account_locks, AccountLocks},
         accounts_db::{
             AccountStorageEntry, AccountsAddRootTiming, AccountsDb, LoadHint, LoadedAccount,
             ScanAccountStorageData, ScanStorageResult, VerifyAccountsHashAndLamportsConfig,
@@ -520,7 +520,7 @@ impl Accounts {
         // Validate the account locks, then get iterator if successful validation.
         let tx_account_locks_results: Vec<Result<_>> = txs
             .map(|tx| {
-                SanitizedTransaction::validate_account_locks(tx.message(), tx_account_lock_limit)
+                validate_account_locks(tx.account_keys(), tx_account_lock_limit)
                     .map(|_| TransactionAccountLocksIterator::new(tx))
             })
             .collect();
@@ -530,19 +530,16 @@ impl Accounts {
     #[must_use]
     pub fn lock_accounts_with_results<'a>(
         &self,
-        txs: impl Iterator<Item = &'a SanitizedTransaction>,
+        txs: impl Iterator<Item = &'a (impl SVMMessage + 'a)>,
         results: impl Iterator<Item = Result<()>>,
         tx_account_lock_limit: usize,
     ) -> Vec<Result<()>> {
         // Validate the account locks, then get iterator if successful validation.
         let tx_account_locks_results: Vec<Result<_>> = txs
             .zip(results)
             .map(|(tx, result)| match result {
-                Ok(()) => SanitizedTransaction::validate_account_locks(
-                    tx.message(),
-                    tx_account_lock_limit,
-                )
-                .map(|_| TransactionAccountLocksIterator::new(tx)),
+                Ok(()) => validate_account_locks(tx.account_keys(), tx_account_lock_limit)
+                    .map(|_| TransactionAccountLocksIterator::new(tx)),
                 Err(err) => Err(err),
             })
             .collect();

sdk/program/src/message/account_keys.rs

@@ -17,6 +17,7 @@ pub struct AccountKeys<'a> {
 
 impl Index<usize> for AccountKeys<'_> {
     type Output = Pubkey;
+    #[inline]
     fn index(&self, index: usize) -> &Self::Output {
         self.get(index).expect("index is invalid")
     }
@@ -33,6 +34,7 @@ impl<'a> AccountKeys<'a> {
     /// 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.
+    #[inline]
     fn key_segment_iter(&self) -> impl Iterator<Item = &'a [Pubkey]> + Clone {
         if let Some(dynamic_keys) = self.dynamic_keys {
             [
@@ -51,6 +53,7 @@ impl<'a> AccountKeys<'a> {
     /// message account keys constructed from static keys, followed by dynamically
     /// loaded writable addresses, and lastly the list of dynamically loaded
     /// readonly addresses.
+    #[inline]
     pub fn get(&self, mut index: usize) -> Option<&'a Pubkey> {
         for key_segment in self.key_segment_iter() {
             if index < key_segment.len() {
@@ -63,6 +66,7 @@ impl<'a> AccountKeys<'a> {
     }
 
     /// Returns the total length of loaded accounts for a message
+    #[inline]
     pub fn len(&self) -> usize {
         let mut len = 0usize;
         for key_segment in self.key_segment_iter() {
@@ -77,6 +81,7 @@ impl<'a> AccountKeys<'a> {
     }
 
     /// Iterator for the addresses of the loaded accounts for a message
+    #[inline]
     pub fn iter(&self) -> impl Iterator<Item = &'a Pubkey> + Clone {
         self.key_segment_iter().flatten()
     }