Check bank capitalization (#11927)

* Check bank capitalization

* Simplify and unify capitalization calculation

* Improve and add tests

* Avoid overflow and inhibit automatic restart

* Fix test

* Tweak checked sum for cap. and add tests

* Fix broken build after merge conflicts..

* Rename to ClusterType

* Rename confusing method

* Clarify comment

* Verify cap. in rent and inflation tests

Co-authored-by: Stephen Akridge <sakridge@gmail.com>

Cargo.lock

@@ -3272,6 +3272,7 @@ dependencies = [
  "solana-measure",
  "solana-runtime",
  "solana-sdk 1.4.0",
+ "solana-version",
 ]
 
 [[package]]

accounts-bench/Cargo.toml

@@ -14,6 +14,7 @@ solana-logger = { path = "../logger", version = "1.4.0" }
 solana-runtime = { path = "../runtime", version = "1.4.0" }
 solana-measure = { path = "../measure", version = "1.4.0" }
 solana-sdk = { path = "../sdk", version = "1.4.0" }
+solana-version = { path = "../version", version = "1.4.0" }
 rand = "0.7.0"
 clap = "2.33.1"
 crossbeam-channel = "0.4"

accounts-bench/src/main.rs

@@ -1,4 +1,4 @@
-use clap::{value_t, App, Arg};
+use clap::{crate_description, crate_name, value_t, App, Arg};
 use rayon::prelude::*;
 use solana_measure::measure::Measure;
 use solana_runtime::{
@@ -6,15 +6,16 @@ use solana_runtime::{
     accounts_index::Ancestors,
 };
 use solana_sdk::{genesis_config::ClusterType, pubkey::Pubkey};
+use std::env;
 use std::fs;
 use std::path::PathBuf;
 
 fn main() {
     solana_logger::setup();
 
-    let matches = App::new("crate")
+    let matches = App::new(crate_name!())
-        .about("about")
+        .about(crate_description!())
-        .version("version")
+        .version(solana_version::version!())
         .arg(
             Arg::with_name("num_slots")
                 .long("num_slots")
@@ -50,7 +51,8 @@ fn main() {
     let clean = matches.is_present("clean");
     println!("clean: {:?}", clean);
 
-    let path = PathBuf::from("farf/accounts-bench");
+    let path = PathBuf::from(env::var("FARF_DIR").unwrap_or_else(|_| "farf".to_owned()))
+        .join("accounts-bench");
     if fs::remove_dir_all(path.clone()).is_err() {
         println!("Warning: Couldn't remove {:?}", path);
     }
@@ -96,7 +98,7 @@ fn main() {
         } else {
             let mut pubkeys: Vec<Pubkey> = vec![];
             let mut time = Measure::start("hash");
-            let hash = accounts.accounts_db.update_accounts_hash(0, &ancestors);
+            let hash = accounts.accounts_db.update_accounts_hash(0, &ancestors).0;
             time.stop();
             println!("hash: {} {}", hash, time);
             create_test_accounts(&accounts, &mut pubkeys, 1, 0);

core/src/non_circulating_supply.rs

@@ -18,7 +18,7 @@ pub fn calculate_non_circulating_supply(bank: &Arc<Bank>) -> NonCirculatingSuppl
     let withdraw_authority_list = withdraw_authority();
 
     let clock = bank.clock();
-    let stake_accounts = bank.get_program_accounts(Some(&solana_stake_program::id()));
+    let stake_accounts = bank.get_program_accounts(&solana_stake_program::id());
     for (pubkey, account) in stake_accounts.iter() {
         let stake_account = StakeState::from(&account).unwrap_or_default();
         match stake_account {

core/src/rpc.rs

@@ -1291,7 +1291,7 @@ fn get_filtered_program_accounts(
     program_id: &Pubkey,
     filters: Vec<RpcFilterType>,
 ) -> impl Iterator<Item = (Pubkey, Account)> {
-    bank.get_program_accounts(Some(&program_id))
+    bank.get_program_accounts(&program_id)
         .into_iter()
         .filter(move |(_, account)| {
             filters.iter().all(|filter_type| match filter_type {

ledger-tool/src/main.rs

@@ -44,7 +44,7 @@ use solana_vote_program::{
 };
 use std::{
     collections::{BTreeMap, BTreeSet, HashMap, HashSet},
-    convert::{TryFrom, TryInto},
+    convert::TryInto,
     ffi::OsStr,
     fs::{self, File},
     io::{self, stdout, BufRead, BufReader, Write},
@@ -708,16 +708,7 @@ fn open_genesis_config_by(ledger_path: &Path, matches: &ArgMatches<'_>) -> Genes
 }
 
 fn assert_capitalization(bank: &Bank) {
-    let calculated_capitalization = bank.calculate_capitalization();
+    assert!(bank.calculate_and_verify_capitalization());
-    assert_eq!(
-        bank.capitalization(),
-        calculated_capitalization,
-        "Capitalization mismatch!?: +/-{}",
-        Sol(u64::try_from(
-            (i128::from(calculated_capitalization) - i128::from(bank.capitalization())).abs()
-        )
-        .unwrap()),
-    );
 }
 
 #[allow(clippy::cognitive_complexity)]
@@ -1695,7 +1686,7 @@ fn main() {
 
                     if remove_stake_accounts {
                         for (address, mut account) in bank
-                            .get_program_accounts(Some(&solana_stake_program::id()))
+                            .get_program_accounts(&solana_stake_program::id())
                             .into_iter()
                         {
                             account.lamports = 0;
@@ -1723,7 +1714,7 @@ fn main() {
 
                         // Delete existing vote accounts
                         for (address, mut account) in bank
-                            .get_program_accounts(Some(&solana_vote_program::id()))
+                            .get_program_accounts(&solana_vote_program::id())
                             .into_iter()
                         {
                             account.lamports = 0;

ledger/src/blockstore_processor.rs

@@ -293,6 +293,9 @@ pub enum BlockstoreProcessorError {
 
     #[error("invalid hard fork")]
     InvalidHardFork(Slot),
+
+    #[error("root bank with mismatched capitalization at {0}")]
+    RootBankWithMismatchedCapitalization(Slot),
 }
 
 /// Callback for accessing bank state while processing the blockstore
@@ -481,6 +484,13 @@ fn do_process_blockstore_from_root(
     );
     assert!(bank_forks.active_banks().is_empty());
 
+    // We might be promptly restarted after bad capitalization was detected while creating newer snapshot.
+    // In that case, we're most likely restored from the last good snapshot and replayed up to this root.
+    // So again check here for the bad capitalization to avoid to continue until the next snapshot creation.
+    if !bank_forks.root_bank().calculate_and_verify_capitalization() {
+        return Err(BlockstoreProcessorError::RootBankWithMismatchedCapitalization(root));
+    }
+
     Ok((bank_forks, leader_schedule_cache))
 }
 

runtime/benches/accounts.rs

@@ -74,10 +74,12 @@ fn test_accounts_hash_bank_hash(bencher: &mut Bencher) {
         &ClusterType::Development,
     );
     let mut pubkeys: Vec<Pubkey> = vec![];
-    create_test_accounts(&accounts, &mut pubkeys, 60000, 0);
+    let num_accounts = 60_000;
+    let slot = 0;
+    create_test_accounts(&accounts, &mut pubkeys, num_accounts, slot);
     let ancestors = vec![(0, 0)].into_iter().collect();
-    accounts.accounts_db.update_accounts_hash(0, &ancestors);
+    let (_, total_lamports) = accounts.accounts_db.update_accounts_hash(0, &ancestors);
-    bencher.iter(|| assert!(accounts.verify_bank_hash(0, &ancestors)));
+    bencher.iter(|| assert!(accounts.verify_bank_hash_and_lamports(0, &ancestors, total_lamports)));
 }
 
 #[bench]

runtime/src/accounts.rs

@@ -428,9 +428,32 @@ impl Accounts {
         accounts_balances
     }
 
+    pub fn calculate_capitalization(&self, ancestors: &Ancestors) -> u64 {
+        let balances = self
+            .load_all(ancestors)
+            .into_iter()
+            .map(|(_pubkey, account, _slot)| {
+                AccountsDB::account_balance_for_capitalization(
+                    account.lamports,
+                    &account.owner,
+                    account.executable,
+                )
+            });
+
+        AccountsDB::checked_sum_for_capitalization(balances)
+    }
+
     #[must_use]
-    pub fn verify_bank_hash(&self, slot: Slot, ancestors: &Ancestors) -> bool {
+    pub fn verify_bank_hash_and_lamports(
-        if let Err(err) = self.accounts_db.verify_bank_hash(slot, ancestors) {
+        &self,
+        slot: Slot,
+        ancestors: &Ancestors,
+        total_lamports: u64,
+    ) -> bool {
+        if let Err(err) =
+            self.accounts_db
+                .verify_bank_hash_and_lamports(slot, ancestors, total_lamports)
+        {
             warn!("verify_bank_hash failed: {:?}", err);
             false
         } else {
@@ -460,13 +483,13 @@ impl Accounts {
     pub fn load_by_program(
         &self,
         ancestors: &Ancestors,
-        program_id: Option<&Pubkey>,
+        program_id: &Pubkey,
     ) -> Vec<(Pubkey, Account)> {
         self.accounts_db.scan_accounts(
             ancestors,
             |collector: &mut Vec<(Pubkey, Account)>, some_account_tuple| {
                 Self::load_while_filtering(collector, some_account_tuple, |account| {
-                    program_id.is_none() || Some(&account.owner) == program_id
+                    account.owner == *program_id
                 })
             },
         )

runtime/src/accounts_db.rs

@@ -40,6 +40,7 @@ use solana_sdk::{
 use std::convert::TryFrom;
 use std::{
     collections::{HashMap, HashSet},
+    convert::TryInto,
     io::{Error as IOError, Result as IOResult},
     iter::FromIterator,
     ops::RangeBounds,
@@ -154,6 +155,7 @@ pub enum BankHashVerificationError {
     MismatchedAccountHash,
     MismatchedBankHash,
     MissingBankHash,
+    MismatchedTotalLamports(u64, u64),
 }
 
 /// Persistent storage structure holding the accounts
@@ -1614,8 +1616,11 @@ impl AccountsDB {
         );
     }
 
-    pub fn compute_merkle_root(hashes: Vec<(Pubkey, Hash)>, fanout: usize) -> Hash {
+    pub fn compute_merkle_root(hashes: Vec<(Pubkey, Hash, u64)>, fanout: usize) -> Hash {
-        let hashes: Vec<_> = hashes.into_iter().map(|(_pubkey, hash)| hash).collect();
+        let hashes: Vec<_> = hashes
+            .into_iter()
+            .map(|(_pubkey, hash, _lamports)| hash)
+            .collect();
         let mut hashes: Vec<_> = hashes.chunks(fanout).map(|x| x.to_vec()).collect();
         while hashes.len() > 1 {
             let mut time = Measure::start("time");
@@ -1640,27 +1645,77 @@ impl AccountsDB {
         hasher.result()
     }
 
-    fn accumulate_account_hashes(mut hashes: Vec<(Pubkey, Hash)>) -> Hash {
+    fn accumulate_account_hashes(hashes: Vec<(Pubkey, Hash, u64)>) -> Hash {
-        let mut sort = Measure::start("sort");
+        let (hash, ..) = Self::do_accumulate_account_hashes_and_capitalization(hashes, false);
+        hash
+    }
+
+    fn accumulate_account_hashes_and_capitalization(
+        hashes: Vec<(Pubkey, Hash, u64)>,
+    ) -> (Hash, u64) {
+        let (hash, cap) = Self::do_accumulate_account_hashes_and_capitalization(hashes, true);
+        (hash, cap.unwrap())
+    }
+
+    fn do_accumulate_account_hashes_and_capitalization(
+        mut hashes: Vec<(Pubkey, Hash, u64)>,
+        calculate_cap: bool,
+    ) -> (Hash, Option<u64>) {
+        let mut sort_time = Measure::start("sort");
         hashes.par_sort_by(|a, b| a.0.cmp(&b.0));
-        sort.stop();
+        sort_time.stop();
+
+        let mut sum_time = Measure::start("cap");
+        let cap = if calculate_cap {
+            Some(Self::checked_sum_for_capitalization(
+                hashes.iter().map(|(_, _, lamports)| *lamports),
+            ))
+        } else {
+            None
+        };
+        sum_time.stop();
+
         let mut hash_time = Measure::start("hash");
-
         let fanout = 16;
-
         let res = Self::compute_merkle_root(hashes, fanout);
-
         hash_time.stop();
-        debug!("{} {}", sort, hash_time);
 
-        res
+        debug!("{} {} {}", sort_time, hash_time, sum_time);
+
+        (res, cap)
+    }
+
+    pub fn checked_sum_for_capitalization<T: Iterator<Item = u64>>(balances: T) -> u64 {
+        balances
+            .map(|b| b as u128)
+            .sum::<u128>()
+            .try_into()
+            .expect("overflow is detected while summing capitalization")
+    }
+
+    pub fn account_balance_for_capitalization(
+        lamports: u64,
+        owner: &Pubkey,
+        executable: bool,
+    ) -> u64 {
+        let is_specially_retained = (solana_sdk::native_loader::check_id(owner) && executable)
+            || solana_sdk::sysvar::check_id(owner);
+
+        if is_specially_retained {
+            // specially retained accounts always have an initial 1 lamport
+            // balance, but could be modified by transfers which increase
+            // the balance but don't affect the capitalization.
+            lamports - 1
+        } else {
+            lamports
+        }
     }
 
     fn calculate_accounts_hash(
         &self,
         ancestors: &Ancestors,
         check_hash: bool,
-    ) -> Result<Hash, BankHashVerificationError> {
+    ) -> Result<(Hash, u64), BankHashVerificationError> {
         use BankHashVerificationError::*;
         let mut scan = Measure::start("scan");
         let accounts_index = self.accounts_index.read().unwrap();
@@ -1679,6 +1734,11 @@ impl AccountsDB {
                             .and_then(|storage_map| storage_map.get(&account_info.store_id))
                             .and_then(|store| {
                                 let account = store.accounts.get_account(account_info.offset)?.0;
+                                let balance = Self::account_balance_for_capitalization(
+                                    account_info.lamports,
+                                    &account.account_meta.owner,
+                                    account.account_meta.executable,
+                                );
 
                                 if check_hash {
                                     let hash = Self::hash_stored_account(
@@ -1694,7 +1754,7 @@ impl AccountsDB {
                                     }
                                 }
 
-                                Some((**pubkey, *account.hash))
+                                Some((**pubkey, *account.hash, balance))
                             })
                     } else {
                         None
@@ -1716,7 +1776,8 @@ impl AccountsDB {
         let hash_total = hashes.len();
 
         let mut accumulate = Measure::start("accumulate");
-        let accumulated_hash = Self::accumulate_account_hashes(hashes);
+        let (accumulated_hash, total_lamports) =
+            Self::accumulate_account_hashes_and_capitalization(hashes);
         accumulate.stop();
         datapoint_info!(
             "update_accounts_hash",
@@ -1724,7 +1785,7 @@ impl AccountsDB {
             ("hash_accumulate", accumulate.as_us(), i64),
             ("hash_total", hash_total, i64),
         );
-        Ok(accumulated_hash)
+        Ok((accumulated_hash, total_lamports))
     }
 
     pub fn get_accounts_hash(&self, slot: Slot) -> Hash {
@@ -1733,22 +1794,32 @@ impl AccountsDB {
         bank_hash_info.snapshot_hash
     }
 
-    pub fn update_accounts_hash(&self, slot: Slot, ancestors: &Ancestors) -> Hash {
+    pub fn update_accounts_hash(&self, slot: Slot, ancestors: &Ancestors) -> (Hash, u64) {
-        let hash = self.calculate_accounts_hash(ancestors, false).unwrap();
+        let (hash, total_lamports) = self.calculate_accounts_hash(ancestors, false).unwrap();
         let mut bank_hashes = self.bank_hashes.write().unwrap();
         let mut bank_hash_info = bank_hashes.get_mut(&slot).unwrap();
         bank_hash_info.snapshot_hash = hash;
-        hash
+        (hash, total_lamports)
     }
 
-    pub fn verify_bank_hash(
+    pub fn verify_bank_hash_and_lamports(
         &self,
         slot: Slot,
         ancestors: &Ancestors,
+        total_lamports: u64,
     ) -> Result<(), BankHashVerificationError> {
         use BankHashVerificationError::*;
 
-        let calculated_hash = self.calculate_accounts_hash(ancestors, true)?;
+        let (calculated_hash, calculated_lamports) =
+            self.calculate_accounts_hash(ancestors, true)?;
+
+        if calculated_lamports != total_lamports {
+            warn!(
+                "Mismatched total lamports: {} calculated: {}",
+                total_lamports, calculated_lamports
+            );
+            return Err(MismatchedTotalLamports(calculated_lamports, total_lamports));
+        }
 
         let bank_hashes = self.bank_hashes.read().unwrap();
         if let Some(found_hash_info) = bank_hashes.get(&slot) {
@@ -1789,7 +1860,7 @@ impl AccountsDB {
         let mut accumulate = Measure::start("accumulate");
         let hashes: Vec<_> = account_maps
             .into_iter()
-            .map(|(pubkey, (_, hash))| (pubkey, hash))
+            .map(|(pubkey, (_, hash))| (pubkey, hash, 0))
             .collect();
         let ret = Self::accumulate_account_hashes(hashes);
         accumulate.stop();
@@ -3305,7 +3376,9 @@ pub mod tests {
         assert_load_account(&accounts, current_slot, purged_pubkey2, 0);
         assert_load_account(&accounts, current_slot, dummy_pubkey, dummy_lamport);
 
-        accounts.verify_bank_hash(4, &HashMap::default()).unwrap();
+        accounts
+            .verify_bank_hash_and_lamports(4, &HashMap::default(), 1222)
+            .unwrap();
     }
 
     #[test]
@@ -3694,7 +3767,7 @@ pub mod tests {
         solana_logger::setup();
         let db = AccountsDB::new(Vec::new(), &ClusterType::Development);
 
-        let key = Pubkey::default();
+        let key = Pubkey::new_rand();
         let some_data_len = 0;
         let some_slot: Slot = 0;
         let account = Account::new(1, some_data_len, &key);
@@ -3703,11 +3776,14 @@ pub mod tests {
         db.store(some_slot, &[(&key, &account)]);
         db.add_root(some_slot);
         db.update_accounts_hash(some_slot, &ancestors);
-        assert_matches!(db.verify_bank_hash(some_slot, &ancestors), Ok(_));
+        assert_matches!(
+            db.verify_bank_hash_and_lamports(some_slot, &ancestors, 1),
+            Ok(_)
+        );
 
         db.bank_hashes.write().unwrap().remove(&some_slot).unwrap();
         assert_matches!(
-            db.verify_bank_hash(some_slot, &ancestors),
+            db.verify_bank_hash_and_lamports(some_slot, &ancestors, 1),
             Err(MissingBankHash)
         );
 
@@ -3722,11 +3798,51 @@ pub mod tests {
             .unwrap()
             .insert(some_slot, bank_hash_info);
         assert_matches!(
-            db.verify_bank_hash(some_slot, &ancestors),
+            db.verify_bank_hash_and_lamports(some_slot, &ancestors, 1),
             Err(MismatchedBankHash)
         );
     }
 
+    #[test]
+    fn test_verify_bank_capitalization() {
+        use BankHashVerificationError::*;
+        solana_logger::setup();
+        let db = AccountsDB::new(Vec::new(), &ClusterType::Development);
+
+        let key = Pubkey::new_rand();
+        let some_data_len = 0;
+        let some_slot: Slot = 0;
+        let account = Account::new(1, some_data_len, &key);
+        let ancestors = vec![(some_slot, 0)].into_iter().collect();
+
+        db.store(some_slot, &[(&key, &account)]);
+        db.add_root(some_slot);
+        db.update_accounts_hash(some_slot, &ancestors);
+        assert_matches!(
+            db.verify_bank_hash_and_lamports(some_slot, &ancestors, 1),
+            Ok(_)
+        );
+
+        let native_account_pubkey = Pubkey::new_rand();
+        db.store(
+            some_slot,
+            &[(
+                &native_account_pubkey,
+                &solana_sdk::native_loader::create_loadable_account("foo"),
+            )],
+        );
+        db.update_accounts_hash(some_slot, &ancestors);
+        assert_matches!(
+            db.verify_bank_hash_and_lamports(some_slot, &ancestors, 1),
+            Ok(_)
+        );
+
+        assert_matches!(
+            db.verify_bank_hash_and_lamports(some_slot, &ancestors, 10),
+            Err(MismatchedTotalLamports(expected, actual)) if expected == 1 && actual == 10
+        );
+    }
+
     #[test]
     fn test_verify_bank_hash_no_account() {
         solana_logger::setup();
@@ -3741,7 +3857,10 @@ pub mod tests {
             .insert(some_slot, BankHashInfo::default());
         db.add_root(some_slot);
         db.update_accounts_hash(some_slot, &ancestors);
-        assert_matches!(db.verify_bank_hash(some_slot, &ancestors), Ok(_));
+        assert_matches!(
+            db.verify_bank_hash_and_lamports(some_slot, &ancestors, 0),
+            Ok(_)
+        );
     }
 
     #[test]
@@ -3764,7 +3883,7 @@ pub mod tests {
         db.store_with_hashes(some_slot, accounts, &[some_hash]);
         db.add_root(some_slot);
         assert_matches!(
-            db.verify_bank_hash(some_slot, &ancestors),
+            db.verify_bank_hash_and_lamports(some_slot, &ancestors, 1),
             Err(MismatchedAccountHash)
         );
     }
@@ -4184,12 +4303,12 @@ pub mod tests {
         let no_ancestors = HashMap::default();
         accounts.update_accounts_hash(current_slot, &no_ancestors);
         accounts
-            .verify_bank_hash(current_slot, &no_ancestors)
+            .verify_bank_hash_and_lamports(current_slot, &no_ancestors, 22300)
             .unwrap();
 
         let accounts = reconstruct_accounts_db_via_serialization(&accounts, current_slot);
         accounts
-            .verify_bank_hash(current_slot, &no_ancestors)
+            .verify_bank_hash_and_lamports(current_slot, &no_ancestors, 22300)
             .unwrap();
 
         // repeating should be no-op
@@ -4376,4 +4495,79 @@ pub mod tests {
             shrink_thread.join().unwrap();
         }
     }
+
+    #[test]
+    fn test_account_balance_for_capitalization_normal() {
+        // system accounts
+        assert_eq!(
+            AccountsDB::account_balance_for_capitalization(10, &Pubkey::default(), false),
+            10
+        );
+        // any random program data accounts
+        assert_eq!(
+            AccountsDB::account_balance_for_capitalization(10, &Pubkey::new_rand(), false),
+            10
+        );
+    }
+
+    #[test]
+    fn test_account_balance_for_capitalization_sysvar() {
+        use solana_sdk::sysvar::Sysvar;
+
+        let normal_sysvar = solana_sdk::slot_history::SlotHistory::default().create_account(1);
+        assert_eq!(
+            AccountsDB::account_balance_for_capitalization(
+                normal_sysvar.lamports,
+                &normal_sysvar.owner,
+                normal_sysvar.executable
+            ),
+            0
+        );
+
+        // currently transactions can send any lamports to sysvars although this is not sensible.
+        assert_eq!(
+            AccountsDB::account_balance_for_capitalization(10, &solana_sdk::sysvar::id(), false),
+            9
+        );
+    }
+
+    #[test]
+    fn test_account_balance_for_capitalization_native_program() {
+        let normal_native_program = solana_sdk::native_loader::create_loadable_account("foo");
+        assert_eq!(
+            AccountsDB::account_balance_for_capitalization(
+                normal_native_program.lamports,
+                &normal_native_program.owner,
+                normal_native_program.executable
+            ),
+            0
+        );
+
+        // test maliciously assigned bogus native loader account
+        assert_eq!(
+            AccountsDB::account_balance_for_capitalization(
+                1,
+                &solana_sdk::native_loader::id(),
+                false
+            ),
+            1
+        )
+    }
+
+    #[test]
+    fn test_checked_sum_for_capitalization_normal() {
+        assert_eq!(
+            AccountsDB::checked_sum_for_capitalization(vec![1, 2].into_iter()),
+            3
+        );
+    }
+
+    #[test]
+    #[should_panic(expected = "overflow is detected while summing capitalization")]
+    fn test_checked_sum_for_capitalization_overflow() {
+        assert_eq!(
+            AccountsDB::checked_sum_for_capitalization(vec![1, u64::max_value()].into_iter()),
+            3
+        );
+    }
 }

runtime/src/bank.rs

@@ -2625,6 +2625,28 @@ impl Bank {
         }
     }
 
+    #[cfg(test)]
+    fn add_account_and_update_capitalization(&self, pubkey: &Pubkey, new_account: &Account) {
+        if let Some(old_account) = self.get_account(&pubkey) {
+            if new_account.lamports > old_account.lamports {
+                self.capitalization.fetch_add(
+                    new_account.lamports - old_account.lamports,
+                    Ordering::Relaxed,
+                );
+            } else {
+                self.capitalization.fetch_sub(
+                    old_account.lamports - new_account.lamports,
+                    Ordering::Relaxed,
+                );
+            }
+        } else {
+            self.capitalization
+                .fetch_add(new_account.lamports, Ordering::Relaxed);
+        }
+
+        self.store_account(pubkey, new_account);
+    }
+
     pub fn withdraw(&self, pubkey: &Pubkey, lamports: u64) -> Result<()> {
         match self.get_account(pubkey) {
             Some(mut account) => {
@@ -2739,7 +2761,7 @@ impl Bank {
             .map(|(acc, _slot)| acc)
     }
 
-    pub fn get_program_accounts(&self, program_id: Option<&Pubkey>) -> Vec<(Pubkey, Account)> {
+    pub fn get_program_accounts(&self, program_id: &Pubkey) -> Vec<(Pubkey, Account)> {
         self.rc
             .accounts
             .load_by_program(&self.ancestors, program_id)
@@ -2873,9 +2895,11 @@ impl Bank {
     /// snapshot.
     #[must_use]
     fn verify_bank_hash(&self) -> bool {
-        self.rc
+        self.rc.accounts.verify_bank_hash_and_lamports(
-            .accounts
+            self.slot(),
-            .verify_bank_hash(self.slot(), &self.ancestors)
+            &self.ancestors,
+            self.capitalization(),
+        )
     }
 
     pub fn get_snapshot_storages(&self) -> SnapshotStorages {
@@ -2905,22 +2929,21 @@ impl Bank {
     }
 
     pub fn calculate_capitalization(&self) -> u64 {
-        self.get_program_accounts(None)
+        self.rc.accounts.calculate_capitalization(&self.ancestors)
-            .into_iter()
+    }
-            .map(|(_pubkey, account)| {
+
-                let is_specially_retained = solana_sdk::native_loader::check_id(&account.owner)
+    pub fn calculate_and_verify_capitalization(&self) -> bool {
-                    || solana_sdk::sysvar::check_id(&account.owner);
+        let calculated = self.calculate_capitalization();
-
+        let expected = self.capitalization();
-                if is_specially_retained {
+        if calculated == expected {
-                    // specially retained accounts are ensured to exist by
+            true
-                    // always having a balance of 1 lamports, which is
+        } else {
-                    // outside the capitalization calculation.
+            warn!(
-                    account.lamports - 1
+                "Capitalization mismatch: calculated: {} != expected: {}",
-                } else {
+                calculated, expected
-                    account.lamports
+            );
+            false
         }
-            })
-            .sum()
     }
 
     /// Forcibly overwrites current capitalization by actually recalculating accounts' balances.
@@ -2937,10 +2960,13 @@ impl Bank {
     }
 
     pub fn update_accounts_hash(&self) -> Hash {
-        self.rc
+        let (hash, total_lamports) = self
+            .rc
             .accounts
             .accounts_db
-            .update_accounts_hash(self.slot(), &self.ancestors)
+            .update_accounts_hash(self.slot(), &self.ancestors);
+        assert_eq!(total_lamports, self.capitalization());
+        hash
     }
 
     /// A snapshot bank should be purged of 0 lamport accounts which are not part of the hash
@@ -3900,6 +3926,8 @@ mod tests {
 
     #[test]
     fn test_rent_distribution() {
+        solana_logger::setup();
+
         let bootstrap_validator_pubkey = Pubkey::new_rand();
         let bootstrap_validator_stake_lamports = 30;
         let mut genesis_config = create_genesis_config_with_leader(
@@ -4035,11 +4063,11 @@ mod tests {
 
         let payer = Keypair::new();
         let payer_account = Account::new(400, 0, &system_program::id());
-        bank.store_account(&payer.pubkey(), &payer_account);
+        bank.add_account_and_update_capitalization(&payer.pubkey(), &payer_account);
 
         let payee = Keypair::new();
         let payee_account = Account::new(70, 1, &system_program::id());
-        bank.store_account(&payee.pubkey(), &payee_account);
+        bank.add_account_and_update_capitalization(&payee.pubkey(), &payee_account);
 
         let bootstrap_validator_initial_balance = bank.get_balance(&bootstrap_validator_pubkey);
 
@@ -4120,6 +4148,8 @@ mod tests {
             previous_capitalization - current_capitalization,
             burned_portion
         );
+        bank.freeze();
+        assert!(bank.calculate_and_verify_capitalization());
     }
 
     #[test]
@@ -5068,6 +5098,8 @@ mod tests {
 
     #[test]
     fn test_bank_update_rewards() {
+        solana_logger::setup();
+
         // create a bank that ticks really slowly...
         let bank = Arc::new(Bank::new(&GenesisConfig {
             accounts: (0..42)
@@ -5104,7 +5136,7 @@ mod tests {
             crate::stakes::tests::create_staked_node_accounts(1_0000);
 
         // set up accounts
-        bank.store_account(&stake_id, &stake_account);
+        bank.add_account_and_update_capitalization(&stake_id, &stake_account);
 
         // generate some rewards
         let mut vote_state = Some(VoteState::from(&vote_account).unwrap());
@@ -5114,7 +5146,7 @@ mod tests {
             }
             let versioned = VoteStateVersions::Current(Box::new(vote_state.take().unwrap()));
             VoteState::to(&versioned, &mut vote_account).unwrap();
-            bank.store_account(&vote_id, &vote_account);
+            bank.add_account_and_update_capitalization(&vote_id, &vote_account);
             match versioned {
                 VoteStateVersions::Current(v) => {
                     vote_state = Some(*v);
@@ -5122,7 +5154,7 @@ mod tests {
                 _ => panic!("Has to be of type Current"),
             };
         }
-        bank.store_account(&vote_id, &vote_account);
+        bank.add_account_and_update_capitalization(&vote_id, &vote_account);
 
         let validator_points: u128 = bank
             .stake_delegation_accounts()
@@ -5177,6 +5209,8 @@ mod tests {
                 (rewards.validator_point_value * validator_points as f64) as i64
             )])
         );
+        bank1.freeze();
+        assert!(bank1.calculate_and_verify_capitalization());
     }
 
     fn do_test_bank_update_rewards_determinism() -> u64 {
@@ -5218,8 +5252,8 @@ mod tests {
         let (stake_id2, stake_account2) = crate::stakes::tests::create_stake_account(456, &vote_id);
 
         // set up accounts
-        bank.store_account(&stake_id1, &stake_account1);
+        bank.add_account_and_update_capitalization(&stake_id1, &stake_account1);
-        bank.store_account(&stake_id2, &stake_account2);
+        bank.add_account_and_update_capitalization(&stake_id2, &stake_account2);
 
         // generate some rewards
         let mut vote_state = Some(VoteState::from(&vote_account).unwrap());
@@ -5229,7 +5263,7 @@ mod tests {
             }
             let versioned = VoteStateVersions::Current(Box::new(vote_state.take().unwrap()));
             VoteState::to(&versioned, &mut vote_account).unwrap();
-            bank.store_account(&vote_id, &vote_account);
+            bank.add_account_and_update_capitalization(&vote_id, &vote_account);
             match versioned {
                 VoteStateVersions::Current(v) => {
                     vote_state = Some(*v);
@@ -5237,7 +5271,7 @@ mod tests {
                 _ => panic!("Has to be of type Current"),
             };
         }
-        bank.store_account(&vote_id, &vote_account);
+        bank.add_account_and_update_capitalization(&vote_id, &vote_account);
 
         // put a child bank in epoch 1, which calls update_rewards()...
         let bank1 = Bank::new_from_parent(
@@ -5248,11 +5282,15 @@ mod tests {
         // verify that there's inflation
         assert_ne!(bank1.capitalization(), bank.capitalization());
 
+        bank1.freeze();
+        assert!(bank1.calculate_and_verify_capitalization());
         bank1.capitalization()
     }
 
     #[test]
     fn test_bank_update_rewards_determinism() {
+        solana_logger::setup();
+
         // The same reward should be distributed given same credits
         let expected_capitalization = do_test_bank_update_rewards_determinism();
         // Repeat somewhat large number of iterations to expose possible different behavior
@@ -6860,17 +6898,17 @@ mod tests {
         let parent = Arc::new(Bank::new(&genesis_config));
         parent.lazy_rent_collection.store(true, Ordering::Relaxed);
 
-        let genesis_accounts: Vec<_> = parent.get_program_accounts(None);
+        let genesis_accounts: Vec<_> = parent.get_all_accounts_with_modified_slots();
         assert!(
             genesis_accounts
                 .iter()
-                .any(|(pubkey, _)| *pubkey == mint_keypair.pubkey()),
+                .any(|(pubkey, _, _)| *pubkey == mint_keypair.pubkey()),
             "mint pubkey not found"
         );
         assert!(
             genesis_accounts
                 .iter()
-                .any(|(pubkey, _)| solana_sdk::sysvar::is_sysvar_id(pubkey)),
+                .any(|(pubkey, _, _)| solana_sdk::sysvar::is_sysvar_id(pubkey)),
             "no sysvars found"
         );
 
@@ -6888,11 +6926,11 @@ mod tests {
         let bank1 = Arc::new(new_from_parent(&bank0));
         bank1.squash();
         assert_eq!(
-            bank0.get_program_accounts(Some(&program_id)),
+            bank0.get_program_accounts(&program_id),
             vec![(pubkey0, account0.clone())]
         );
         assert_eq!(
-            bank1.get_program_accounts(Some(&program_id)),
+            bank1.get_program_accounts(&program_id),
             vec![(pubkey0, account0)]
         );
         assert_eq!(
@@ -6911,8 +6949,8 @@ mod tests {
 
         let bank3 = Arc::new(new_from_parent(&bank2));
         bank3.squash();
-        assert_eq!(bank1.get_program_accounts(Some(&program_id)).len(), 2);
+        assert_eq!(bank1.get_program_accounts(&program_id).len(), 2);
-        assert_eq!(bank3.get_program_accounts(Some(&program_id)).len(), 2);
+        assert_eq!(bank3.get_program_accounts(&program_id).len(), 2);
     }
 
     #[test]

sdk/src/genesis_config.rs

@@ -202,10 +202,6 @@ impl GenesisConfig {
         self.native_instruction_processors.push((name, program_id));
     }
 
-    pub fn add_rewards_pool(&mut self, pubkey: Pubkey, account: Account) {
-        self.rewards_pools.insert(pubkey, account);
-    }
-
     pub fn hashes_per_tick(&self) -> Option<u64> {
         self.poh_config.hashes_per_tick
     }
@@ -244,6 +240,8 @@ impl fmt::Display for GenesisConfig {
              {:?}\n\
              {:?}\n\
              Capitalization: {} SOL in {} accounts\n\
+             Native instruction processors: {:#?}\n\
+             Rewards pool: {:#?}\n\
              ",
             Utc.timestamp(self.creation_time, 0).to_rfc3339(),
             self.cluster_type,
@@ -272,6 +270,8 @@ impl fmt::Display for GenesisConfig {
                     .sum::<u64>()
             ),
             self.accounts.len(),
+            self.native_instruction_processors,
+            self.rewards_pools,
         )
     }
 }