Verifies epoch stakes after rebuilding bank from snapshot (solana-lab…

…s#924)

runtime/src/snapshot_bank_utils.rs

@@ -1,6 +1,7 @@
 use {
     crate::{
         bank::{builtins::BuiltinPrototype, Bank, BankFieldsToDeserialize, BankSlotDelta},
+        epoch_stakes::EpochStakes,
         runtime_config::RuntimeConfig,
         serde_snapshot::{
             bank_from_streams, bank_to_stream, fields_from_streams,
@@ -20,7 +21,7 @@ use {
             verify_and_unarchive_snapshots, verify_unpacked_snapshots_dir_and_version,
             AddBankSnapshotError, ArchiveFormat, BankSnapshotInfo, BankSnapshotKind, SnapshotError,
             SnapshotRootPaths, SnapshotVersion, StorageAndNextAccountsFileId,
-            UnpackedSnapshotsDirAndVersion, VerifySlotDeltasError,
+            UnpackedSnapshotsDirAndVersion, VerifyEpochStakesError, VerifySlotDeltasError,
         },
         status_cache,
     },
@@ -38,18 +39,19 @@ use {
     },
     solana_measure::{measure, measure::Measure},
     solana_sdk::{
-        clock::Slot,
+        clock::{Epoch, Slot},
         feature_set,
         genesis_config::GenesisConfig,
         hash::Hash,
         pubkey::Pubkey,
         slot_history::{Check, SlotHistory},
     },
     std::{
-        collections::HashSet,
+        collections::{HashMap, HashSet},
         fs,
         io::{BufWriter, Write},
         num::NonZeroUsize,
+        ops::RangeInclusive,
         path::{Path, PathBuf},
         sync::{atomic::AtomicBool, Arc},
     },
@@ -758,6 +760,8 @@ fn rebuild_bank_from_unarchived_snapshots(
         )
     })?;
 
+    verify_epoch_stakes(&bank)?;
+
     // The status cache is rebuilt from the latest snapshot.  So, if there's an incremental
     // snapshot, use that.  Otherwise use the full snapshot.
     let status_cache_path = incremental_snapshot_unpacked_snapshots_dir_and_version
@@ -831,6 +835,8 @@ fn rebuild_bank_from_snapshot(
         )?)
     })?;
 
+    verify_epoch_stakes(&bank)?;
+
     let status_cache_path = bank_snapshot
         .snapshot_dir
         .join(snapshot_utils::SNAPSHOT_STATUS_CACHE_FILENAME);
@@ -944,6 +950,51 @@ fn verify_slot_deltas_with_history(
     Ok(())
 }
 
