fix: protect module account from freeze and burn (#33)

* protect module account from freeze and burn

* add test

* fmt

crates/natives/src/account.rs

@@ -44,6 +44,16 @@ pub struct NativeAccountContext<'a> {
     api: &'a dyn AccountAPI,
     new_accounts: BTreeMap<AccountAddress, (u64 /* account_number */, u8 /* account_type */)>,
     next_account_number: u64,
+
+    #[cfg(feature = "testing")]
+    test_accounts: BTreeMap<
+        AccountAddress,
+        (
+            u64, /* account_number */
+            u64, /* sequence */
+            u8,  /* account_type */
+        ),
+    >,
 }
 
 impl<'a> NativeAccountContext<'a> {
@@ -52,6 +62,9 @@ impl<'a> NativeAccountContext<'a> {
             api,
             new_accounts: Default::default(),
             next_account_number,
+
+            #[cfg(feature = "testing")]
+            test_accounts: Default::default(),
         }
     }
 
@@ -63,6 +76,18 @@ impl<'a> NativeAccountContext<'a> {
                 .collect::<Vec<(AccountAddress, u64, u8)>>(),
         )
     }
+
+    #[cfg(feature = "testing")]
+    pub fn set_account_info(
+        &mut self,
+        addr: AccountAddress,
+        account_number: u64,
+        sequence: u64,
+        account_type: u8,
+    ) {
+        self.test_accounts
+            .insert(addr, (account_number, sequence, account_type));
+    }
 }
 
 /***************************************************************************************************
@@ -94,6 +119,18 @@ fn native_get_account_info(
         if let Some(new_account) = account_context.new_accounts.get(&address) {
             (true, new_account.0, 0, new_account.1)
         } else {
+            #[cfg(feature = "testing")]
+            if let Some((account_number, sequence, account_type)) =
+                account_context.test_accounts.get(&address)
+            {
+                return Ok(smallvec![
+                    Value::bool(true),
+                    Value::u64(*account_number),
+                    Value::u64(*sequence),
+                    Value::u8(*account_type)
+                ]);
+            }
+
             account_context
                 .api
                 .get_account_info(address)
@@ -135,7 +172,7 @@ fn native_create_account(
         .create_account;
 
     debug_assert!(ty_args.is_empty());
-    debug_assert!(arguments.len() == 2);
+    debug_assert!(arguments.len() == 3);
 
     context.charge(gas_params.base_cost)?;
 
@@ -145,11 +182,17 @@ fn native_create_account(
             abort_code: UNKNOWN_ACCOUNT_TYPE,
         });
     }
+    let mut account_number = safely_pop_arg!(arguments, u64);
     let address = safely_pop_arg!(arguments, AccountAddress);
 
     let account_context = context.extensions_mut().get_mut::<NativeAccountContext>();
-    let account_number = account_context.next_account_number;
-    account_context.next_account_number += 1;
+
+    // if the account is not specified, use the next account number
+    if account_number == 0 {
+        account_number = account_context.next_account_number;
+        account_context.next_account_number += 1;
+    }
+
     account_context
         .new_accounts
         .insert(address, (account_number, account_type));
@@ -225,16 +268,49 @@ fn native_create_signer(
 pub fn make_all(
     builder: &SafeNativeBuilder,
 ) -> impl Iterator<Item = (String, NativeFunction)> + '_ {
-    let natives = [
+    let mut natives = vec![];
+    natives.extend([
         ("get_account_info", native_get_account_info as RawSafeNative),
         ("request_create_account", native_create_account),
         ("create_address", native_create_address),
         ("create_signer", native_create_signer),
-    ];
+    ]);
+
+    #[cfg(feature = "testing")]
+    natives.extend([(
+        "set_account_info",
+        native_test_only_set_account_info as RawSafeNative,
+    )]);
 
     builder.make_named_natives(natives)
 }
 
+#[cfg(feature = "testing")]
+fn native_test_only_set_account_info(
+    context: &mut SafeNativeContext,
+    ty_args: Vec<Type>,
+    mut arguments: VecDeque<Value>,
+) -> SafeNativeResult<SmallVec<[Value; 1]>> {
+    debug_assert!(ty_args.is_empty());
+    debug_assert!(arguments.len() == 4);
+
+    let account_type = safely_pop_arg!(arguments, u8);
+    let sequence = safely_pop_arg!(arguments, u64);
+    let account_number = safely_pop_arg!(arguments, u64);
+    let addr = safely_pop_arg!(arguments, AccountAddress);
+
+    let account_context = context.extensions_mut().get_mut::<NativeAccountContext>();
+    NativeAccountContext::set_account_info(
+        account_context,
+        addr,
+        account_number,
+        sequence,
+        account_type,
+    );
+
+    Ok(smallvec![])
+}
+
 // =========================================================================================
 // Helpers
 

crates/natives/src/staking.rs

@@ -157,10 +157,8 @@ fn native_delegate(
     }
 
     #[cfg(feature = "testing")]
-    if staking_context.share_ratios.contains_key(&validator) {
-        let ratios = staking_context.share_ratios.get(&validator).unwrap();
-        if ratios.contains_key(&metadata) {
-            let ratio = ratios.get(&metadata).unwrap();
+    if let Some(ratios) = staking_context.share_ratios.get(&validator) {
+        if let Some(ratio) = ratios.get(&metadata) {
             return Ok(smallvec![Value::u64(amount * ratio.0 / ratio.1)]);
         }
     }

precompile/modules/initia_stdlib/sources/account.move

@@ -28,23 +28,30 @@ module initia_std::account {
         let (found, _, _, _) = get_account_info(addr);
         assert!(!found, error::already_exists(EACCOUNT_ALREADY_EXISTS));
 
-        request_create_account(addr, ACCOUNT_TYPE_BASE)
+        request_create_account(addr, 0, ACCOUNT_TYPE_BASE)
     }
 
-   /// TableAccount is similar to CosmosSDK's ModuleAccount in concept, 
+    /// TableAccount is similar to CosmosSDK's ModuleAccount in concept, 
     /// as both cannot have a pubkey, there is no way to use the account externally.
     public(friend) fun create_table_account(addr: address): u64 {
-        let (found, _, _, _) = get_account_info(addr);
-        assert!(!found, error::already_exists(EACCOUNT_ALREADY_EXISTS));
+        let (found, account_number, sequence, account_type) = get_account_info(addr);
+        assert!(
+            !found || (account_type == ACCOUNT_TYPE_BASE && sequence == 0),
+            error::already_exists(EACCOUNT_ALREADY_EXISTS),
+        );
 
-        request_create_account(addr, ACCOUNT_TYPE_TABLE)
+        request_create_account(addr, account_number, ACCOUNT_TYPE_TABLE)
     }
 
     /// ObjectAccount is similar to CosmosSDK's ModuleAccount in concept, 
     /// as both cannot have a pubkey, there is no way to use the account externally.
     public(friend) fun create_object_account(addr: address): u64 {
-        let (found, account_number, _, account_type) = get_account_info(addr);
-        if (found) {
+        let (found, account_number, sequence, account_type) = get_account_info(addr);
+
+        // base account with sequence 0 is considered as not created.
+        if (!found || (account_type == ACCOUNT_TYPE_BASE && sequence == 0)) {
+            request_create_account(addr, account_number, ACCOUNT_TYPE_OBJECT)
+        } else {
             // When an Object is deleted, the ObjectAccount in CosmosSDK is designed 
             // not to be deleted in order to prevent unexpected issues. Therefore, 
             // in this case, the creation of an account is omitted.
@@ -55,8 +62,6 @@ module initia_std::account {
             } else {
                 abort(error::already_exists(EACCOUNT_ALREADY_EXISTS))
             }
-        } else {
-            request_create_account(addr, ACCOUNT_TYPE_OBJECT)
         }
     }
 
@@ -114,11 +119,14 @@ module initia_std::account {
         account_type == ACCOUNT_TYPE_MODULE
     }
 
-    native fun request_create_account(addr: address, account_type: u8): u64;
+    native fun request_create_account(addr: address, account_number: u64, account_type: u8): u64;
     native public fun get_account_info(addr: address): (bool /* found */, u64 /* account_number */, u64 /* sequence_number */, u8 /* account_type */);
     native public(friend) fun create_address(bytes: vector<u8>): address;
     native public(friend) fun create_signer(addr: address): signer;
 
