[move] Module publishing improvements (#15654)

- `init_module` can no longer publish modules
- new feature flag to gate this change

aptos-move/aptos-release-builder/src/components/feature_flags.rs

@@ -133,6 +133,7 @@ pub enum FeatureFlag {
     CollectionOwner,
     NativeMemoryOperations,
     EnableLoaderV2,
+    DisallowInitModuleToPublishModules,
 }
 
 fn generate_features_blob(writer: &CodeWriter, data: &[u64]) {
@@ -353,6 +354,9 @@ impl From<FeatureFlag> for AptosFeatureFlag {
             FeatureFlag::CollectionOwner => AptosFeatureFlag::COLLECTION_OWNER,
             FeatureFlag::NativeMemoryOperations => AptosFeatureFlag::NATIVE_MEMORY_OPERATIONS,
             FeatureFlag::EnableLoaderV2 => AptosFeatureFlag::ENABLE_LOADER_V2,
+            FeatureFlag::DisallowInitModuleToPublishModules => {
+                AptosFeatureFlag::DISALLOW_INIT_MODULE_TO_PUBLISH_MODULES
+            },
         }
     }
 }
@@ -500,6 +504,9 @@ impl From<AptosFeatureFlag> for FeatureFlag {
             AptosFeatureFlag::COLLECTION_OWNER => FeatureFlag::CollectionOwner,
             AptosFeatureFlag::NATIVE_MEMORY_OPERATIONS => FeatureFlag::NativeMemoryOperations,
             AptosFeatureFlag::ENABLE_LOADER_V2 => FeatureFlag::EnableLoaderV2,
+            AptosFeatureFlag::DISALLOW_INIT_MODULE_TO_PUBLISH_MODULES => {
+                FeatureFlag::DisallowInitModuleToPublishModules
+            },
         }
     }
 }

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

