[move] Re-enabled bytecode verification of test code (MystenLabs#7792)

This is a new version of MystenLabs#6564 and
it also closes MystenLabs#7759.

Bytecode verification of of test code was introduced in
MystenLabs#5732 but then temporarily
disabled in MystenLabs#6435 due to developer
complaints about inability to test `init` functions.

This PR special-cases the bytecode verifier to avoid running checks
related to calling `init` functions and to instantiating a
one-time-witness struct by-hand in the test code.

It also enables bytecode verification when running Move unit tests.

crates/sui-adapter/src/adapter.rs

@@ -504,7 +504,7 @@ pub fn verify_and_link<
     // run the Sui verifier
     for module in modules.iter() {
         // Run Sui bytecode verifier, which runs some additional checks that assume the Move bytecode verifier has passed.
-        verifier::verify_module(module)?;
+        verifier::verify_module(module, &BTreeMap::new())?;
     }
     Ok(vm)
 }

crates/sui-framework-build/src/compiled_package.rs

@@ -2,7 +2,8 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use std::{
-    collections::{BTreeSet, HashSet},
+    collections::{BTreeMap, BTreeSet, HashSet},
+    io::Write,
     path::PathBuf,
 };
 
@@ -13,18 +14,29 @@ use move_binary_format::{
     CompiledModule,
 };
 use move_bytecode_utils::{layout::SerdeLayoutBuilder, module_cache::GetModule, Modules};
-use move_compiler::compiled_unit::{CompiledUnitEnum, NamedCompiledModule};
+use move_compiler::{
+    compiled_unit::{
+        AnnotatedCompiledModule, AnnotatedCompiledScript, CompiledUnitEnum, NamedCompiledModule,
+    },
+    diagnostics::{report_diagnostics_to_color_buffer, report_warnings},
+    expansion::ast::{AttributeName_, Attributes},
+    shared::known_attributes::KnownAttribute,
+};
 use move_core_types::{
     account_address::AccountAddress,
     language_storage::{ModuleId, StructTag, TypeTag},
 };
 use move_package::{
-    compilation::compiled_package::CompiledPackage as MoveCompiledPackage,
+    compilation::{
+        build_plan::BuildPlan, compiled_package::CompiledPackage as MoveCompiledPackage,
+    },
+    resolution::resolution_graph::ResolvedGraph,
     BuildConfig as MoveBuildConfig,
 };
 use serde_reflection::Registry;
 use sui_types::{
     error::{SuiError, SuiResult},
+    move_package::{FnInfo, FnInfoKey, FnInfoMap},
     MOVE_STDLIB_ADDRESS, SUI_FRAMEWORK_ADDRESS,
 };
 use sui_verifier::verifier as sui_bytecode_verifier;
@@ -54,25 +66,90 @@ impl BuildConfig {
         build_config
     }
 