+    #[test_only]
+    native public fun set_account_info(addr: address, account_number: u64, sequence: u64, account_type: u8);
+
     #[test_only]
     /// Create signer for testing
     public fun create_signer_for_test(addr: address): signer { create_signer(addr) }
@@ -174,4 +182,35 @@ module initia_std::account {
         let authentication_key = bcs::to_bytes(&new_address);
         assert!(vector::length(&authentication_key) == 32, 0);
     }
+
+    #[test(new_address = @0x41, new_address2 = @0x42, new_address3 = @0x43, new_address4 = @0x44)]
+    public fun test_create_table_account_and_object_account(
+        new_address: address, new_address2: address, new_address3: address, new_address4: address,
+    ) {
+        let table_account_num = create_table_account(new_address);
+        assert!(table_account_num == get_account_number(new_address), 0);
+        assert!(is_table_account(new_address), 1);
+        assert!(exists_at(new_address), 2);
+
+        // set base account with 0 sequence
+        set_account_info(new_address2, 100, 0, ACCOUNT_TYPE_BASE);
+        let table_account_num = create_table_account(new_address2);
+        assert!(table_account_num == get_account_number(new_address2), 0);
+        assert!(table_account_num == 100, 0);
+        assert!(is_table_account(new_address2), 1);
+        assert!(exists_at(new_address2), 2);
+
+        let object_account_num = create_object_account(new_address3);
+        assert!(object_account_num == get_account_number(new_address3), 3);
+        assert!(is_object_account(new_address3), 4);
+        assert!(exists_at(new_address3), 5);
+
+        // set base account with 0 sequence
+        set_account_info(new_address4, 200, 0, ACCOUNT_TYPE_BASE);
+        let object_account_num = create_object_account(new_address4);
+        assert!(object_account_num == get_account_number(new_address4), 0);
+        assert!(object_account_num == 200, 0);
+        assert!(is_object_account(new_address4), 1);
+        assert!(exists_at(new_address4), 2);
+    }
 }

