[replay-verify] Add an option to re-validate all move modules in the snapshot

aptos-move/aptos-vm/src/move_vm_ext/mod.rs

@@ -10,5 +10,6 @@ mod vm;
 pub use crate::move_vm_ext::{
     resolver::MoveResolverExt,
     session::{SessionExt, SessionId, SessionOutput},
+    vm::verifier_config,
     vm::MoveVmExt,
 };

aptos-move/aptos-vm/src/move_vm_ext/vm.rs

@@ -30,22 +30,17 @@ impl MoveVmExt {
         treat_friend_as_private: bool,
         chain_id: u8,
     ) -> VMResult<Self> {
+        let mut config = verifier_config();
+        config.treat_friend_as_private = treat_friend_as_private;
+
         Ok(Self {
             inner: MoveVM::new_with_configs(
                 aptos_natives(
                     native_gas_params,
                     abs_val_size_gas_params,
                     gas_feature_version,
                 ),
-                VerifierConfig {
-                    max_loop_depth: Some(5),
-                    treat_friend_as_private,
-                    max_generic_instantiation_length: Some(32),
-                    max_function_parameters: Some(128),
-                    max_basic_blocks: Some(1024),
-                    max_value_stack_size: 1024,
-                    max_type_nodes: Some(256),
-                },
+                config,
                 crate::AptosVM::get_runtime_config(),
             )?,
             chain_id,
@@ -96,3 +91,15 @@ impl Deref for MoveVmExt {
         &self.inner
     }
 }
+
+pub fn verifier_config() -> VerifierConfig {
+    VerifierConfig {
+        max_loop_depth: Some(5),
+        treat_friend_as_private: false,
+        max_generic_instantiation_length: Some(32),
+        max_function_parameters: Some(128),
+        max_basic_blocks: Some(1024),
+        max_value_stack_size: 1024,
+        max_type_nodes: Some(256),
+    }
+}

storage/backup/backup-cli/Cargo.toml

@@ -33,6 +33,8 @@ executor-test-helpers = { workspace = true }
 executor-types = { workspace = true }
 futures = { workspace = true }
 itertools = { workspace = true }
+move-binary-format = { workspace = true }
+move-bytecode-verifier = { workspace = true }
 num_cpus = { workspace = true }
 once_cell = { workspace = true }
 pin-project = { workspace = true }

storage/backup/backup-cli/src/backup_types/state_snapshot/restore.rs

@@ -26,13 +26,17 @@ use anyhow::{anyhow, ensure, Result};
 use aptos_infallible::Mutex;
 use aptos_logger::prelude::*;
 use aptos_types::{
+    access_path::Path,
     ledger_info::LedgerInfoWithSignatures,
     proof::TransactionInfoWithProof,
     state_store::{state_key::StateKey, state_value::StateValue},
     transaction::Version,
 };
+use aptos_vm::move_vm_ext::verifier_config;
 use clap::Parser;
 use futures::{stream, TryStreamExt};
+use move_binary_format::CompiledModule;
+use move_bytecode_verifier::verify_module_with_config;
 use std::sync::Arc;
 use storage_interface::StateSnapshotReceiver;
 use tokio::time::Instant;
@@ -43,6 +47,8 @@ pub struct StateSnapshotRestoreOpt {
     pub manifest_handle: FileHandle,
     #[clap(long = "state-into-version")]
     pub version: Version,
+    #[clap(long)]
+    pub validate_modules: bool,
 }
 
 pub struct StateSnapshotRestoreController {
@@ -56,6 +62,7 @@ pub struct StateSnapshotRestoreController {
     target_version: Version,
     epoch_history: Option<Arc<EpochHistory>>,
     concurrent_downloads: usize,
+    validate_modules: bool,
 }
 
 impl StateSnapshotRestoreController {
@@ -73,6 +80,7 @@ impl StateSnapshotRestoreController {
             target_version: global_opt.target_version,
             epoch_history,
             concurrent_downloads: global_opt.concurrent_downloads,
+            validate_modules: opt.validate_modules,
         }
     }
 
@@ -186,6 +194,9 @@ impl StateSnapshotRestoreController {
                 .with_label_values(&["add_state_chunk"])
                 .start_timer();
             let receiver = receiver.clone();
+            if self.validate_modules {
+                Self::validate_modules(&blobs);
+            }
             tokio::task::spawn_blocking(move || {
                 receiver.lock().as_mut().unwrap().add_chunk(blobs, proof)
             })
@@ -207,6 +218,23 @@ impl StateSnapshotRestoreController {
         Ok(())
     }
 
+    fn validate_modules(blob: &[(StateKey, StateValue)]) {
+        let config = verifier_config();
+        for (key, value) in blob {
+            if let StateKey::AccessPath(p) = key {
+                if let Path::Code(module_id) = p.get_path() {
+                    if let Ok(module) = CompiledModule::deserialize(value.bytes()) {
+                        if let Err(err) = verify_module_with_config(&config, &module) {
+                            error!("Module {:?} failed validation: {:?}", module_id, err);
+                        }
+                    } else {
+                        error!("Module {:?} failed to deserialize", module_id);
+                    }
+                }
+            }
+        }
+    }
+
     async fn read_state_value(
         storage: &Arc<dyn BackupStorage>,
         file_handle: FileHandle,

storage/backup/backup-cli/src/backup_types/state_snapshot/tests.rs

@@ -76,6 +76,7 @@ fn end_to_end() {
             StateSnapshotRestoreOpt {
                 manifest_handle,
                 version,
+                validate_modules: false,
             },
             GlobalRestoreOpt {
                 dry_run: false,

storage/backup/backup-cli/src/backup_types/tests.rs

@@ -133,6 +133,7 @@ fn test_end_to_end_impl(d: TestData) {
                 StateSnapshotRestoreOpt {
                     manifest_handle: state_snapshot_manifest.unwrap(),
                     version,
+                    validate_modules: false,
                 },
                 global_restore_opt.clone(),
                 Arc::clone(&store),

storage/backup/backup-cli/src/bin/replay-verify.rs

@@ -46,6 +46,8 @@ struct Opt {
         in the backup). [Defaults to the latest version available] "
     )]
     end_version: Option<Version>,
+    #[clap(long)]
+    validate_modules: bool,
 }
 
 #[tokio::main]
@@ -79,6 +81,7 @@ async fn main_impl() -> Result<()> {
         restore_handler,
         opt.start_version.unwrap_or(0),
         opt.end_version.unwrap_or(Version::MAX),
+        opt.validate_modules,
     )?
     .run()
     .await

storage/backup/backup-cli/src/coordinators/replay_verify.rs

@@ -27,6 +27,7 @@ pub struct ReplayVerifyCoordinator {
     restore_handler: RestoreHandler,
     start_version: Version,
     end_version: Version,
+    validate_modules: bool,
 }
 
 impl ReplayVerifyCoordinator {
@@ -39,6 +40,7 @@ impl ReplayVerifyCoordinator {
         restore_handler: RestoreHandler,
         start_version: Version,
         end_version: Version,
+        validate_modules: bool,
     ) -> Result<Self> {
         Ok(Self {
             storage,
@@ -49,6 +51,7 @@ impl ReplayVerifyCoordinator {
             restore_handler,
             start_version,
             end_version,
+            validate_modules,
         })
     }
 
@@ -112,6 +115,7 @@ impl ReplayVerifyCoordinator {
                 StateSnapshotRestoreOpt {
                     manifest_handle: backup.manifest,
                     version: backup.version,
+                    validate_modules: self.validate_modules,
                 },
                 global_opt.clone(),
                 Arc::clone(&self.storage),

storage/backup/backup-cli/src/coordinators/restore.rs

@@ -170,6 +170,7 @@ impl RestoreCoordinator {
             StateSnapshotRestoreOpt {
                 manifest_handle: state_snapshot_backup.manifest,
                 version,
+                validate_modules: false,
             },
             self.global_opt.clone(),
             Arc::clone(&self.storage),

storage/backup/backup-cli/src/coordinators/verify.rs

@@ -100,6 +100,7 @@ impl VerifyCoordinator {
                 StateSnapshotRestoreOpt {
                     manifest_handle: backup.manifest,
                     version: backup.version,
+                    validate_modules: false,
                 },
                 global_opt.clone(),
                 Arc::clone(&self.storage),