@@ -1511,7 +1511,12 @@ impl AptosVM {
         new_published_modules_loaded: &mut bool,
         change_set_configs: &ChangeSetConfigs,
     ) -> Result<UserSessionChangeSet, VMStatus> {
-        let maybe_publish_request = session.execute(|session| session.extract_publish_request());
+        let maybe_publish_request = session.execute(|session| {
+            let disable_publishing = self
+                .features()
+                .is_disallow_init_module_to_publish_modules_enabled();
+            session.extract_publish_request(disable_publishing)
+        });
         if maybe_publish_request.is_none() {
             let user_change_set = session.finish(change_set_configs, module_storage)?;
 

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

@@ -101,7 +101,7 @@ impl<'r, 'l> SessionExt<'r, 'l> {
             chain_id.id(),
             maybe_user_transaction_context,
         ));
-        extensions.add(NativeCodeContext::default());
+        extensions.add(NativeCodeContext::new());
         extensions.add(NativeStateStorageContext::new(resolver));
         extensions.add(NativeEventContext::default());
         extensions.add(NativeObjectContext::default());
@@ -197,9 +197,11 @@ impl<'r, 'l> SessionExt<'r, 'l> {
         Ok((change_set, module_write_set))
     }
 
-    pub fn extract_publish_request(&mut self) -> Option<PublishRequest> {
+    /// Returns the publish request if it exists. If the provided flag is set to true, disables any
+    /// subsequent module publish requests.
+    pub fn extract_publish_request(&mut self, disable: bool) -> Option<PublishRequest> {
         let ctx = self.get_native_extensions().get_mut::<NativeCodeContext>();
-        ctx.requested_module_bundle.take()
+        ctx.extract_publish_request(disable)
     }
 
     fn populate_v0_resource_group_change_set(

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

@@ -203,7 +203,7 @@ fn unit_test_extensions_hook(exts: &mut NativeContextExtensions) {
     use aptos_table_natives::NativeTableContext;
 
     exts.add(NativeTableContext::new([0u8; 32], &*DUMMY_RESOLVER));
-    exts.add(NativeCodeContext::default());
+    exts.add(NativeCodeContext::new());
     exts.add(NativeTransactionContext::new(
         vec![1],
         vec![1],

aptos-move/e2e-move-tests/src/tests/code_publishing.data/pack_init_module_code_publish/Move.toml

@@ -0,0 +1,7 @@
+[package]
+name = "test_package"
+version = "0.0.0"
+upgrade_policy = "immutable"
+
+[dependencies]
+AptosFramework = { local = "../../../../../framework/aptos-framework" }

aptos-move/e2e-move-tests/src/tests/code_publishing.data/pack_init_module_code_publish/sources/test.move

@@ -0,0 +1,15 @@
+module 0xcafe::test {
+    use aptos_framework::code;
+
+    fun init_module(s: &signer) {
+        // The following metadata and code corresponds to an immutable package called `Package` with compatibility
+        // checks. Code:
+        //   module 0xcafe::m {
+        //       public fun f() {}
+        //   }
+        let metadata: vector<u8> = vector[7, 80, 97, 99, 107, 97, 103, 101, 1, 0, 0, 0, 0, 0, 0, 0, 0, 64, 68, 56, 49, 69, 55, 68, 70, 69, 70, 54, 51, 52, 66, 50, 56, 56, 49, 69, 48, 48, 51, 69, 67, 70, 49, 54, 66, 54, 66, 69, 53, 53, 66, 69, 57, 49, 54, 54, 55, 53, 65, 65, 66, 66, 50, 67, 57, 52, 70, 55, 56, 52, 54, 67, 56, 70, 57, 55, 68, 49, 50, 57, 54, 65, 107, 31, 139, 8, 0, 0, 0, 0, 0, 2, 255, 37, 138, 203, 9, 128, 48, 16, 68, 239, 91, 133, 164, 0, 177, 1, 123, 240, 30, 68, 214, 236, 32, 193, 124, 150, 68, 5, 187, 55, 65, 230, 52, 239, 61, 171, 236, 78, 62, 176, 82, 226, 136, 97, 30, 204, 242, 3, 67, 15, 74, 245, 57, 117, 54, 141, 109, 134, 110, 61, 10, 11, 54, 205, 193, 187, 183, 11, 151, 163, 242, 229, 247, 208, 122, 203, 34, 5, 181, 162, 174, 68, 86, 160, 72, 130, 228, 124, 255, 31, 84, 65, 171, 55, 103, 0, 0, 0, 1, 1, 109, 0, 0, 0, 0, 0];
+        let code: vector<vector<u8>> = vector[vector[161, 28, 235, 11, 7, 0, 0, 10, 6, 1, 0, 2, 3, 2, 6, 5, 8, 1, 7, 9, 4, 8, 13, 32, 12, 45, 7, 0, 0, 0, 1, 0, 0, 0, 1, 0, 1, 109, 1, 102, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 202, 254, 0, 1, 0, 0, 0, 1, 2, 0]];
+
+        code::publish_package_txn(s, metadata, code)
+    }
+}

aptos-move/e2e-move-tests/src/tests/code_publishing.rs

@@ -10,8 +10,8 @@ use aptos_package_builder::PackageBuilder;
 use aptos_types::{
     account_address::{create_resource_address, AccountAddress},
     move_utils::MemberId,
-    on_chain_config::FeatureFlag,
-    transaction::{ExecutionStatus, TransactionPayload},
+    on_chain_config::{FeatureFlag, OnChainConfig},
+    transaction::{ExecutionStatus, TransactionPayload, TransactionStatus},
 };
 use claims::assert_ok;
 use move_core_types::{
@@ -22,7 +22,7 @@ use move_core_types::{
 };
 use rstest::rstest;
 use serde::{Deserialize, Serialize};
-use std::str::FromStr;
+use std::{collections::BTreeSet, str::FromStr};
 
 // Note: this module uses parameterized tests via the
 // [`rstest` crate](https://crates.io/crates/rstest)
@@ -564,20 +564,7 @@ fn test_module_publishing_does_not_leak_speculative_information() {
         let status = output.status().clone();
         match maybe_abort_code {
             Some(abort_code) => {
-                let status = assert_ok!(status.as_kept_status());
-                if let ExecutionStatus::MoveAbort {
-                    location: _,
-                    code,
-                    info: _,
-                } = status
-                {
-                    assert_eq!(code, abort_code);
-                } else {
-                    panic!(
-                        "Transaction should succeed with abort code {}, got {:?}",
-                        abort_code, status
-                    );
-                }
+                assert_move_abort(status, abort_code);
             },
             None => {
                 assert_success!(status);
@@ -588,3 +575,69 @@ fn test_module_publishing_does_not_leak_speculative_information() {
         }
     }
 }
+
+fn assert_move_abort(status: TransactionStatus, expected_abort_code: u64) {
+    let status = assert_ok!(status.as_kept_status());
+    if let ExecutionStatus::MoveAbort {
+        location: _,
+        code,
+        info: _,
+    } = status
+    {
+        assert_eq!(code, expected_abort_code);
+    } else {
+        panic!(
+            "Transaction should succeed with abort code {}, got {:?}",
+            expected_abort_code, status
+        );
+    }
+}
+
+#[test]
+fn test_init_module_should_not_publish_modules() {
+    let mut h = MoveHarness::new_with_features(
+        vec![FeatureFlag::DISALLOW_INIT_MODULE_TO_PUBLISH_MODULES],
+        vec![],
+    );
+    let addr = AccountAddress::from_hex_literal("0xcafe").unwrap();
+    let account = h.new_account_at(addr);
+
+    let txn = h.create_publish_package_cache_building(
+        &account,
+        &common::test_dir_path("code_publishing.data/pack_init_module_code_publish"),
+        |_| {},
+    );
+    let output = h.run_block_get_output(vec![txn]).pop().unwrap();
+    // The abort code corresponds to EALREADY_REQUESTED.
+    assert_move_abort(output.status().clone(), 0x03_0000);
+}
+
+#[test]
+fn test_init_module_can_publish_modules() {
+    let mut h = MoveHarness::new_with_features(vec![], vec![
+        FeatureFlag::DISALLOW_INIT_MODULE_TO_PUBLISH_MODULES,
+    ]);
+    let addr = AccountAddress::from_hex_literal("0xcafe").unwrap();
+    let account = h.new_account_at(addr);
+
+    let txn = h.create_publish_package_cache_building(
+        &account,
+        &common::test_dir_path("code_publishing.data/pack_init_module_code_publish"),
+        |_| {},
+    );
+    let output = h.run_block_get_output(vec![txn]).pop().unwrap();
+    assert_success!(output.status().clone());
+
+    let registry: PackageRegistry = h
+        .read_resource(&addr, PackageRegistry::struct_tag())
+        .unwrap();
+    assert_eq!(registry.packages.len(), 2);
+    let mut package_names = BTreeSet::new();
+    for package in registry.packages {
+        package_names.insert(package.name);
+    }
+    // The package metadata containing init_module exists.
+    assert!(package_names.remove("test_package"));
+    // The package metadata published by init_module call also exists.
+    assert!(package_names.remove("Package"));
+}

aptos-move/framework/src/natives/code.rs

@@ -192,11 +192,36 @@ const EALREADY_REQUESTED: u64 = 0x03_0000;
 const ARBITRARY_POLICY: u8 = 0;
 
 /// The native code context.
-#[derive(Tid, Default)]
+#[derive(Tid)]
 pub struct NativeCodeContext {
-    /// Remembers whether the publishing of a module bundle was requested during transaction
-    /// execution.
-    pub requested_module_bundle: Option<PublishRequest>,
+    /// If false, publish requests are ignored and any attempts to publish code result in runtime
+    /// errors.
+    enabled: bool,
+    /// Possibly stores (if not [None]) the request to publish a module bundle. The request is made
+    /// using the native code defined in this context. It is later extracted by the VM for further
+    /// checks and processing the actual publish.
+    requested_module_bundle: Option<PublishRequest>,
+}
+
+impl NativeCodeContext {
+    #[allow(clippy::new_without_default)]
+    pub fn new() -> Self {
+        Self {
+            enabled: true,
+            requested_module_bundle: None,
+        }
+    }
+
+    pub fn extract_publish_request(&mut self, disable: bool) -> Option<PublishRequest> {
+        if !self.enabled {
+            return None;
+        }
+
+        if disable {
+            self.enabled = false;
+        }
+        self.requested_module_bundle.take()
+    }
 }
 
 /// Represents a request for code publishing made from a native call and to be processed
@@ -316,8 +341,8 @@ fn native_request_publish(
     });
 
     let code_context = context.extensions_mut().get_mut::<NativeCodeContext>();
-    if code_context.requested_module_bundle.is_some() {
-        // Can't request second time.
+    if code_context.requested_module_bundle.is_some() || !code_context.enabled {
+        // Can't request second time or if publish requests are not allowed.
         return Err(SafeNativeError::Abort {
             abort_code: EALREADY_REQUESTED,
         });
@@ -329,7 +354,7 @@ fn native_request_publish(
         allowed_deps,
         check_compat: policy != ARBITRARY_POLICY,
     });
-    // TODO(Gas): charge gas for requesting code load (charge for actual code loading done elsewhere)
+
     Ok(smallvec![])
 }
 

types/src/on_chain_config/aptos_features.rs

@@ -100,6 +100,10 @@ pub enum FeatureFlag {
     /// AIP-105 (https://github.com/aptos-foundation/AIPs/blob/main/aips/aip-105.md)
     NATIVE_MEMORY_OPERATIONS = 80,
     ENABLE_LOADER_V2 = 81,
+    /// Prior to this feature flag, it was possible to attempt 'init_module' to publish modules
+    /// that results in a new package created but without any code. With this feature, it is no
+    /// longer possible and an explicit error is returned if publishing is attempted.
+    DISALLOW_INIT_MODULE_TO_PUBLISH_MODULES = 82,
 }
 
 impl FeatureFlag {
@@ -181,6 +185,7 @@ impl FeatureFlag {
             FeatureFlag::NATIVE_MEMORY_OPERATIONS,
             FeatureFlag::COLLECTION_OWNER,
             FeatureFlag::ENABLE_LOADER_V2,
+            FeatureFlag::DISALLOW_INIT_MODULE_TO_PUBLISH_MODULES,
         ]
     }
 }
@@ -329,6 +334,10 @@ impl Features {
         self.is_enabled(FeatureFlag::ENABLE_LOADER_V2)
     }
 
+    pub fn is_disallow_init_module_to_publish_modules_enabled(&self) -> bool {
+        self.is_enabled(FeatureFlag::DISALLOW_INIT_MODULE_TO_PUBLISH_MODULES)
+    }
+
     pub fn get_max_identifier_size(&self) -> u64 {
         if self.is_enabled(FeatureFlag::LIMIT_MAX_IDENTIFIER_LENGTH) {
             IDENTIFIER_SIZE_MAX