+    fn is_test(attributes: &Attributes) -> bool {
+        attributes
+            .iter()
+            .any(|(_, name, _)| matches!(name, AttributeName_::Known(KnownAttribute::Testing(_))))
+    }
+
+    fn fn_info(
+        units: &[CompiledUnitEnum<AnnotatedCompiledModule, AnnotatedCompiledScript>],
+    ) -> FnInfoMap {
+        let mut fn_info_map = BTreeMap::new();
+        for u in units {
+            match u {
+                CompiledUnitEnum::Module(m) => {
+                    let mod_addr = m.named_module.address.into_inner();
+                    for (_, s, info) in &m.function_infos {
+                        let fn_name = s.as_str().to_string();
+                        let is_test = Self::is_test(&info.attributes);
+                        fn_info_map.insert(FnInfoKey { fn_name, mod_addr }, FnInfo { is_test });
+                    }
+                }
+                CompiledUnitEnum::Script(_) => continue,
+            }
+        }
+
+        fn_info_map
+    }
+
+    fn compile_package<W: Write>(
+        resolution_graph: ResolvedGraph,
+        writer: &mut W,
+    ) -> anyhow::Result<(MoveCompiledPackage, FnInfoMap)> {
+        let build_plan = BuildPlan::create(resolution_graph)?;
+        let mut fn_info = None;
+        let compiled_pkg = build_plan.compile_with_driver(writer, |compiler| {
+            let (files, units_res) = compiler.build()?;
+            match units_res {
+                Ok((units, warning_diags)) => {
+                    report_warnings(&files, warning_diags);
+                    fn_info = Some(Self::fn_info(&units));
+                    Ok((files, units))
+                }
+                Err(error_diags) => {
+                    assert!(!error_diags.is_empty());
+                    let diags_buf = report_diagnostics_to_color_buffer(&files, error_diags);
+                    if let Err(err) = std::io::stdout().write_all(&diags_buf) {
+                        anyhow::bail!("Cannot output compiler diagnostics: {}", err);
+                    }
+                    anyhow::bail!("Compilation error");
+                }
+            }
+        })?;
+        Ok((compiled_pkg, fn_info.unwrap()))
+    }
+
     /// Given a `path` and a `build_config`, build the package in that path, including its dependencies.
     /// If we are building the Sui framework, we skip the check that the addresses should be 0
     pub fn build(self, path: PathBuf) -> SuiResult<CompiledPackage> {
         let res = if self.print_diags_to_stderr {
-            self.config
-                .compile_package_no_exit(&path, &mut std::io::stderr())
+            let resolution_graph = self
+                .config
+                .resolution_graph_for_package(&path, &mut std::io::stderr())
+                .map_err(|err| SuiError::ModuleBuildFailure {
+                    error: format!("{:?}", err),
+                })?;
+            Self::compile_package(resolution_graph, &mut std::io::stderr())
         } else {
-            self.config.compile_package_no_exit(&path, &mut Vec::new())
+            let resolution_graph = self
+                .config
+                .resolution_graph_for_package(&path, &mut Vec::new())
+                .map_err(|err| SuiError::ModuleBuildFailure {
+                    error: format!("{:?}", err),
+                })?;
+            Self::compile_package(resolution_graph, &mut Vec::new())
         };
 
         // write build failure diagnostics to stderr, convert `error` to `String` using `Debug`
         // format to include anyhow's error context chain.
-        let package = match res {
+        let (package, fn_info) = match res {
             Err(error) => {
                 return Err(SuiError::ModuleBuildFailure {
                     error: format!("{:?}", error),
                 })
             }
-            Ok(package) => package,
+            Ok((package, fn_info)) => (package, fn_info),
         };
         let compiled_modules = package.root_modules_map();
         if self.run_bytecode_verifier {
@@ -82,7 +159,7 @@ impl BuildConfig {
                         error: err.to_string(),
                     }
                 })?;
-                sui_bytecode_verifier::verify_module(m)?;
+                sui_bytecode_verifier::verify_module(m, &fn_info)?;
             }
             // TODO(https://github.com/MystenLabs/sui/issues/69): Run Move linker
         }

crates/sui-framework/docs/verifier_tests.md

