[randomness] Add entry function annotation (aptos-labs#12134)

Adds `#[randomness]` annotation which should be used to mark
entry functions which use randomness. For example
```
#[randomness]
entry fun foo() {
  let _ = randomness::u64_integer();
}
```

The attribute has the following semantics:
1. It can only be used for entry functions. Using it on other type
    of functions is not allowed.
2. It can be used on entry functions even if they don't use randomness.
3. If an entry function doesn't have an annotation, but uses randomness,
    the randomness call fails at runtime.

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

@@ -20,13 +20,15 @@ use crate::{
     sharded_block_executor::{executor_client::ExecutorClient, ShardedBlockExecutor},
     system_module_names::*,
     transaction_metadata::TransactionMetadata,
-    transaction_validation, verifier, VMExecutor, VMValidator,
+    transaction_validation, verifier,
+    verifier::randomness::has_randomness_attribute,
+    VMExecutor, VMValidator,
 };
 use anyhow::anyhow;
 use aptos_block_executor::txn_commit_hook::NoOpTransactionCommitHook;
 use aptos_crypto::HashValue;
 use aptos_framework::{
-    natives::{code::PublishRequest, transaction_context::NativeTransactionContext},
+    natives::{code::PublishRequest, randomness::RandomnessContext},
     RuntimeModuleMetadataV1,
 };
 use aptos_gas_algebra::{Gas, GasQuantity, NumBytes, Octa};
@@ -733,6 +735,7 @@ impl AptosVM {
 
     fn validate_and_execute_entry_function(
         &self,
+        resolver: &impl AptosMoveResolver,
         session: &mut SessionExt,
         gas_meter: &mut impl AptosGasMeter,
         traversal_context: &mut TraversalContext,
@@ -752,26 +755,27 @@ impl AptosVM {
             )])?;
         }
 
