Register the sub account on the main account on creation (#55)

* add sub_accounts to storage on manager

* sub_accounts_move_ownership test

* update in state in case of moved ownership

* update state of sub-accounts

* add new tests, update old ones

* update comment

* fix the verify method

* changelog update

* update fix log

* remove reply, wrap sub-account actions for manager

* clean up of unused stuff

* rename variant to be more clear

framework/CHANGELOG.md

@@ -9,13 +9,15 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
 
 ### Added
 - Install modules on account or Sub-account creation
+- Manager stores his sub-accounts and sub-accounts can register or unregister in case of ownership change
 
 ### Changed
 - Updated fetch_data arguments of CwStakingCommand
 - Owner of the sub-accounts now Proxy, allowing modules to interact with sub-accounts
 
 ### Fixed
 - Partially fixed cw-staking for Osmosis
+- Manager governance now changes only after new "owner" claimed ownership
 
 ## [0.17.2] - 2023-07-27
 

framework/contracts/account/manager/src/commands.rs

@@ -1,15 +1,18 @@
 use crate::{contract::ManagerResult, error::ManagerError, queries::query_module_cw2};
 use crate::{validation, versioning};
-use abstract_core::manager::state::{ACCOUNT_FACTORY, MODULE_QUEUE};
+use abstract_core::manager::state::{
+    ACCOUNT_FACTORY, MODULE_QUEUE, PENDING_GOVERNANCE, SUB_ACCOUNTS,
+};
 
 use abstract_core::adapter::{
     AuthorizedAddressesResponse, BaseExecuteMsg, BaseQueryMsg, ExecuteMsg as AdapterExecMsg,
     QueryMsg as AdapterQuery,
 };
-use abstract_core::manager::InternalConfigAction;
+use abstract_core::manager::{InternalConfigAction, UpdateSubAccountAction};
 use abstract_core::objects::gov_type::GovernanceDetails;
 use abstract_core::objects::AssetEntry;
 
+use abstract_core::proxy::state::ACCOUNT_ID;
 use abstract_core::version_control::ModuleResponse;
 use abstract_macros::abstract_response;
 use abstract_sdk::cw_helpers::AbstractAttributes;
@@ -40,6 +43,7 @@ use cosmwasm_std::{
     DepsMut, Empty, Env, MessageInfo, Response, StdError, StdResult, Storage, WasmMsg,
 };
 use cw2::{get_contract_version, ContractVersion};
+use cw_ownable::OwnershipError;
 use cw_storage_plus::Item;
 use semver::Version;
 
@@ -211,7 +215,7 @@ pub fn exec_on_module(
 
 #[allow(clippy::too_many_arguments)]
 /// Creates a sub-account for this account,
-pub fn create_subaccount(
+pub fn create_sub_account(
     deps: DepsMut,
     msg_info: MessageInfo,
     env: Env,
@@ -258,6 +262,64 @@ pub fn create_subaccount(
     Ok(response)
 }
 
+pub fn handle_sub_account_action(
+    deps: DepsMut,
+    info: MessageInfo,
+    action: UpdateSubAccountAction,
+) -> ManagerResult {
+    match action {
+        UpdateSubAccountAction::UnregisterSubAccount { id } => {
+            unregister_sub_account(deps, info, id)
+        }
+        UpdateSubAccountAction::RegisterSubAccount { id } => register_sub_account(deps, info, id),
+        _ => unimplemented!(),
+    }
+}
+
+// Unregister sub-account from the state
+fn unregister_sub_account(deps: DepsMut, info: MessageInfo, id: u32) -> ManagerResult {
+    let config = CONFIG.load(deps.storage)?;
+
+    let account = abstract_core::version_control::state::ACCOUNT_ADDRESSES.query(
+        &deps.querier,
+        config.version_control_address,
+        id,
+    )?;
+
+    if account.is_some_and(|a| a.manager == info.sender) {
+        SUB_ACCOUNTS.remove(deps.storage, id);
+
+        Ok(ManagerResponse::new(
+            "unregister_sub_account",
+            vec![("sub_account_removed", id.to_string())],
+        ))
+    } else {
+        Err(ManagerError::SubAccountRemovalFailed {})
+    }
+}
+
+// Register sub-account to the state
+fn register_sub_account(deps: DepsMut, info: MessageInfo, id: u32) -> ManagerResult {
+    let config = CONFIG.load(deps.storage)?;
+
+    let account = abstract_core::version_control::state::ACCOUNT_ADDRESSES.query(
+        &deps.querier,
+        config.version_control_address,
+        id,
+    )?;
+
+    if account.is_some_and(|a| a.manager == info.sender) {
+        SUB_ACCOUNTS.save(deps.storage, id, &Empty {})?;
+
+        Ok(ManagerResponse::new(
+            "register_sub_account",
+            vec![("sub_account_added", id.to_string())],
+        ))
+    } else {
+        Err(ManagerError::SubAccountRegisterFailed {})
+    }
+}
+
 /// Checked load of a module address
 fn load_module_addr(storage: &dyn Storage, module_id: &String) -> Result<Addr, ManagerError> {
     ACCOUNT_MODULES
@@ -306,36 +368,84 @@ pub fn set_owner(
     info: MessageInfo,
     new_owner: GovernanceDetails<String>,
 ) -> ManagerResult {
+    assert_admin_right(deps.as_ref(), &info.sender)?;
+    // In case it's a top level owner we need to pass current owner into update_ownership method
+    let owner = cw_ownable::get_ownership(deps.storage)?
+        .owner
+        .ok_or(cw_ownable::OwnershipError::NoOwner)?;
     // verify the provided governance details
-    let verified_gov = new_owner.verify(deps.api)?;
+    let config = CONFIG.load(deps.storage)?;
+    let verified_gov = new_owner.verify(deps.as_ref(), config.version_control_address)?;
     let new_owner_addr = verified_gov.owner_address();
 
-    // Update the account information
-    let mut acc_info = INFO.load(deps.storage)?;
-
     // Check that there are changes
+    let acc_info = INFO.load(deps.storage)?;
     if acc_info.governance_details == verified_gov {
         return Err(ManagerError::NoUpdates {});
     }
 
-    acc_info.governance_details = verified_gov.clone();
-    INFO.save(deps.storage, &acc_info)?;
+    let mut response = ManagerResponse::new(
+        "update_owner",
+        vec![("governance_type", verified_gov.to_string())],
+    );
 
+    PENDING_GOVERNANCE.save(deps.storage, &verified_gov)?;
     // Update the Owner of the Account
     let ownership = cw_ownable::update_ownership(
         deps,
         &env.block,
-        &info.sender,
+        &owner,
         cw_ownable::Action::TransferOwnership {
             new_owner: new_owner_addr.into_string(),
             expiry: None,
         },
     )?;
+    response = response.add_attributes(ownership.into_attributes());
+    Ok(response)
+}
 
-    let mut attrs = vec![("governance_type", verified_gov.to_string()).into()];
-    attrs.extend(ownership.into_attributes());
+/// Update governance of this account after claim
+pub(crate) fn update_governance(storage: &mut dyn Storage) -> ManagerResult<Vec<CosmosMsg>> {
+    let mut msgs = vec![];
+    let mut acc_info = INFO.load(storage)?;
+    let mut account_id = None;
+    // Get pending governance and clear it
+    let pending_governance = PENDING_GOVERNANCE
+        .may_load(storage)?
+        .ok_or(OwnershipError::TransferNotFound)?;
+    PENDING_GOVERNANCE.remove(storage);
+
+    // Clear state for previous manager if it was sub-account
+    if let GovernanceDetails::SubAccount { manager, .. } = acc_info.governance_details {
+        let id = ACCOUNT_ID.load(storage)?;
+        // For optimizing the gas we save it, in case new owner is sub-account as well
+        account_id = Some(id);
+        let unregister_message = wasm_execute(
+            manager,
+            &ExecuteMsg::UpdateSubAccount(UpdateSubAccountAction::UnregisterSubAccount { id }),
+            vec![],
+        )?;
+        msgs.push(unregister_message.into());
+    }
 
-    Ok(ManagerResponse::new("update_owner", attrs))
+    // Update state for new manager if owner will be the sub-account
+    if let GovernanceDetails::SubAccount { manager, .. } = &pending_governance {
+        let id = if let Some(id) = account_id {
+            id
+        } else {
+            ACCOUNT_ID.load(storage)?
+        };
+        let register_message = wasm_execute(
+            manager,
+            &ExecuteMsg::UpdateSubAccount(UpdateSubAccountAction::RegisterSubAccount { id }),
+            vec![],
+        )?;
+        msgs.push(register_message.into());
+    }
+    // Update governance of this account
+    acc_info.governance_details = pending_governance;
+    INFO.save(storage, &acc_info)?;
+    Ok(msgs)
 }
 
 /// Migrate modules through address updates or contract migrations
@@ -839,12 +949,26 @@ pub fn update_internal_config(
             })
             .collect::<ManagerResult<Vec<CosmosMsg>>>();
 
+        let mut response =
+            ManagerResponse::action("manager_after_init").add_messages(install_msgs?);
+
+        let account_info = INFO.load(deps.storage)?;
+        if let GovernanceDetails::SubAccount { manager, .. } = account_info.governance_details {
+            response = response.add_message(wasm_execute(
+                manager,
+                &ExecuteMsg::UpdateSubAccount(UpdateSubAccountAction::RegisterSubAccount {
+                    id: ACCOUNT_ID.load(deps.storage)?,
+                }),
+                vec![],
+            )?);
+        }
+
         // clear the queue
         MODULE_QUEUE.remove(deps.storage);
         // Remove account factory from storage
         ACCOUNT_FACTORY.set(deps, None)?;
 
-        Ok(ManagerResponse::action("manager_after_init").add_messages(install_msgs?))
+        Ok(response)
     } else {
         assert_admin_right(deps.as_ref(), &info.sender)?;
         update_module_addresses(deps, add, remove)
@@ -1041,11 +1165,20 @@ mod tests {
             let new_gov = "new_gov".to_string();
 
             let msg = ExecuteMsg::SetOwner {
-                owner: GovernanceDetails::Monarchy { monarch: new_gov },
+                owner: GovernanceDetails::Monarchy {
+                    monarch: new_gov.clone(),
+                },
             };
 
             execute_as_owner(deps.as_mut(), msg)?;
 
+            let actual_info = INFO.load(deps.as_ref().storage)?;
+            assert_that!(&actual_info.governance_details.owner_address().to_string())
+                .is_equal_to("owner".to_string());
+
+            let accept_msg = ExecuteMsg::UpdateOwnership(cw_ownable::Action::AcceptOwnership);
+            execute_as(deps.as_mut(), &new_gov, accept_msg)?;
+
             let actual_info = INFO.load(deps.as_ref().storage)?;
             assert_that!(&actual_info.governance_details.owner_address().to_string())
                 .is_equal_to("new_gov".to_string());
@@ -1880,6 +2013,13 @@ mod tests {
                     pending_owner: Some(Addr::unchecked(pending_owner)),
                 },
             )?;
+            // mock pending governance
+            Item::new("pgov").save(
+                deps.as_mut().storage,
+                &GovernanceDetails::Monarchy {
+                    monarch: pending_owner.to_owned(),
+                },
+            )?;
 
             let msg = ExecuteMsg::UpdateOwnership(cw_ownable::Action::AcceptOwnership {});
 

framework/contracts/account/manager/src/contract.rs

@@ -1,7 +1,7 @@
 use crate::{
     commands::*,
     error::ManagerError,
-    queries,
+    queries::{self, handle_sub_accounts_query},
     queries::{handle_account_info_query, handle_config_query, handle_module_info_query},
     versioning,
 };
@@ -44,11 +44,13 @@ pub fn instantiate(
 ) -> ManagerResult {
     set_contract_version(deps.storage, MANAGER, CONTRACT_VERSION)?;
     let module_factory_address = deps.api.addr_validate(&msg.module_factory_address)?;
+    let version_control_address = deps.api.addr_validate(&msg.version_control_address)?;
+
     ACCOUNT_ID.save(deps.storage, &msg.account_id)?;
     CONFIG.save(
         deps.storage,
         &Config {
-            version_control_address: deps.api.addr_validate(&msg.version_control_address)?,
+            version_control_address: version_control_address.clone(),
             module_factory_address: module_factory_address.clone(),
         },
     )?;
@@ -58,12 +60,12 @@ pub fn instantiate(
     validate_link(&msg.link)?;
     validate_name(&msg.name)?;
 
-    let governance_details = msg.owner.verify(deps.api)?;
+    let governance_details = msg.owner.verify(deps.as_ref(), version_control_address)?;
     let owner = governance_details.owner_address();
 
     let account_info = AccountInfo {
         name: msg.name,
-        governance_details,
+        governance_details: governance_details.clone(),
         chain_id: env.block.chain_id,
         description: msg.description,
         link: msg.link,
@@ -131,7 +133,7 @@ pub fn execute(deps: DepsMut, env: Env, info: MessageInfo, msg: ExecuteMsg) -> M
                     base_asset,
                     namespace,
                     install_modules,
-                } => create_subaccount(
+                } => create_sub_account(
                     deps,
                     info,
                     env,
@@ -161,22 +163,28 @@ pub fn execute(deps: DepsMut, env: Env, info: MessageInfo, msg: ExecuteMsg) -> M
 
                     Ok(response)
                 }
+                ExecuteMsg::UpdateSubAccount(action) => {
+                    handle_sub_account_action(deps, info, action)
+                }
                 ExecuteMsg::Callback(CallbackMsg {}) => handle_callback(deps, env, info),
-                ExecuteMsg::UpdateOwnership(action) => match action {
-                    // Disallow the user from using the TransferOwnership action
-                    cw_ownable::Action::TransferOwnership { .. } => {
-                        Err(ManagerError::MustUseSetOwner {})
-                    }
-                    _ => {
-                        abstract_sdk::execute_update_ownership!(
-                            ManagerResponse,
-                            deps,
-                            env,
-                            info,
-                            action
-                        )
-                    }
-                },
+                ExecuteMsg::UpdateOwnership(action) => {
+                    let msgs = match action {
+                        // Disallow the user from using the TransferOwnership action
+                        cw_ownable::Action::TransferOwnership { .. } => {
+                            return Err(ManagerError::MustUseSetOwner {});
+                        }
+                        cw_ownable::Action::AcceptOwnership => update_governance(deps.storage)?,
+                        _ => vec![],
+                    };
+                    let result: ManagerResult = abstract_sdk::execute_update_ownership!(
+                        ManagerResponse,
+                        deps,
+                        env,
+                        info,
+                        action
+                    );
+                    Ok(result?.add_messages(msgs))
+                }
                 _ => panic!(),
             }
         }
@@ -194,6 +202,9 @@ pub fn query(deps: Deps, env: Env, msg: QueryMsg) -> StdResult<Binary> {
         QueryMsg::Info {} => handle_account_info_query(deps),
         QueryMsg::Config {} => handle_config_query(deps),
         QueryMsg::Ownership {} => abstract_sdk::query_ownership!(deps),
+        QueryMsg::SubAccountIds { start_after, limit } => {
+            handle_sub_accounts_query(deps, start_after, limit)
+        }
     }
 }
 

framework/contracts/account/manager/src/error.rs

@@ -93,4 +93,10 @@ pub enum ManagerError {
             You either have the an error in your sub-account configuration or you are not authorized to make this call.
     ")]
     SubAccountAdminVerification,
+
+    #[error("Removing sub account failed")]
+    SubAccountRemovalFailed {},
+
+    #[error("Register of sub account failed")]
+    SubAccountRegisterFailed {},
 }