API to match account owner against a set of owners (#30154)

Author: pgarg66

runtime/src/accounts_db.rs

@@ -42,9 +42,9 @@ use {
             get_ancient_append_vec_capacity, is_ancient, AccountsToStore, StorageSelector,
         },
         append_vec::{
-            aligned_stored_size, AppendVec, StorableAccountsWithHashesAndWriteVersions,
-            StoredAccountMeta, StoredMetaWriteVersion, APPEND_VEC_MMAPPED_FILES_OPEN,
-            STORE_META_OVERHEAD,
+            aligned_stored_size, AppendVec, MatchAccountOwnerError,
+            StorableAccountsWithHashesAndWriteVersions, StoredAccountMeta, StoredMetaWriteVersion,
+            APPEND_VEC_MMAPPED_FILES_OPEN, STORE_META_OVERHEAD,
         },
         cache_hash_data::{CacheHashData, CacheHashDataFile},
         contains::Contains,
@@ -815,6 +815,32 @@ impl<'a> LoadedAccountAccessor<'a> {
             }
         }
     }
+
+    fn account_matches_owners(&self, owners: &[&Pubkey]) -> Result<(), MatchAccountOwnerError> {
+        match self {
+            LoadedAccountAccessor::Cached(cached_account) => cached_account
+                .as_ref()
+                .and_then(|cached_account| {
+                    (!cached_account.account.is_zero_lamport()
+                        && owners.contains(&cached_account.account.owner()))
+                    .then_some(())
+                })
+                .ok_or(MatchAccountOwnerError::NoMatch),
+            LoadedAccountAccessor::Stored(maybe_storage_entry) => {
+                // storage entry may not be present if slot was cleaned up in
+                // between reading the accounts index and calling this function to
+                // get account meta from the storage entry here
+                maybe_storage_entry
+                    .as_ref()
+                    .map(|(storage_entry, offset)| {
+                        storage_entry
+                            .accounts
+                            .account_matches_owners(*offset, owners)
+                    })
+                    .unwrap_or(Err(MatchAccountOwnerError::UnableToLoad))
+            }
+        }
+    }
 }
 
 pub enum LoadedAccount<'a> {
@@ -4918,6 +4944,38 @@ impl AccountsDb {
         self.do_load(ancestors, pubkey, None, load_hint, LoadZeroLamports::None)
     }
 
