[aptos-vm] V2 loader uses the same publishing checks now

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

@@ -1665,6 +1665,16 @@ impl AptosVM {
                 let compatability_checks =
                     Compatibility::new(check_struct_layout, check_friend_linking);
 
+                // TODO(loader_v2): Make sure to `validate_publish_request` passes to environment
+                //                  verifier extension in V2 design. This works for now for tests.
+                self.validate_publish_request(
+                    session,
+                    module_storage,
+                    modules,
+                    expected_modules,
+                    allowed_deps,
+                )?;
+
                 if self.features().is_loader_v2_enabled() {
                     // Create a temporary storage. If this fails, it means publishing
                     // is not possible. We use a new VM here so that struct index map
@@ -1674,8 +1684,6 @@ impl AptosVM {
                     let tmp_vm = self.move_vm.clone();
                     let tmp_module_storage = TemporaryModuleStorage::new_with_compat_config(
                         &destination,
-                        // TODO(loader_v2): Make sure to `validate_publish_request` passes to environment
-                        //                  verifier extension.
                         tmp_vm.runtime_environment(),
                         compatability_checks,
                         module_storage,
@@ -1739,14 +1747,6 @@ impl AptosVM {
 
                     Ok(Some((module_write_set, init_module_changes)))
                 } else {
-                    // Validate the module bundle
-                    self.validate_publish_request(
-                        session,
-                        modules,
-                        expected_modules,
-                        allowed_deps,
-                    )?;
-
                     // Check what modules exist before publishing.
                     let mut exists = BTreeSet::new();
                     for m in modules {
@@ -1798,6 +1798,7 @@ impl AptosVM {
     fn validate_publish_request(
         &self,
         session: &mut SessionExt,
+        module_storage: &impl AptosModuleStorage,
         modules: &[CompiledModule],
         mut expected_modules: BTreeSet<String>,
         allowed_deps: Option<BTreeMap<AccountAddress, BTreeSet<String>>>,
@@ -1842,13 +1843,22 @@ impl AptosVM {
             aptos_framework::verify_module_metadata(m, self.features(), self.timed_features())
                 .map_err(|err| Self::metadata_validation_error(&err.to_string()))?;
         }
+
+        let use_loader_v2 = self.features().is_enabled(FeatureFlag::ENABLE_LOADER_V2);
         verifier::resource_groups::validate_resource_groups(
             session,
+            module_storage,
             modules,
             self.features()
                 .is_enabled(FeatureFlag::SAFER_RESOURCE_GROUPS),
+            use_loader_v2,
+        )?;
+        verifier::event_validation::validate_module_events(
+            session,
+            module_storage,
+            modules,
+            use_loader_v2,
         )?;
-        verifier::event_validation::validate_module_events(session, modules)?;
 
         if !expected_modules.is_empty() {
             return Err(Self::metadata_validation_error(

aptos-move/aptos-vm/src/verifier/event_validation.rs

@@ -3,6 +3,7 @@
 
 use crate::move_vm_ext::SessionExt;
 use aptos_framework::RuntimeModuleMetadataV1;
+use aptos_vm_types::module_and_script_storage::module_storage::AptosModuleStorage;
 use move_binary_format::{
     access::{ModuleAccess, ScriptAccess},
     errors::{Location, PartialVMError, VMError, VMResult},
@@ -35,7 +36,9 @@ fn metadata_validation_error(msg: &str) -> VMError {
 /// * Verify all changes are compatible upgrades (existing event attributes cannot be removed)
 pub(crate) fn validate_module_events(
     session: &mut SessionExt,
+    module_storage: &impl AptosModuleStorage,
     modules: &[CompiledModule],
+    use_loader_v2: bool,
 ) -> VMResult<()> {
     for module in modules {
         let mut new_event_structs =
@@ -48,8 +51,12 @@ pub(crate) fn validate_module_events(
         // Check all the emit calls have the correct struct with event attribute.
         validate_emit_calls(&new_event_structs, module)?;
 
-        let original_event_structs =
-            extract_event_metadata_from_module(session, &module.self_id())?;
+        let original_event_structs = extract_event_metadata_from_module(
+            session,
+            module_storage,
+            &module.self_id(),
+            use_loader_v2,
+        )?;
 
         for member in original_event_structs {
             // Fail if we see a removal of an event attribute.
@@ -118,23 +125,37 @@ pub(crate) fn validate_emit_calls(
 /// Given a module id extract all event metadata
 pub(crate) fn extract_event_metadata_from_module(
     session: &mut SessionExt,
+    module_storage: &impl AptosModuleStorage,
     module_id: &ModuleId,
+    use_loader_v2: bool,
 ) -> VMResult<HashSet<String>> {
-    #[allow(deprecated)]
-    let metadata = session
-        .fetch_module_from_data_store(module_id)
-        .map(|module| {
-            CompiledModule::deserialize_with_config(
-                &module,
-                &session.get_vm_config().deserializer_config,
-            )
-            .map(|module| aptos_framework::get_metadata_from_compiled_module(&module))
-        });
-
-    if let Ok(Ok(Some(metadata))) = metadata {
-        extract_event_metadata(&metadata)
+    if use_loader_v2 {
+        // TODO(loader_v2): We can optimize metadata calls as well.
+        let metadata = module_storage
+            .fetch_deserialized_module(module_id.address(), module_id.name())
+            .map(|m| aptos_framework::get_metadata_from_compiled_module(m.as_ref()));
+        if let Ok(Some(metadata)) = metadata {
+            extract_event_metadata(&metadata)
+        } else {
+            Ok(HashSet::new())
+        }
     } else {
-        Ok(HashSet::new())
+        #[allow(deprecated)]
+        let metadata = session
+            .fetch_module_from_data_store(module_id)
+            .map(|module| {
+                CompiledModule::deserialize_with_config(
+                    &module,
+                    &session.get_vm_config().deserializer_config,
+                )
+                .map(|module| aptos_framework::get_metadata_from_compiled_module(&module))
+            });
+
+        if let Ok(Ok(Some(metadata))) = metadata {
+            extract_event_metadata(&metadata)
+        } else {
+            Ok(HashSet::new())
+        }
     }
 }
 

aptos-move/aptos-vm/src/verifier/resource_groups.rs

@@ -3,6 +3,7 @@
 
 use crate::move_vm_ext::SessionExt;
 use aptos_framework::{ResourceGroupScope, RuntimeModuleMetadataV1};
+use aptos_vm_types::module_and_script_storage::module_storage::AptosModuleStorage;
 use move_binary_format::{
     access::ModuleAccess,
     errors::{Location, PartialVMError, VMError, VMResult},
@@ -12,7 +13,10 @@ use move_core_types::{
     language_storage::{ModuleId, StructTag},
     vm_status::StatusCode,
 };
-use std::collections::{BTreeMap, BTreeSet};
+use std::{
+    collections::{BTreeMap, BTreeSet},
+    sync::Arc,
+};
 
 fn metadata_validation_err(msg: &str) -> Result<(), VMError> {
     Err(metadata_validation_error(msg))
@@ -32,15 +36,22 @@ fn metadata_validation_error(msg: &str) -> VMError {
 /// * For any new members, verify that they are in a valid resource group
 pub(crate) fn validate_resource_groups(
     session: &mut SessionExt,
+    module_storage: &impl AptosModuleStorage,
     modules: &[CompiledModule],
     safer_resource_groups: bool,
+    use_loader_v2: bool,
 ) -> Result<(), VMError> {
     let mut groups = BTreeMap::new();
     let mut members = BTreeMap::new();
 
     for module in modules {
-        let (new_groups, new_members) =
-            validate_module_and_extract_new_entries(session, module, safer_resource_groups)?;
+        let (new_groups, new_members) = validate_module_and_extract_new_entries(
+            session,
+            module_storage,
+            module,
+            safer_resource_groups,
+            use_loader_v2,
+        )?;
         groups.insert(module.self_id(), new_groups);
         members.insert(module.self_id(), new_members);
     }
@@ -49,8 +60,12 @@ pub(crate) fn validate_resource_groups(
         for value in inner_members.values() {
             let value_module_id = value.module_id();
             if !groups.contains_key(&value_module_id) {
-                let (inner_groups, _, _) =
-                    extract_resource_group_metadata_from_module(session, &value_module_id)?;
+                let (inner_groups, _, _) = extract_resource_group_metadata_from_module(
+                    session,
+                    module_storage,
+                    &value_module_id,
+                    use_loader_v2,
+                )?;
                 groups.insert(value.module_id(), inner_groups);
             }
 
@@ -77,8 +92,10 @@ pub(crate) fn validate_resource_groups(
 /// * Return any new members to validate correctness and all groups to assist in validation
 pub(crate) fn validate_module_and_extract_new_entries(
     session: &mut SessionExt,
+    module_storage: &impl AptosModuleStorage,
     module: &CompiledModule,
     safer_resource_groups: bool,
+    use_loader_v2: bool,
 ) -> VMResult<(
     BTreeMap<String, ResourceGroupScope>,
     BTreeMap<String, StructTag>,
@@ -91,7 +108,12 @@ pub(crate) fn validate_module_and_extract_new_entries(
         };
 
     let (original_groups, original_members, mut structs) =
-        extract_resource_group_metadata_from_module(session, &module.self_id())?;
+        extract_resource_group_metadata_from_module(
+            session,
+            module_storage,
+            &module.self_id(),
+            use_loader_v2,
+        )?;
 
     for (member, value) in original_members {
         // We don't need to re-validate new_members above.
@@ -144,22 +166,16 @@ pub(crate) fn validate_module_and_extract_new_entries(
 /// Given a module id extract all resource group metadata
 pub(crate) fn extract_resource_group_metadata_from_module(
     session: &mut SessionExt,
+    module_storage: &impl AptosModuleStorage,
     module_id: &ModuleId,
+    use_loader_v2: bool,
 ) -> VMResult<(
     BTreeMap<String, ResourceGroupScope>,
     BTreeMap<String, StructTag>,
     BTreeSet<String>,
 )> {
-    #[allow(deprecated)]
-    let module = session
-        .fetch_module_from_data_store(module_id)
-        .map(|module| {
-            CompiledModule::deserialize_with_config(
-                &module,
-                &session.get_vm_config().deserializer_config,
-            )
-        });
-    let (metadata, module) = if let Ok(Ok(module)) = module {
+    let module = fetch_module(session, module_storage, module_id, use_loader_v2);
+    let (metadata, module) = if let Ok(module) = module {
         (
             aptos_framework::get_metadata_from_compiled_module(&module),
             module,
@@ -186,6 +202,26 @@ pub(crate) fn extract_resource_group_metadata_from_module(
     }
 }
 
+fn fetch_module(
+    session: &mut SessionExt,
+    module_storage: &impl AptosModuleStorage,
+    module_id: &ModuleId,
+    use_loader_v2: bool,
+) -> VMResult<Arc<CompiledModule>> {
+    if use_loader_v2 {
+        module_storage.fetch_deserialized_module(module_id.address(), module_id.name())
+    } else {
+        #[allow(deprecated)]
+        let bytes = session.fetch_module_from_data_store(module_id)?;
+        let module = CompiledModule::deserialize_with_config(
+            &bytes,
+            &session.get_vm_config().deserializer_config,
+        )
+        .map_err(|e| e.finish(Location::Undefined))?;
+        Ok(Arc::new(module))
+    }
+}
+
 /// Given a module id extract all resource group metadata
 pub(crate) fn extract_resource_group_metadata(
     metadata: &RuntimeModuleMetadataV1,