From da088681ba8720cf56f8e3d1352b106b3b611135 Mon Sep 17 00:00:00 2001
From: Brooks <brooks@solana.com>
Date: Tue, 27 Feb 2024 18:08:25 -0500
Subject: [PATCH] Adds safer alternatives to get_internal() (#35325)

---
 accounts-db/src/accounts_index.rs        | 17 ++++++-----
 accounts-db/src/in_mem_accounts_index.rs | 39 ++++++++++++++++++++++--
 2 files changed, 45 insertions(+), 11 deletions(-)

diff --git a/accounts-db/src/accounts_index.rs b/accounts-db/src/accounts_index.rs
index 5ac5d9a7a..7a5c75669 100644
--- a/accounts-db/src/accounts_index.rs
+++ b/accounts-db/src/accounts_index.rs
@@ -1130,9 +1130,9 @@ impl<T: IndexValue, U: DiskIndexValue + From<T> + Into<T>> AccountsIndex<T, U> {
     pub fn get_and_then<R>(
         &self,
         pubkey: &Pubkey,
-        callback: impl FnOnce(Option<&AccountMapEntry<T>>) -> (bool, R),
+        callback: impl FnOnce(Option<&AccountMapEntryInner<T>>) -> (bool, R),
     ) -> R {
-        self.get_bin(pubkey).get_internal(pubkey, callback)
+        self.get_bin(pubkey).get_internal_inner(pubkey, callback)
     }
 
     /// Gets the index's entry for `pubkey`, with `ancestors` and `max_root`,
@@ -1159,10 +1159,8 @@ impl<T: IndexValue, U: DiskIndexValue + From<T> + Into<T>> AccountsIndex<T, U> {
     ///
     /// Prefer `get_and_then()` whenever possible.
     pub fn get_cloned(&self, pubkey: &Pubkey) -> Option<AccountMapEntry<T>> {
-        // We *must* add the index entry to the in-mem cache!
-        // If the index entry is only on-disk, returning a clone would allow the entry
-        // to be modified, but those modifications would be lost on drop!
-        self.get_and_then(pubkey, |entry| (true, entry.cloned()))
+        self.get_bin(pubkey)
+            .get_internal_cloned(pubkey, |entry| entry)
     }
 
     /// Is `pubkey` in the index?
@@ -1443,6 +1441,9 @@ impl<T: IndexValue, U: DiskIndexValue + From<T> + Into<T>> AccountsIndex<T, U> {
                 lock = Some(&self.account_maps[bin]);
                 last_bin = bin;
             }
+            // SAFETY: The caller must ensure that if `provide_entry_in_callback` is true, and
+            // if it's possible for `callback` to clone the entry Arc, then it must also add
+            // the entry to the in-mem cache if the entry is made dirty.
             lock.as_ref().unwrap().get_internal(pubkey, |entry| {
                 let mut cache = false;
                 match entry {
@@ -1830,7 +1831,7 @@ impl<T: IndexValue, U: DiskIndexValue + From<T> + Into<T>> AccountsIndex<T, U> {
 
     pub fn ref_count_from_storage(&self, pubkey: &Pubkey) -> RefCount {
         let map = self.get_bin(pubkey);
-        map.get_internal(pubkey, |entry| {
+        map.get_internal_inner(pubkey, |entry| {
             (
                 false,
                 entry.map(|entry| entry.ref_count()).unwrap_or_default(),
@@ -4073,7 +4074,7 @@ pub mod tests {
 
         let map = index.get_bin(&key);
         for expected in [false, true] {
-            assert!(map.get_internal(&key, |entry| {
+            assert!(map.get_internal_inner(&key, |entry| {
                 // check refcount BEFORE the unref
                 assert_eq!(u64::from(!expected), entry.unwrap().ref_count());
                 // first time, ref count was at 1, we can unref once. Unref should return false.
diff --git a/accounts-db/src/in_mem_accounts_index.rs b/accounts-db/src/in_mem_accounts_index.rs
index 1e8e8a8fd..054fd7589 100644
--- a/accounts-db/src/in_mem_accounts_index.rs
+++ b/accounts-db/src/in_mem_accounts_index.rs
@@ -320,7 +320,7 @@ impl<T: IndexValue, U: DiskIndexValue + From<T> + Into<T>> InMemAccountsIndex<T,
 
     /// lookup 'pubkey' in index (in mem or on disk)
     pub fn get(&self, pubkey: &K) -> Option<AccountMapEntry<T>> {
-        self.get_internal(pubkey, |entry| (true, entry.map(Arc::clone)))
+        self.get_internal_cloned(pubkey, |entry| entry)
     }
 
     /// set age of 'entry' to the future
@@ -331,7 +331,40 @@ impl<T: IndexValue, U: DiskIndexValue + From<T> + Into<T>> InMemAccountsIndex<T,
 
     /// lookup 'pubkey' in index (in_mem or disk).
     /// call 'callback' whether found or not
-    pub fn get_internal<RT>(
+    pub(crate) fn get_internal_inner<RT>(
+        &self,
+        pubkey: &K,
+        // return true if item should be added to in_mem cache
+        callback: impl for<'a> FnOnce(Option<&AccountMapEntryInner<T>>) -> (bool, RT),
+    ) -> RT {
+        // SAFETY: The entry Arc is not passed to `callback`, so
+        // it cannot live beyond this function call.
+        self.get_internal(pubkey, |entry| callback(entry.map(Arc::as_ref)))
+    }
+
+    /// lookup 'pubkey' in the index (in_mem or disk).
+    /// call 'callback' whether found or not
+    pub(crate) fn get_internal_cloned<RT>(
+        &self,
+        pubkey: &K,
+        callback: impl for<'a> FnOnce(Option<AccountMapEntry<T>>) -> RT,
+    ) -> RT {
+        // SAFETY: Since we're passing the entry Arc clone to `callback`, we must
+        // also add the entry to the in-mem cache.
+        self.get_internal(pubkey, |entry| (true, callback(entry.map(Arc::clone))))
+    }
+
+    /// lookup 'pubkey' in index (in_mem or disk).
+    /// call 'callback' whether found or not
+    ///
+    /// # Safety
+    ///
+    /// If the item is on-disk (and not in-mem), add if the item is/could be made dirty
+    /// *after* `callback` finishes (e.g. the entry Arc is cloned and saved by the caller),
+    /// then the disk entry *must* also be added to the in-mem cache.
+    ///
+    /// Prefer `get_internal_inner()` or `get_internal_cloned()` for safe alternatives.
+    pub(crate) fn get_internal<RT>(
         &self,
         pubkey: &K,
         // return true if item should be added to in_mem cache
@@ -446,7 +479,7 @@ impl<T: IndexValue, U: DiskIndexValue + From<T> + Into<T>> InMemAccountsIndex<T,
         pubkey: &Pubkey,
         user: impl for<'a> FnOnce(&mut RwLockWriteGuard<'a, SlotList<T>>) -> RT,
     ) -> Option<RT> {
-        self.get_internal(pubkey, |entry| {
+        self.get_internal_inner(pubkey, |entry| {
             (
                 true,
                 entry.map(|entry| {