Validate Move String transaction arguments

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

@@ -11,17 +11,25 @@ use crate::{
     VMStatus,
 };
 use move_deps::{
+    move_binary_format::file_format_common::read_uleb128_as_u64,
     move_core_types::{account_address::AccountAddress, value::MoveValue, vm_status::StatusCode},
     move_vm_runtime::session::LoadedFunctionInstantiation,
     move_vm_types::loaded_data::runtime_types::Type,
 };
 use once_cell::sync::Lazy;
-use std::collections::BTreeSet;
+use std::collections::BTreeMap;
+use std::io::{Cursor, Read};
 
-static ALLOWED_STRUCTS: Lazy<BTreeSet<String>> = Lazy::new(|| {
-    ["0x1::string::String"]
-        .iter()
-        .map(|s| s.to_string())
+// A map which contains the structs allowed as transaction input and the
+// validation function for those, if one was needed (None otherwise).
+// The validation function takes the serialized argument and returns
+// an error if the validation fails.
+type ValidateArg = fn(&[u8]) -> Result<(), VMStatus>;
+
+static ALLOWED_STRUCTS: Lazy<BTreeMap<String, Option<ValidateArg>>> = Lazy::new(|| {
+    [("0x1::string::String", Some(check_string as ValidateArg))]
+        .into_iter()
+        .map(|(s, validator)| (s.to_string(), validator))
         .collect()
 });
 
@@ -56,8 +64,11 @@ pub(crate) fn validate_combine_signer_and_txn_args<S: MoveResolverExt>(
         }
     }
     // validate all non_signer params
+    let mut needs_validation = false;
     for ty in func.parameters[signer_param_cnt..].iter() {
-        if !is_valid_txn_arg(session, ty) {
+        let (valid, validation) = is_valid_txn_arg(session, ty);
+        needs_validation = needs_validation || validation;
+        if !valid {
             return Err(VMStatus::Error(StatusCode::INVALID_MAIN_FUNCTION_SIGNATURE));
         }
     }
@@ -83,22 +94,116 @@ pub(crate) fn validate_combine_signer_and_txn_args<S: MoveResolverExt>(
             .chain(args)
             .collect()
     };
+    if needs_validation {
+        validate_args(session, &combined_args, func)?;
+    }
     Ok(combined_args)
 }
 
-fn is_valid_txn_arg<S: MoveResolverExt>(session: &SessionExt<S>, typ: &Type) -> bool {
+// Return whether the argument is valid/allowed and whether it needs validation.
+// Validation is only needed for String arguments at the moment and vectors of them.
+fn is_valid_txn_arg<S: MoveResolverExt>(session: &SessionExt<S>, typ: &Type) -> (bool, bool) {
     use move_deps::move_vm_types::loaded_data::runtime_types::Type::*;
     match typ {
-        Bool | U8 | U64 | U128 | Address => true,
+        Bool | U8 | U64 | U128 | Address => (true, false),
         Vector(inner) => is_valid_txn_arg(session, inner),
         Struct(idx) | StructInstantiation(idx, _) => {
             if let Some(st) = session.get_struct_type(*idx) {
                 let full_name = format!("{}::{}", st.module.short_str_lossless(), st.name);
-                ALLOWED_STRUCTS.contains(&full_name)
+                match ALLOWED_STRUCTS.get(&full_name) {
+                    None => (false, false),
+                    Some(validator) => (true, validator.is_some()),
+                }
             } else {
-                false
+                (false, false)
+            }
+        }
+        Signer | Reference(_) | MutableReference(_) | TyParam(_) => (false, false),
+    }
+}
+
+// Validate arguments. Walk through the arguments and according to the signature
+// validate arguments that require so.
+// TODO: This needs a more solid story and a tighter integration with the VM.
+// Validation at the moment is only for Strings and Vector of them, so we
+// manually walk the serialized args until we find a string.
+// This is obviously brittle and something to change at some point soon.
+fn validate_args<S: MoveResolverExt>(
+    session: &SessionExt<S>,
+    args: &[Vec<u8>],
+    func: &LoadedFunctionInstantiation,
+) -> Result<(), VMStatus> {
+    for (ty, arg) in func.parameters.iter().zip(args.iter()) {
+        let mut cursor = Cursor::new(&arg[..]);
+        validate_arg(session, ty, &mut cursor)?;
+    }
+    Ok(())
+}
+
+// Validate a single arg. A Cursor is used to walk the serialized arg manually and correctly.
+fn validate_arg<S: MoveResolverExt>(
+    session: &SessionExt<S>,
+    ty: &Type,
+    cursor: &mut Cursor<&[u8]>,
+) -> Result<(), VMStatus> {
+    use move_deps::move_vm_types::loaded_data::runtime_types::Type::*;
+    Ok(match ty {
+        Vector(inner) => {
+            // get the vector length and iterate over each element
+            let mut len = get_len(cursor)?;
+            while len > 0 {
+                validate_arg(session, inner, cursor)?;
+                len -= 1;
+            }
+        }
+        // only strings are validated, and given we are here only if one was present
+        // (`is_valid_txn_arg`), this match arm must be for a string
+        Struct(idx) | StructInstantiation(idx, _) => {
+            // load the struct name, we use `expect()` because that check was already
+            // performed in `is_valid_txn_arg`
+            let st = session
+                .get_struct_type(*idx)
+                .expect("unrachable, type must exist");
+            let full_name = format!("{}::{}", st.module.short_str_lossless(), st.name);
+            // load the serialized string
+            let len = get_len(cursor)?;
+            let mut s = vec![0u8; len];
+            cursor
+                .read_exact(&mut s)
+                .map_err(|_| VMStatus::Error(StatusCode::FAILED_TO_DESERIALIZE_ARGUMENT))?;
+            // validate the struct value, we use `expect()` because that check was already
+            // performed in `is_valid_txn_arg`
+            let option_validator = ALLOWED_STRUCTS
+                .get(&full_name)
+                .expect("unreachable: struct must be allowed");
+            if let Some(validator) = option_validator {
+                validator(&s)?;
             }
         }
-        Signer | Reference(_) | MutableReference(_) | TyParam(_) => false,
+        // nothing to validate
+        Bool | U8 | U64 | U128 | Address | Signer | Reference(_) | MutableReference(_)
+        | TyParam(_) => (),
+    })
+}
+
+// String is a vector of bytes, so both string and vector carry a length in the serialized format.
+// Length of vectors in BCS uses uleb128 as a compression format.
+fn get_len(cursor: &mut Cursor<&[u8]>) -> Result<usize, VMStatus> {
+    match read_uleb128_as_u64(cursor) {
+        Err(_) => Err(VMStatus::Error(StatusCode::FAILED_TO_DESERIALIZE_ARGUMENT)),
+        Ok(len) => Ok(len as usize),
+    }
+}
+
+//
+// Argument validation functions
+//
+
+// Check if a string is valid. This code is copied from string.rs in the stdlib.
+// TODO: change the move VM code (string.rs) to expose a function that does validation.
+fn check_string(s: &[u8]) -> Result<(), VMStatus> {
+    match std::str::from_utf8(s) {
+        Ok(_) => Ok(()),
+        Err(_) => Err(VMStatus::Error(StatusCode::FAILED_TO_DESERIALIZE_ARGUMENT)),
     }
 }

aptos-move/e2e-move-tests/tests/string_args.data/pack/sources/hello.move

@@ -1,6 +1,7 @@
 module 0xCAFE::test {
     use std::signer;
     use std::string::String;
+    use std::vector;
 
     struct ModuleData has key, store {
         state: String,
@@ -14,4 +15,22 @@ module 0xCAFE::test {
             borrow_global_mut<ModuleData>(addr).state = msg;
         }
     }
+
+    public entry fun hi_vec(sender: &signer, msgs: vector<String>, i: u64) acquires ModuleData {
+        let addr = signer::address_of(sender);
+        if (!exists<ModuleData>(addr)) {
+            move_to(sender, ModuleData{state: *vector::borrow(&msgs, i)})
+        } else {
+            borrow_global_mut<ModuleData>(addr).state = *vector::borrow(&msgs, i);
+        }
+    }
+
+    public entry fun more_hi_vec(sender: &signer, msgs: vector<vector<String>>, i: u64, j: u64) acquires ModuleData {
+        let addr = signer::address_of(sender);
+        if (!exists<ModuleData>(addr)) {
+            move_to(sender, ModuleData{state: *vector::borrow(vector::borrow(&msgs, i), j)})
+        } else {
+            borrow_global_mut<ModuleData>(addr).state = *vector::borrow(vector::borrow(&msgs, i), j);
+        }
+    }
 }