+    pub fn account_matches_owners(
+        &self,
+        ancestors: &Ancestors,
+        account: &Pubkey,
+        owners: &[&Pubkey],
+    ) -> Result<(), MatchAccountOwnerError> {
+        let (slot, storage_location, _maybe_account_accesor) = self
+            .read_index_for_accessor_or_load_slow(ancestors, account, None, false)
+            .ok_or(MatchAccountOwnerError::UnableToLoad)?;
+
+        if !storage_location.is_cached() {
+            let result = self.read_only_accounts_cache.load(*account, slot);
+            if let Some(account) = result {
+                return (!account.is_zero_lamport() && owners.contains(&account.owner()))
+                    .then_some(())
+                    .ok_or(MatchAccountOwnerError::NoMatch);
+            }
+        }
+
+        let (account_accessor, _slot) = self
+            .retry_to_get_account_accessor(
+                slot,
+                storage_location,
+                ancestors,
+                account,
+                None,
+                LoadHint::Unspecified,
+            )
+            .ok_or(MatchAccountOwnerError::UnableToLoad)?;
+        account_accessor.account_matches_owners(owners)
+    }
+
     pub fn load_account_into_read_cache(&self, ancestors: &Ancestors, pubkey: &Pubkey) {
         self.do_load_with_populate_read_cache(
             ancestors,
@@ -14068,6 +14126,101 @@ pub mod tests {
         assert_eq!(db.read_only_accounts_cache.cache_len(), 1);
     }
 
+    #[test]
+    fn test_account_matches_owners() {
+        let db = Arc::new(AccountsDb::new_with_config_for_tests(
+            Vec::new(),
+            &ClusterType::Development,
+            AccountSecondaryIndexes::default(),
+            AccountShrinkThreshold::default(),
+        ));
+
+        let owners: Vec<Pubkey> = (0..2).map(|_| Pubkey::new_unique()).collect();
+        let owners_refs: Vec<&Pubkey> = owners.iter().collect();
+
+        let account1_key = Pubkey::new_unique();
+        let account1 = AccountSharedData::new(321, 10, &owners[0]);
+
+        let account2_key = Pubkey::new_unique();
+        let account2 = AccountSharedData::new(1, 1, &owners[1]);
+
+        let account3_key = Pubkey::new_unique();
+        let account3 = AccountSharedData::new(1, 1, &Pubkey::new_unique());
+
+        // Account with 0 lamports
+        let account4_key = Pubkey::new_unique();
+        let account4 = AccountSharedData::new(0, 1, &owners[1]);
+
+        db.store_cached((0, &[(&account1_key, &account1)][..]), None);
+        db.store_cached((1, &[(&account2_key, &account2)][..]), None);
+        db.store_cached((2, &[(&account3_key, &account3)][..]), None);
+        db.store_cached((3, &[(&account4_key, &account4)][..]), None);
+
+        db.add_root(0);
+        db.add_root(1);
+        db.add_root(2);
+        db.add_root(3);
+
+        // Flush the cache so that the account meta will be read from the storage
+        db.flush_accounts_cache(true, None);
+        db.clean_accounts_for_tests();
+
+        assert_eq!(
+            db.account_matches_owners(&Ancestors::default(), &account1_key, &owners_refs),
+            Ok(())
+        );
+        assert_eq!(
+            db.account_matches_owners(&Ancestors::default(), &account2_key, &owners_refs),
+            Ok(())
+        );
+        assert_eq!(
+            db.account_matches_owners(&Ancestors::default(), &account3_key, &owners_refs),
+            Err(MatchAccountOwnerError::NoMatch)
+        );
+        assert_eq!(
+            db.account_matches_owners(&Ancestors::default(), &account4_key, &owners_refs),
+            Err(MatchAccountOwnerError::NoMatch)
+        );
+        assert_eq!(
+            db.account_matches_owners(&Ancestors::default(), &Pubkey::new_unique(), &owners_refs),
+            Err(MatchAccountOwnerError::UnableToLoad)
+        );
+
+        // Flush the cache and load account1 (so that it's in the cache)
+        db.flush_accounts_cache(true, None);
+        db.clean_accounts_for_tests();
+        let _ = db
+            .do_load(
+                &Ancestors::default(),
+                &account1_key,
+                Some(0),
+                LoadHint::Unspecified,
+                LoadZeroLamports::SomeWithZeroLamportAccountForTests,
+            )
+            .unwrap();
+
+        assert_eq!(
+            db.account_matches_owners(&Ancestors::default(), &account1_key, &owners_refs),
+            Ok(())
+        );
+        assert_eq!(
+            db.account_matches_owners(&Ancestors::default(), &account2_key, &owners_refs),
+            Ok(())
+        );
+        assert_eq!(
+            db.account_matches_owners(&Ancestors::default(), &account3_key, &owners_refs),
+            Err(MatchAccountOwnerError::NoMatch)
+        );
+        assert_eq!(
+            db.account_matches_owners(&Ancestors::default(), &account4_key, &owners_refs),
+            Err(MatchAccountOwnerError::NoMatch)
+        );
+        assert_eq!(
+            db.account_matches_owners(&Ancestors::default(), &Pubkey::new_unique(), &owners_refs),
+            Err(MatchAccountOwnerError::UnableToLoad)
+        );
+    }
+
     /// a test that will accept either answer
     const LOAD_ZERO_LAMPORTS_ANY_TESTS: LoadZeroLamports = LoadZeroLamports::None;
 

runtime/src/append_vec.rs

@@ -28,6 +28,7 @@ use {
             Mutex,
         },
     },
+    thiserror::Error,
 };
 
 pub mod test_utils;
@@ -270,6 +271,14 @@ impl<'a> Iterator for AppendVecAccountsIter<'a> {
     }
 }
 
+#[derive(Error, Debug, PartialEq, Eq)]
+pub enum MatchAccountOwnerError {
+    #[error("The account owner does not match with the provided list")]
+    NoMatch,
+    #[error("Unable to load the account")]
+    UnableToLoad,
+}
+
 /// A thread-safe, file-backed block of memory used to store `Account` instances. Append operations
 /// are serialized such that only one thread updates the internal `append_lock` at a time. No
 /// restrictions are placed on reading. That is, one may read items from one thread while another
