cherry-pick two improvements (#11020)

* [vm][loader] optimize memory usage for types * Fix OOM for string transaction arguments (#189) --------- Co-authored-by: Victor Gao <[email protected]> Co-authored-by: George Mitenkov <[email protected]>
aptos-labs · Nov 21, 2023 · 903e482 · 903e482
1 parent 1a6c758
commit 903e482
Show file tree

Hide file tree

Showing 12 changed files with 274 additions and 191 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -530,6 +530,7 @@ log = "0.4.17"
 lru = "0.7.5"
 lz4 = "1.24.0"
 maplit = "1.0.2"
+memory-stats = "1.1.0"
 merlin = "3"
 mime = "0.3.16"
 mirai-annotations = "1.12.0"

diff --git a/aptos-move/aptos-vm/src/verifier/transaction_arg_validation.rs b/aptos-move/aptos-vm/src/verifier/transaction_arg_validation.rs
@@ -255,6 +255,7 @@ fn construct_arg(
     match ty {
         Bool | U8 | U16 | U32 | U64 | U128 | U256 | Address => Ok(arg),
         Vector(_) | Struct(_) | StructInstantiation(_, _) => {
+            let initial_cursor_len = arg.len();
             let mut cursor = Cursor::new(&arg[..]);
             let mut new_arg = vec![];
             let mut max_invocations = 10; // Read from config in the future
@@ -263,13 +264,14 @@ fn construct_arg(
                 ty,
                 allowed_structs,
                 &mut cursor,
+                initial_cursor_len,
                 gas_meter,
                 &mut max_invocations,
                 &mut new_arg,
             )?;
             // Check cursor has parsed everything
             // Unfortunately, is_empty is only enabled in nightly, so we check this way.
-            if cursor.position() != arg.len() as u64 {
+            if cursor.position() != initial_cursor_len as u64 {
                 return Err(VMStatus::error(
                     StatusCode::FAILED_TO_DESERIALIZE_ARGUMENT,
                     Some(String::from(
@@ -298,6 +300,7 @@ pub(crate) fn recursively_construct_arg(
     ty: &Type,
     allowed_structs: &ConstructorMap,
     cursor: &mut Cursor<&[u8]>,
+    initial_cursor_len: usize,
     gas_meter: &mut impl GasMeter,
     max_invocations: &mut u64,
     arg: &mut Vec<u8>,
@@ -315,6 +318,7 @@ pub(crate) fn recursively_construct_arg(
                     inner,
                     allowed_structs,
                     cursor,
+                    initial_cursor_len,
                     gas_meter,
                     max_invocations,
                     arg,
@@ -340,6 +344,7 @@ pub(crate) fn recursively_construct_arg(
                 constructor,
                 allowed_structs,
                 cursor,
+                initial_cursor_len,
                 gas_meter,
                 max_invocations,
             )?);
@@ -365,6 +370,7 @@ fn validate_and_construct(
     constructor: &FunctionId,
     allowed_structs: &ConstructorMap,
     cursor: &mut Cursor<&[u8]>,
+    initial_cursor_len: usize,
     gas_meter: &mut impl GasMeter,
     max_invocations: &mut u64,
 ) -> Result<Vec<u8>, VMStatus> {
@@ -394,8 +400,21 @@ fn validate_and_construct(
                 VMStatus::error(StatusCode::FAILED_TO_DESERIALIZE_ARGUMENT, None)
             }
         };
-        // short cut for the utf8 constructor, which is a special case
+        // Short cut for the utf8 constructor, which is a special case.
         let len = get_len(cursor)?;
+        if !cursor
+            .position()
+            .checked_add(len as u64)
+            .is_some_and(|l| l <= initial_cursor_len as u64)
+        {
+            // We need to make sure we do not allocate more bytes than
+            // needed.
+            return Err(VMStatus::error(
+                StatusCode::FAILED_TO_DESERIALIZE_ARGUMENT,
+                Some("String argument is too long".to_string()),
+            ));
+        }
+
         let mut arg = vec![];
         read_n_bytes(len, cursor, &mut arg)?;
         std::str::from_utf8(&arg).map_err(|_| constructor_error())?;
@@ -418,6 +437,7 @@ fn validate_and_construct(
             &param_type.subst(&instantiation.type_arguments).unwrap(),
             allowed_structs,
             cursor,
+            initial_cursor_len,
             gas_meter,
             max_invocations,
             &mut arg,
@@ -457,12 +477,28 @@ fn serialize_uleb128(mut x: usize, dest: &mut Vec<u8>) {
 }
 
 fn read_n_bytes(n: usize, src: &mut Cursor<&[u8]>, dest: &mut Vec<u8>) -> Result<(), VMStatus> {
-    let len = dest.len();
-    dest.resize(len + n, 0);
-    src.read_exact(&mut dest[len..]).map_err(|_| {
+    let deserialization_error = |msg: &str| -> VMStatus {
         VMStatus::error(
             StatusCode::FAILED_TO_DESERIALIZE_ARGUMENT,
-            Some(String::from("Couldn't read bytes")),
+            Some(msg.to_string()),
         )
-    })
+    };
+    let len = dest.len();
+
+    // It is safer to limit the length under some big (but still reasonable
+    // number).
+    const MAX_NUM_BYTES: usize = 1_000_000;
+    if !len.checked_add(n).is_some_and(|s| s <= MAX_NUM_BYTES) {
+        return Err(deserialization_error(&format!(
+            "Couldn't read bytes: maximum limit of {} bytes exceeded",
+            MAX_NUM_BYTES
+        )));
+    }
+
+    // Ensure we have enough capacity for resizing.
+    dest.try_reserve(len + n)
+        .map_err(|e| deserialization_error(&format!("Couldn't read bytes: {}", e)))?;
+    dest.resize(len + n, 0);
+    src.read_exact(&mut dest[len..])
+        .map_err(|_| deserialization_error("Couldn't read bytes"))
 }
diff --git a/aptos-move/e2e-move-tests/Cargo.toml b/aptos-move/e2e-move-tests/Cargo.toml
@@ -34,10 +34,14 @@ aptos-writeset-generator = { workspace = true }
 bcs = { workspace = true }
 hex = { workspace = true }
 itertools = { workspace = true }
+memory-stats = { workspace = true }
 move-binary-format = { workspace = true }
 move-core-types = { workspace = true }
 move-package = { workspace = true }
 move-symbol-pool = { workspace = true }
+move-vm-runtime = { workspace = true }
+move-vm-test-utils = { workspace = true }
+move-vm-types = { workspace = true }
 once_cell = { workspace = true }
 project-root = { workspace = true }
 proptest = { workspace = true }

diff --git a/aptos-move/e2e-move-tests/src/tests/mod.rs b/aptos-move/e2e-move-tests/src/tests/mod.rs
@@ -23,6 +23,7 @@ mod memory_quota;
 mod metadata;
 mod mint_nft;
 mod missing_gas_parameter;
+mod module_bomb;
 mod module_event;
 mod new_integer_types;
 mod nft_dao;

diff --git a/aptos-move/e2e-move-tests/src/tests/module_bomb.mv b/aptos-move/e2e-move-tests/src/tests/module_bomb.mv
diff --git a/aptos-move/e2e-move-tests/src/tests/module_bomb.rs b/aptos-move/e2e-move-tests/src/tests/module_bomb.rs
@@ -0,0 +1,36 @@
+// Copyright © Aptos Foundation
+
+use move_binary_format::CompiledModule;
+use move_core_types::identifier::Identifier;
+use move_vm_runtime::move_vm::MoveVM;
+use move_vm_test_utils::BlankStorage;
+use move_vm_types::gas::UnmeteredGasMeter;
+
+const BLOB: &[u8] = include_bytes!("module_bomb.mv");
+
+#[test]
+fn test_module_bomb() {
+    let m = CompiledModule::deserialize(BLOB).unwrap();
+
+    let mut vms = vec![];
+
+    for _i in 0..2 {
+        let vm = MoveVM::new(vec![]).unwrap();
+        let storage = BlankStorage;
+        let mut sess = vm.new_session(&storage);
+        sess.publish_module(
+            BLOB.to_vec(),
+            *m.self_id().address(),
+            &mut UnmeteredGasMeter,
+        )
+        .unwrap();
+
+        sess.load_function(&m.self_id(), &Identifier::new("f1").unwrap(), &[])
+            .unwrap();
+        vms.push(vm);
+    }
+
+    let stats = memory_stats::memory_stats().unwrap();
+    println!("Physical: {}", stats.physical_mem);
+    println!("Virtual: {}", stats.virtual_mem);
+}
diff --git a/aptos-move/e2e-move-tests/src/tests/string_args.rs b/aptos-move/e2e-move-tests/src/tests/string_args.rs
@@ -653,3 +653,26 @@ fn string_args_generic_instantiation() {
 
     success_generic(vec![string_type, address_type], tests);
 }
+
+#[test]
+fn huge_string_args_are_not_allowed() {
+    let mut tests = vec![];
+    let mut len: u64 = 1_000_000_000_000;
+    let mut big_str_arg = vec![];
+    loop {
+        let cur = len & 0x7F;
+        if cur != len {
+            big_str_arg.push((cur | 0x80) as u8);
+            len >>= 7;
+        } else {
+            big_str_arg.push(cur as u8);
+            break;
+        }
+    }
+    tests.push((
+        "0xcafe::test::hi",
+        vec![big_str_arg],
+        deserialization_failure(),
+    ));
+    fail(tests);
+}
diff --git a/third_party/move/move-vm/runtime/Cargo.toml b/third_party/move/move-vm/runtime/Cargo.toml
@@ -11,6 +11,7 @@ edition = "2021"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
+backtrace = "0.3.69"
 better_any = "0.1.1"
 fail = "0.4.0"
 move-binary-format = { path = "../../move-binary-format" }

diff --git a/third_party/move/move-vm/runtime/src/interpreter.rs b/third_party/move/move-vm/runtime/src/interpreter.rs
@@ -1040,7 +1040,8 @@ fn check_depth_of_type_impl(
         | Type::Signer => check_depth!(0),
         // Even though this is recursive this is OK since the depth of this recursion is
         // bounded by the depth of the type arguments, which we have already checked.
-        Type::Reference(ty) | Type::MutableReference(ty) | Type::Vector(ty) => {
+        Type::Vector(ty) => check_depth_of_type_impl(resolver, ty, max_depth, check_depth!(1))?,
+        Type::Reference(ty) | Type::MutableReference(ty) => {
             check_depth_of_type_impl(resolver, ty, max_depth, check_depth!(1))?
         },
         Type::Struct(si) => {
@@ -1658,7 +1659,7 @@ impl Frame {
                 }
                 interpreter
                     .operand_stack
-                    .push_ty(Type::Vector(Box::new(ty)))?;
+                    .push_ty(Type::Vector(Arc::new(ty)))?;
             },
             Bytecode::VecLen(si) => {
                 let ty = resolver.instantiate_single_type(*si, ty_args)?;