@@ -0,0 +1,93 @@
+
+<a name="0x2_verifier_tests"></a>
+
+# Module `0x2::verifier_tests`
+
+Tests if normally illegal (in terms of Sui bytecode verification) code is allowed in tests.
+
+
+-  [Struct `VERIFIER_TESTS`](#0x2_verifier_tests_VERIFIER_TESTS)
+-  [Function `init`](#0x2_verifier_tests_init)
+-  [Function `is_otw`](#0x2_verifier_tests_is_otw)
+
+
+<pre><code><b>use</b> <a href="tx_context.md#0x2_tx_context">0x2::tx_context</a>;
+<b>use</b> <a href="types.md#0x2_types">0x2::types</a>;
+</code></pre>
+
+
+
+<a name="0x2_verifier_tests_VERIFIER_TESTS"></a>
+
+## Struct `VERIFIER_TESTS`
+
+
+
+<pre><code><b>struct</b> <a href="verifier_tests.md#0x2_verifier_tests_VERIFIER_TESTS">VERIFIER_TESTS</a> <b>has</b> drop
+</code></pre>
+
+
+
+<details>
+<summary>Fields</summary>
+
+
+<dl>
+<dt>
+<code>dummy_field: bool</code>
+</dt>
+<dd>
+
+</dd>
+</dl>
+
+
+</details>
+
+<a name="0x2_verifier_tests_init"></a>
+
+## Function `init`
+
+
+
+<pre><code><b>fun</b> <a href="verifier_tests.md#0x2_verifier_tests_init">init</a>(otw: <a href="verifier_tests.md#0x2_verifier_tests_VERIFIER_TESTS">verifier_tests::VERIFIER_TESTS</a>, _: &<b>mut</b> <a href="tx_context.md#0x2_tx_context_TxContext">tx_context::TxContext</a>)
+</code></pre>
+
+
+
+<details>
+<summary>Implementation</summary>
+
+
+<pre><code><b>fun</b> <a href="verifier_tests.md#0x2_verifier_tests_init">init</a>(otw: <a href="verifier_tests.md#0x2_verifier_tests_VERIFIER_TESTS">VERIFIER_TESTS</a>, _: &<b>mut</b> sui::tx_context::TxContext) {
+    <b>assert</b>!(sui::types::is_one_time_witness(&otw), 0);
+}
+</code></pre>
+
+
+
+</details>
+
+<a name="0x2_verifier_tests_is_otw"></a>
+
+## Function `is_otw`
+
+
+
+<pre><code><b>fun</b> <a href="verifier_tests.md#0x2_verifier_tests_is_otw">is_otw</a>(witness: <a href="verifier_tests.md#0x2_verifier_tests_VERIFIER_TESTS">verifier_tests::VERIFIER_TESTS</a>): bool
+</code></pre>
+
+
+
+<details>
+<summary>Implementation</summary>
+
+
+<pre><code><b>fun</b> <a href="verifier_tests.md#0x2_verifier_tests_is_otw">is_otw</a>(witness: <a href="verifier_tests.md#0x2_verifier_tests_VERIFIER_TESTS">VERIFIER_TESTS</a>): bool {
+    sui::types::is_one_time_witness(&witness)
+}
+</code></pre>
+
+
+
+</details>

crates/sui-framework/src/lib.rs

@@ -209,6 +209,21 @@ mod tests {
     }
 
     fn check_move_unit_tests(path: &Path) {
+        // build tests first to enable Sui-specific test code verification
+        matches!(
+            build_move_package(
+                path,
+                BuildConfig {
+                    config: MoveBuildConfig {
+                        test_mode: true, // make sure to verify tests
+                        ..MoveBuildConfig::default()
+                    },
+                    run_bytecode_verifier: true,
+                    print_diags_to_stderr: true,
+                },
+            ),
+            Ok(_)
+        );
         assert_eq!(
             run_move_unit_tests(path, MoveBuildConfig::default(), None, false).unwrap(),
             UnitTestResult::Success

crates/sui-framework/tests/verifier_tests.move

@@ -0,0 +1,35 @@
+// Copyright (c) Mysten Labs, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+/// Tests if normally illegal (in terms of Sui bytecode verification) code is allowed in tests.
+module sui::verifier_tests {
+    struct VERIFIER_TESTS has drop {}
+
+    fun init(otw: VERIFIER_TESTS, _: &mut sui::tx_context::TxContext) {
+        assert!(sui::types::is_one_time_witness(&otw), 0);
+    }
+
+    #[test]
+    fun test_init() {
+        use sui::test_scenario;
+        let admin = @0xBABE;
+
+        let scenario_val = test_scenario::begin(admin);
+        let scenario = &mut scenario_val;
+        let otw = VERIFIER_TESTS{};
+        init(otw, test_scenario::ctx(scenario));
+        test_scenario::end(scenario_val);
+    }
+
+    fun is_otw(witness: VERIFIER_TESTS): bool {
+        sui::types::is_one_time_witness(&witness)
+    }
+
+    #[test]
+    fun test_otw() {
+        // we should be able to construct otw in test code
+        let otw = VERIFIER_TESTS{};
+        assert!(is_otw(otw), 0);
+    }
+
+}

crates/sui-types/src/move_package.rs

@@ -9,7 +9,7 @@ use move_binary_format::access::ModuleAccess;
 use move_binary_format::binary_views::BinaryIndexedView;
 use move_binary_format::file_format::CompiledModule;
 use move_binary_format::normalized;
-use move_core_types::identifier::Identifier;
+use move_core_types::{account_address::AccountAddress, identifier::Identifier};
 use move_disassembler::disassembler::Disassembler;
 use move_ir_types::location::Spanned;
 use serde::{Deserialize, Serialize};
@@ -24,6 +24,23 @@ use sui_protocol_constants::*;
 // #[path = "unit_tests/move_package.rs"]
 // mod base_types_tests;
 
+#[derive(Clone, Debug)]
+/// Additional information about a function
+pub struct FnInfo {
+    /// If true, it's a function involved in testing (`[test]`, `[test_only]`, `[expected_failure]`)
+    pub is_test: bool,
+}
+
+#[derive(Clone, Debug, Eq, Hash, PartialEq, PartialOrd, Ord)]
+/// Uniquely identifies a function in a module
+pub struct FnInfoKey {
+    pub fn_name: String,
+    pub mod_addr: AccountAddress,
+}
+
+/// A map from function info keys to function info
+pub type FnInfoMap = BTreeMap<FnInfoKey, FnInfo>;
+
 // serde_bytes::ByteBuf is an analog of Vec<u8> with built-in fast serialization.
 #[serde_as]
 #[derive(Eq, PartialEq, Debug, Clone, Deserialize, Serialize, Hash)]

crates/sui-verifier/src/entry_points_verifier.rs

@@ -16,10 +16,13 @@ use sui_types::{
     },
     error::ExecutionError,
     id::{ID_STRUCT_NAME, OBJECT_MODULE_NAME},
+    move_package::FnInfoMap,
     MOVE_STDLIB_ADDRESS, SUI_FRAMEWORK_ADDRESS,
 };
 
-use crate::{format_signature_token, resolve_struct, verification_failure, INIT_FN_NAME};
+use crate::{
+    format_signature_token, is_test_fun, resolve_struct, verification_failure, INIT_FN_NAME,
+};
 
 /// Checks valid rules rules for entry points, both for module initialization and transactions
 ///
@@ -37,12 +40,22 @@ use crate::{format_signature_token, resolve_struct, verification_failure, INIT_F
 /// - The function may have a &mut TxContext or &TxContext (see `is_tx_context`) parameter
 ///   - The transaction context parameter must be the last parameter
 /// - The function cannot have any return values
-pub fn verify_module(module: &CompiledModule) -> Result<(), ExecutionError> {
-    for func_def in &module.function_defs {
-        verify_init_not_called(module, func_def).map_err(verification_failure)?;
+pub fn verify_module(
+    module: &CompiledModule,
+    fn_info_map: &FnInfoMap,
+) -> Result<(), ExecutionError> {
+    // When verifying test functions, a check preventing explicit calls to init functions is
+    // disabled.
 
+    for func_def in &module.function_defs {
         let handle = module.function_handle_at(func_def.function);
         let name = module.identifier_at(handle.name);
+
+        // allow calling init function in the test code
+        if !is_test_fun(name, module, fn_info_map) {
+            verify_init_not_called(module, func_def).map_err(verification_failure)?;
+        }
+
         if name == INIT_FN_NAME {
             verify_init_function(module, func_def).map_err(verification_failure)?;
             continue;