@@ -572,7 +581,7 @@ impl AppendVec {
         Some((unsafe { &*ptr }, next))
     }
 
-    /// Return account metadata for the account at `offset` if its data doesn't overrun
+    /// Return stored account metadata for the account at `offset` if its data doesn't overrun
     /// the internal buffer. Otherwise return None. Also return the offset of the first byte
     /// after the requested data that falls on a 64-byte boundary.
     pub fn get_account<'a>(&'a self, offset: usize) -> Option<(StoredAccountMeta<'a>, usize)> {
@@ -594,6 +603,31 @@ impl AppendVec {
         ))
     }
 
+    fn get_account_meta<'a>(&self, offset: usize) -> Option<&'a AccountMeta> {
+        // Skip over StoredMeta data in the account
+        let offset = offset.checked_add(mem::size_of::<StoredMeta>())?;
+        // u64_align! does an unchecked add for alignment. Check that it won't cause an overflow.
+        offset.checked_add(ALIGN_BOUNDARY_OFFSET - 1)?;
+        let (account_meta, _): (&AccountMeta, _) = self.get_type(u64_align!(offset))?;
+        Some(account_meta)
+    }
+
+    /// Return Some(true) if the account owner at `offset` is one of the pubkeys in `owners`.
+    /// Return Some(false) if the account owner is not one of the pubkeys in `owners`.
+    /// It returns None if the `offset` value causes a data overrun.
+    pub fn account_matches_owners(
+        &self,
+        offset: usize,
+        owners: &[&Pubkey],
+    ) -> Result<(), MatchAccountOwnerError> {
+        let account_meta = self
+            .get_account_meta(offset)
+            .ok_or(MatchAccountOwnerError::UnableToLoad)?;
+        (account_meta.lamports != 0 && owners.contains(&&account_meta.owner))
+            .then_some(())
+            .ok_or(MatchAccountOwnerError::NoMatch)
+    }
+
     #[cfg(test)]
     pub fn get_account_test(&self, offset: usize) -> Option<(StoredMeta, AccountSharedData)> {
         let (stored_account, _) = self.get_account(offset)?;
@@ -1047,6 +1081,47 @@ pub mod tests {
         assert_eq!(av.get_account_test(index1).unwrap(), account1);
     }
 
+    #[test]
+    fn test_account_matches_owners() {
+        let path = get_append_vec_path("test_append_data");
+        let av = AppendVec::new(&path.path, true, 1024 * 1024);
+        let owners: Vec<Pubkey> = (0..2).map(|_| Pubkey::new_unique()).collect();
+        let owners_refs: Vec<&Pubkey> = owners.iter().collect();
+
+        let mut account = create_test_account(5);
+        account.1.set_owner(owners[0]);
+        let index = av.append_account_test(&account).unwrap();
+        assert_eq!(av.account_matches_owners(index, &owners_refs), Ok(()));
+
+        let mut account1 = create_test_account(6);
+        account1.1.set_owner(owners[1]);
+        let index1 = av.append_account_test(&account1).unwrap();
+        assert_eq!(av.account_matches_owners(index1, &owners_refs), Ok(()));
+        assert_eq!(av.account_matches_owners(index, &owners_refs), Ok(()));
+
+        let mut account2 = create_test_account(6);
+        account2.1.set_owner(Pubkey::new_unique());
+        let index2 = av.append_account_test(&account2).unwrap();
+        assert_eq!(
+            av.account_matches_owners(index2, &owners_refs),
+            Err(MatchAccountOwnerError::NoMatch)
+        );
+
+        // tests for overflow
+        assert_eq!(
+            av.account_matches_owners(usize::MAX - mem::size_of::<StoredMeta>(), &owners_refs),
+            Err(MatchAccountOwnerError::UnableToLoad)
+        );
+
+        assert_eq!(
+            av.account_matches_owners(
+                usize::MAX - mem::size_of::<StoredMeta>() - mem::size_of::<AccountMeta>() + 1,
+                &owners_refs
+            ),
+            Err(MatchAccountOwnerError::UnableToLoad)
+        );
+    }
+
     #[test]
     fn test_append_vec_append_many() {
         let path = get_append_vec_path("test_append_many");