+/// Verifies the bank's epoch stakes are valid after rebuilding from a snapshot
+fn verify_epoch_stakes(bank: &Bank) -> std::result::Result<(), VerifyEpochStakesError> {
+    // Stakes are required for epochs from the current epoch up-to-and-including the
+    // leader schedule epoch.  In practice this will only be two epochs: the current and the next.
+    // Using a range mirrors how Bank::new_with_paths() seeds the initial epoch stakes.
+    let current_epoch = bank.epoch();
+    let leader_schedule_epoch = bank.get_leader_schedule_epoch(bank.slot());
+    let required_epochs = current_epoch..=leader_schedule_epoch;
+    _verify_epoch_stakes(bank.epoch_stakes_map(), required_epochs)
+}
+
+/// Verifies the bank's epoch stakes are valid after rebuilding from a snapshot
+///
+/// This version of the function exists to facilitate testing.
+/// Normal callers should use `verify_epoch_stakes()`.
+fn _verify_epoch_stakes(
+    epoch_stakes_map: &HashMap<Epoch, EpochStakes>,
+    required_epochs: RangeInclusive<Epoch>,
+) -> std::result::Result<(), VerifyEpochStakesError> {
+    // Ensure epoch stakes from the snapshot does not contain entries for invalid epochs.
+    // Since epoch stakes are computed for the leader schedule epoch (usually `epoch + 1`),
+    // the snapshot's epoch stakes therefor can have entries for epochs at-or-below the
+    // leader schedule epoch.
+    let max_epoch = *required_epochs.end();
+    if let Some(invalid_epoch) = epoch_stakes_map.keys().find(|epoch| **epoch > max_epoch) {
+        return Err(VerifyEpochStakesError::EpochGreaterThanMax(
+            *invalid_epoch,
+            max_epoch,
+        ));
+    }
+
+    // Ensure epoch stakes contains stakes for all the required epochs
+    if let Some(missing_epoch) = required_epochs
+        .clone()
+        .find(|epoch| !epoch_stakes_map.contains_key(epoch))
+    {
+        return Err(VerifyEpochStakesError::StakesNotFound(
+            missing_epoch,
+            required_epochs,
+        ));
+    }
+
+    Ok(())
+}
+
 /// Get the snapshot storages for this bank
 pub fn get_snapshot_storages(bank: &Bank) -> Vec<Arc<AccountStorageEntry>> {
     let mut measure_snapshot_storages = Measure::start("snapshot-storages");
@@ -1262,6 +1313,7 @@ mod tests {
     use {
         super::*,
         crate::{
+            bank::tests::create_simple_test_bank,
             bank_forks::BankForks,
             genesis_utils,
             snapshot_config::SnapshotConfig,
@@ -2647,6 +2699,55 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_verify_epoch_stakes_good() {
+        let bank = create_simple_test_bank(100 * LAMPORTS_PER_SOL);
+        assert_eq!(verify_epoch_stakes(&bank), Ok(()));
+    }
+
+    #[test]
+    fn test_verify_epoch_stakes_bad() {
+        let bank = create_simple_test_bank(100 * LAMPORTS_PER_SOL);
+        let current_epoch = bank.epoch();
+        let leader_schedule_epoch = bank.get_leader_schedule_epoch(bank.slot());
+        let required_epochs = current_epoch..=leader_schedule_epoch;
+
+        // insert an invalid epoch into the epoch stakes
+        {
+            let mut epoch_stakes_map = bank.epoch_stakes_map().clone();
+            let invalid_epoch = *required_epochs.end() + 1;
+            epoch_stakes_map.insert(
+                invalid_epoch,
+                bank.epoch_stakes(bank.epoch()).cloned().unwrap(),
+            );
+
+            assert_eq!(
+                _verify_epoch_stakes(&epoch_stakes_map, required_epochs.clone()),
+                Err(VerifyEpochStakesError::EpochGreaterThanMax(
+                    invalid_epoch,
+                    *required_epochs.end(),
+                )),
+            );
+        }
+
+        // remove required stakes
+        {
+            for removed_epoch in required_epochs.clone() {
+                let mut epoch_stakes_map = bank.epoch_stakes_map().clone();
+                let removed_stakes = epoch_stakes_map.remove(&removed_epoch);
+                assert!(removed_stakes.is_some());
+
+                assert_eq!(
+                    _verify_epoch_stakes(&epoch_stakes_map, required_epochs.clone()),
+                    Err(VerifyEpochStakesError::StakesNotFound(
+                        removed_epoch,
+                        required_epochs.clone(),
+                    )),
+                );
+            }
+        }
+    }
+
     #[test]
     fn test_get_highest_loadable_bank_snapshot() {
         let bank_snapshots_dir = TempDir::new().unwrap();

runtime/src/snapshot_utils.rs

@@ -26,13 +26,17 @@ use {
         utils::{move_and_async_delete_path, ACCOUNTS_RUN_DIR, ACCOUNTS_SNAPSHOT_DIR},
     },
     solana_measure::{measure, measure::Measure},
-    solana_sdk::{clock::Slot, hash::Hash},
+    solana_sdk::{
+        clock::{Epoch, Slot},
+        hash::Hash,
+    },
     std::{
         cmp::Ordering,
         collections::{HashMap, HashSet},
         fmt, fs,
         io::{BufReader, BufWriter, Error as IoError, Read, Result as IoResult, Seek, Write},
         num::NonZeroUsize,
+        ops::RangeInclusive,
         path::{Path, PathBuf},
         process::ExitStatus,
         str::FromStr,
@@ -332,6 +336,9 @@ pub enum SnapshotError {
     #[error("snapshot slot deltas are invalid: {0}")]
     VerifySlotDeltas(#[from] VerifySlotDeltasError),
 
+    #[error("snapshot epoch stakes are invalid: {0}")]
+    VerifyEpochStakes(#[from] VerifyEpochStakesError),
+
     #[error("bank_snapshot_info new_from_dir failed: {0}")]
     NewFromDir(#[from] SnapshotNewFromDirError),
 
@@ -408,6 +415,16 @@ pub enum VerifySlotDeltasError {
     BadSlotHistory,
 }
 
+/// Errors that can happen in `verify_epoch_stakes()`
+#[derive(Error, Debug, PartialEq, Eq)]
+pub enum VerifyEpochStakesError {
+    #[error("epoch {0} is greater than the max {1}")]
+    EpochGreaterThanMax(Epoch, Epoch),
+
+    #[error("stakes not found for epoch {0} (required epochs: {1:?})")]
+    StakesNotFound(Epoch, RangeInclusive<Epoch>),
+}
+
 /// Errors that can happen in `add_bank_snapshot()`
 #[derive(Error, Debug)]
 pub enum AddBankSnapshotError {