-        let is_friend_or_private = session.load_function_def_is_friend_or_private(
+        let (function, is_friend_or_private) = session.load_function_and_is_friend_or_private_def(
             entry_fn.module(),
             entry_fn.function(),
             entry_fn.ty_args(),
         )?;
-        if is_friend_or_private {
+
+        if is_friend_or_private && has_randomness_attribute(resolver, session, entry_fn)? {
             let txn_context = session
                 .get_native_extensions()
-                .get_mut::<NativeTransactionContext>();
-            txn_context.set_is_friend_or_private_entry_func();
+                .get_mut::<RandomnessContext>();
+            txn_context.mark_unbiasable();
         }
 
-        let function =
-            session.load_function(entry_fn.module(), entry_fn.function(), entry_fn.ty_args())?;
+        let struct_constructors_enabled =
+            self.features().is_enabled(FeatureFlag::STRUCT_CONSTRUCTORS);
         let args = verifier::transaction_arg_validation::validate_combine_signer_and_txn_args(
             session,
             senders,
             entry_fn.args().to_vec(),
             &function,
-            self.features().is_enabled(FeatureFlag::STRUCT_CONSTRUCTORS),
+            struct_constructors_enabled,
         )?;
         session.execute_entry_function(
             entry_fn.module(),
@@ -820,6 +824,7 @@ impl AptosVM {
             TransactionPayload::EntryFunction(entry_fn) => {
                 session.execute(|session| {
                     self.validate_and_execute_entry_function(
+                        resolver,
                         session,
                         gas_meter,
                         traversal_context,
@@ -921,6 +926,7 @@ impl AptosVM {
                         aptos_try!({
                             return_on_failure!(session.execute(|session| self
                                 .execute_multisig_entry_function(
+                                    resolver,
                                     session,
                                     gas_meter,
                                     traversal_context,
@@ -1042,6 +1048,7 @@ impl AptosVM {
             MultisigTransactionPayload::EntryFunction(entry_function) => {
                 session.execute(|session| {
                     self.execute_multisig_entry_function(
+                        resolver,
                         session,
                         gas_meter,
                         traversal_context,
@@ -1142,6 +1149,7 @@ impl AptosVM {
 
     fn execute_multisig_entry_function(
         &self,
+        resolver: &impl AptosMoveResolver,
         session: &mut SessionExt,
         gas_meter: &mut impl AptosGasMeter,
         traversal_context: &mut TraversalContext,
@@ -1152,6 +1160,7 @@ impl AptosVM {
         // If txn args are not valid, we'd still consider the transaction as executed but
         // failed. This is primarily because it's unrecoverable at this point.
         self.validate_and_execute_entry_function(
+            resolver,
             session,
             gas_meter,
             traversal_context,

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

@@ -227,9 +227,11 @@ fn unit_test_extensions_hook(exts: &mut NativeContextExtensions) {
 
     exts.add(NativeTableContext::new([0u8; 32], &*DUMMY_RESOLVER));
     exts.add(NativeCodeContext::default());
-    let mut txn_context = NativeTransactionContext::new(vec![1], vec![1], ChainId::test().id());
-    txn_context.set_is_friend_or_private_entry_func();
-    exts.add(txn_context); // We use the testing environment chain ID here
+    exts.add(NativeTransactionContext::new(
+        vec![1],
+        vec![1],
+        ChainId::test().id(),
+    ));
     exts.add(NativeAggregatorContext::new(
         [0; 32],
         &*DUMMY_RESOLVER,
@@ -238,5 +240,8 @@ fn unit_test_extensions_hook(exts: &mut NativeContextExtensions) {
     exts.add(NativeRistrettoPointContext::new());
     exts.add(AlgebraContext::new());
     exts.add(NativeEventContext::default());
-    exts.add(RandomnessContext::new());
+
+    let mut randomness_ctx = RandomnessContext::new();
+    randomness_ctx.mark_unbiasable();
+    exts.add(randomness_ctx);
 }

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

@@ -2,6 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 pub(crate) mod event_validation;
 pub(crate) mod module_init;
+pub(crate) mod randomness;
 pub(crate) mod resource_groups;
 pub mod transaction_arg_validation;
 pub(crate) mod view_function;

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

@@ -0,0 +1,27 @@
+// Copyright © Aptos Foundation
+// SPDX-License-Identifier: Apache-2.0
+
+use crate::move_vm_ext::{AptosMoveResolver, SessionExt};
+use aptos_types::transaction::EntryFunction;
+use move_binary_format::errors::VMResult;
+
+/// Returns true if function has an attribute that it uses randomness.
+pub(crate) fn has_randomness_attribute(
+    resolver: &impl AptosMoveResolver,
+    session: &mut SessionExt,
+    entry_fn: &EntryFunction,
+) -> VMResult<bool> {
+    let module = session
+        .get_move_vm()
+        .load_module(entry_fn.module(), resolver)?;
+    let metadata = aptos_framework::get_metadata_from_compiled_module(&module);
+    if let Some(metadata) = metadata {
+        Ok(metadata
+            .fun_attributes
+            .get(entry_fn.function().as_str())
+            .map(|attrs| attrs.iter().any(|attr| attr.is_randomness()))
+            .unwrap_or(false))
+    } else {
+        Ok(false)
+    }
+}

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

@@ -35,6 +35,7 @@ mod object_code_deployment;
 mod offer_rotation_capability;
 mod offer_signer_capability;
 mod per_category_gas_limits;
+mod randomness_test_and_abort;
 mod resource_groups;
 mod rotate_auth_key;
 mod scripts;

aptos-move/e2e-move-tests/src/tests/randomness.data/invalid_pack_non_entry/Move.toml

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

aptos-move/e2e-move-tests/src/tests/randomness.data/invalid_pack_non_entry/sources/invalid.move

@@ -0,0 +1,6 @@
+module 0x1::invalid {
+    #[randomness]
+    public entry fun foo() {
+        // Do nothing.
+    }
+}

aptos-move/e2e-move-tests/src/tests/randomness.data/invalid_pack_public_entry/Move.toml

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

aptos-move/e2e-move-tests/src/tests/randomness.data/invalid_pack_public_entry/sources/invalid.move

@@ -0,0 +1,6 @@
+module 0x1::invalid {
+    #[randomness]
+    public fun foo() {
+        // Do nothing.
+    }
+}

aptos-move/e2e-move-tests/src/tests/randomness.data/pack/Move.toml

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

aptos-move/e2e-move-tests/src/tests/randomness.data/pack/sources/test.move

@@ -0,0 +1,21 @@
+module 0x1::test {
+    use aptos_framework::randomness;
+
+    entry fun ok_if_not_annotated_and_not_using_randomness() {
+        // Do nothing.
+    }
+
+    #[randomness]
+    entry fun ok_if_annotated_and_not_using_randomness() {
+        // Do nothing.
+    }
+
+    #[randomness]
+    entry fun ok_if_annotated_and_using_randomness() {
+        let _ = randomness::u64_integer();
+    }
+
+    entry fun fail_if_not_annotated_and_using_randomness() {
+        let _ = randomness::u64_integer();
+    }
+}

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

@@ -0,0 +1,168 @@
+// Copyright © Aptos Foundation
+// SPDX-License-Identifier: Apache-2.0
+
+use crate::{assert_abort, assert_success, build_package, tests::common, MoveHarness};
+use aptos_framework::BuiltPackage;
+use aptos_language_e2e_tests::account::{Account, TransactionBuilder};
+use aptos_types::{
+    account_address::AccountAddress,
+    on_chain_config::OnChainConfig,
+    randomness::PerBlockRandomness,
+    transaction::{ExecutionStatus, Script, TransactionStatus},
+};
+use claims::assert_ok;
+use move_core_types::{ident_str, language_storage::ModuleId, vm_status::AbortLocation};
+
+// Error codes from randomness module.
+const E_API_USE_SUSCEPTIBLE_TO_TEST_AND_ABORT: u64 = 1;
+
+#[test]
+fn test_and_abort_defense_is_sound_and_correct() {
+    let mut h = MoveHarness::new();
+
+    // These scripts call a public entry function and a public function. The randomness API will reject both calls.
+    for dir in [
+        "randomness_unsafe_public_entry.data/pack",
+        "randomness_unsafe_public.data/pack",
+    ] {
+        println!("Testing {dir}");
+        // This will redeploy the package, so backwards compatibility must be maintained in these directories.
+        let (_, package) =
+            deploy_code(AccountAddress::ONE, dir, &mut h).expect("building package must succeed");
+
+        let status = run_script(&mut h, &package);
+        assert_abort!(status, E_API_USE_SUSCEPTIBLE_TO_TEST_AND_ABORT);
+    }
+
+    // The randomness module is initialized, but the randomness seed is not set.
+    set_randomness_seed(&mut h);
+
+    // This is a safe call that the randomness API should allow through.
+    let status = run_entry_func(
+        &mut h,
+        "0xa11ce",
+        "0x1::some_randapp::safe_private_entry_call",
+    );
+    assert_success!(status);
+
+    // This is a safe call that the randomness API should allow through.
+    // (I suppose that, since TXNs with private entry function payloads are okay, increasing the
+    // visibility to public(friend) should not create any problems.)
+    let status = run_entry_func(
+        &mut h,
+        "0xa11ce",
+        "0x1::some_randapp::safe_friend_entry_call",
+    );
+    assert_success!(status);
+}
+
+#[test]
+fn test_only_private_entry_function_can_be_annotated() {
+    // Make sure building a package fails.
+    let mut h = MoveHarness::new();
+    assert!(deploy_code(
+        AccountAddress::ONE,
+        "randomness.data/invalid_pack_non_entry",
+        &mut h
+    )
+    .is_err());
+    assert!(deploy_code(
+        AccountAddress::ONE,
+        "randomness.data/invalid_pack_public_entry",
+        &mut h
+    )
+    .is_err());
+}
+
+#[test]
+fn test_unbiasable_annotation() {
+    let mut h = MoveHarness::new();
+    deploy_code(AccountAddress::ONE, "randomness.data/pack", &mut h)
+        .expect("building package must succeed");
+    set_randomness_seed(&mut h);
+
+    let should_succeed = [
+        "0x1::test::ok_if_not_annotated_and_not_using_randomness",
+        "0x1::test::ok_if_annotated_and_not_using_randomness",
+        "0x1::test::ok_if_annotated_and_using_randomness",
+    ];
+    for entry_func in should_succeed {
+        let status = run_entry_func(&mut h, "0xa11ce", entry_func);
+        assert_success!(status);
+    }
+
+    // Non-annotated functions which use randomness fail at runtime.
+    let entry_func = "0x1::test::fail_if_not_annotated_and_using_randomness";
+    let status = run_entry_func(&mut h, "0xa11ce", entry_func);
+    let status = assert_ok!(status.as_kept_status());
+
+    if let ExecutionStatus::MoveAbort {
+        location,
+        code,
+        info: _,
+    } = status
+    {
+        assert_eq!(
+            location,
+            AbortLocation::Module(ModuleId::new(
+                AccountAddress::ONE,
+                ident_str!("randomness").to_owned()
+            ))
+        );
+        assert_eq!(code, E_API_USE_SUSCEPTIBLE_TO_TEST_AND_ABORT);
+    } else {
+        unreachable!("Non-annotated entry call function should result in Move abort")
+    }
+}
+
+fn set_randomness_seed(h: &mut MoveHarness) {
+    let fx = h.aptos_framework_account();
+    let mut pbr = h
+        .read_resource::<PerBlockRandomness>(fx.address(), PerBlockRandomness::struct_tag())
+        .unwrap();
+    assert!(pbr.seed.is_none());
+
+    pbr.seed = Some((0..32).map(|_| 0u8).collect::<Vec<u8>>());
+    assert_eq!(pbr.seed.as_ref().unwrap().len(), 32);
+    h.set_resource(*fx.address(), PerBlockRandomness::struct_tag(), &pbr);
+}
+
+fn deploy_code(
+    addr: AccountAddress,
+    code_path: &str,
+    harness: &mut MoveHarness,
+) -> anyhow::Result<(Account, BuiltPackage)> {
+    let account = harness.new_account_at(addr);
+
+    let package = build_package(
+        common::test_dir_path(code_path),
+        aptos_framework::BuildOptions::default(),
+    )?;
+
+    let txn = harness.create_publish_built_package(&account, &package, |_| {});
+
+    assert_success!(harness.run(txn));
+    Ok((account, package))
+}
+
+fn run_script(h: &mut MoveHarness, package: &BuiltPackage) -> TransactionStatus {
+    let alice = h.new_account_at(AccountAddress::from_hex_literal("0xa11ce").unwrap());
+    let scripts = package.extract_script_code();
+    let code = scripts[0].clone();
+
+    let txn = TransactionBuilder::new(alice.clone())
+        .script(Script::new(code, vec![], vec![]))
+        .sequence_number(10)
+        .max_gas_amount(1_000_000)
+        .gas_unit_price(1)
+        .sign();
+
+    h.run(txn)
+}
+
+fn run_entry_func(h: &mut MoveHarness, signer: &str, name: &str) -> TransactionStatus {
+    let alice = h.new_account_at(AccountAddress::from_hex_literal(signer).unwrap());
+
+    println!("Running entry function '{name}'");
+    h.run_entry_function(&alice, str::parse(name).unwrap(), vec![], vec![])
+}