precompile/modules/initia_stdlib/sources/fa/fungible_asset.move

@@ -8,6 +8,7 @@ module initia_std::fungible_asset {
     use std::option::{Self, Option};
     use std::signer;
     use std::string::{Self, String};
+    use std::account;
 
     /// The transfer ref and the fungible asset do not match.
     const ETRANSFER_REF_AND_FUNGIBLE_ASSET_MISMATCH: u64 = 2;
@@ -49,6 +50,8 @@ module initia_std::fungible_asset {
     const ESUPPLY_UNDERFLOW: u64 = 20;
     /// Supply resource is not found for a metadata object.
     const ESUPPLY_NOT_FOUND: u64 = 21;
+    /// Module account store cannot be manipulated.
+    const ECONNOT_MANIPULATE_MODULE_ACCOUNT_STORE: u64 = 22;
 
     //
     // Constants
@@ -440,6 +443,10 @@ module initia_std::fungible_asset {
 
         let metadata_addr = object::object_address(ref.metadata);
         let store_addr = object::object_address(store);
+
+        // cannot freeze module account store
+        assert!(!is_module_account_store_addr(store_addr), error::invalid_argument(ECONNOT_MANIPULATE_MODULE_ACCOUNT_STORE));
+
         borrow_global_mut<FungibleStore>(store_addr).frozen = frozen;
 
         // emit event
@@ -470,6 +477,10 @@ module initia_std::fungible_asset {
         assert!(metadata == store_metadata(store), error::invalid_argument(EBURN_REF_AND_STORE_MISMATCH));
 
         let store_addr = object::object_address(store);
+
+        // cannot burn module account funds
+        assert!(!is_module_account_store_addr(store_addr), error::invalid_argument(ECONNOT_MANIPULATE_MODULE_ACCOUNT_STORE));
+
         burn(ref, withdraw_internal(store_addr, amount));
     }
 
@@ -484,6 +495,10 @@ module initia_std::fungible_asset {
             error::invalid_argument(ETRANSFER_REF_AND_STORE_MISMATCH),
         );
 
+        // cannot withdraw module account funds
+        let store_addr = object::object_address(store);
+        assert!(!is_module_account_store_addr(store_addr), error::invalid_argument(ECONNOT_MANIPULATE_MODULE_ACCOUNT_STORE));
+
         withdraw_internal(object::object_address(store), amount)
     }
 
@@ -610,6 +625,12 @@ module initia_std::fungible_asset {
         supply.current = supply.current - (amount as u128);
     }
 
+    fun is_module_account_store_addr(store_addr: address): bool {
+        let fungible_store = object::address_to_object<FungibleStore>(store_addr);
+        let owner_addr = object::owner(fungible_store);
+        account::exists_at(owner_addr) && account::is_module_account(owner_addr)
+    }
+
     inline fun borrow_fungible_metadata<T: key>(
         metadata: Object<T>
     ): &Metadata acquires Metadata {
@@ -628,7 +649,6 @@ module initia_std::fungible_asset {
         borrow_global<FungibleStore>(object::object_address(store))
     }
 
-
     #[test_only]
     struct TestToken has key {}
 
@@ -813,4 +833,45 @@ module initia_std::fungible_asset {
             amount: _
         } = base;
     }
+
+    #[test(creator = @0xcafe, module_acc = @0x123)]
+    #[expected_failure(abort_code = 0x10016, location = Self)]
+    fun test_freeze_module_account_store(creator: &signer, module_acc: &signer) acquires FungibleStore {
+        let (mint_ref, transfer_ref, _burn_ref, _) = create_fungible_asset(creator);
+        let metadata = mint_ref.metadata;
+
+        let module_acc_store = create_test_store(module_acc, metadata);
+        account::set_account_info(signer::address_of(module_acc), 10, 0, 3);
+
+        set_frozen_flag(&transfer_ref, module_acc_store, true);
+    }
+
+    #[test(creator = @0xcafe, module_acc = @0x123)]
+    #[expected_failure(abort_code = 0x10016, location = Self)]
+    fun test_burn_module_account_funds(creator: &signer, module_acc: &signer) acquires FungibleStore, Supply {
+        let (mint_ref, _transfer_ref, burn_ref, _) = create_fungible_asset(creator);
+        let metadata = mint_ref.metadata;
+
+        let module_acc_store = create_test_store(module_acc, metadata);
+        account::set_account_info(signer::address_of(module_acc), 10, 0, 3);
+
+        let fa = mint(&mint_ref, 100);
+        deposit(module_acc_store, fa);
+        burn_from(&burn_ref, module_acc_store, 30);
+    }
+
+    #[test(creator = @0xcafe, module_acc = @0x123)]
+    #[expected_failure(abort_code = 0x10016, location = Self)]
+    fun test_withdraw_module_account_funds_with_ref(creator: &signer, module_acc: &signer) acquires FungibleStore, Supply {
+        let (mint_ref, transfer_ref, _burn_ref, _) = create_fungible_asset(creator);
+        let metadata = mint_ref.metadata;
+
+        let module_acc_store = create_test_store(module_acc, metadata);
+        account::set_account_info(signer::address_of(module_acc), 10, 0, 3);
+
+        let fa = mint(&mint_ref, 100);
+        deposit(module_acc_store, fa);
+        let fa = withdraw_with_ref(&transfer_ref, module_acc_store, 30);
+        deposit(module_acc_store, fa);
+    }
 }