check account code id on rebalancing

README.md

@@ -5,27 +5,21 @@ Valence Services houses the code for the [Valence Rebalancer](https://www.valenc
 
 Learn more about the Rebalancer [here](https://www.valence.zone/blog/Rebalancer-Protocol-Asset-Management).
 
-This repository contains the following contracts:
+This repository contains the smart contracts for the Valence Services. It includes the following contracts:
 - Auctions Manager
 - Auction
 - Oracle
 - Services Manager
 - Rebalancer
 - Account
 
-The following is a simplified diagram of the Valence Services system. Detailed system diagrams are presented in [Valence Services Architecture](./architecture.md)
+Here is a high level of the system for more details please refer to the [Valence Services Architecture](./architecture.md)
 
 ![Top Level](./images/high-level.png)
 
-For documentation on specific contracts, please refer to the README.md file in each contract's directory.
+For a specific contract documentation, please refer to the README.md file of each contract.
 
-### ⚖️ DAO Treasuries
-The Rebalancer is purpose built for treasury managment. If you are a member of DAO that is interested in using the Rebalancer, please send us a [direct message on X](https://x.com/TimewaveLabs).
-
-### 💸 Market Makers
-The Rebalancer sends funds to be auctioned daily. If you are a market maker interested in bidding in these auctions and would like support, please send us a [direct message on X](https://x.com/TimewaveLabs). For more information on auctions:
-- Read the [Auctions contract documentation](./contracts//auction/auction/README.md)
-- Check out Timewave's [Auction Arbitrage Bot](https://github.com/timewave-computer/auction-arbitrage-bot) that can be used to arbitrage auctions and various DEXs
+Parties (e.g. market makers) interested to participate in the Rebalancer Auctions should read the [Auctions contract documentation](./contracts//auction/auction/README.md).
 
 ## Neutron Mainnet Contracts
 
@@ -61,4 +55,4 @@ If you believe you've found a security-related issue with the contracts associat
 
 ## License
 
-All materials in this repository are licensed under [BSL](./LICENSE).
+All materials in this repository are licensed under [BSL](./LICENSE).

contracts/services/rebalancer/src/contract.rs

@@ -683,6 +683,9 @@ pub fn query(deps: Deps, _env: Env, msg: QueryMsg) -> StdResult<Binary> {
                         valence_package::services::rebalancer::PauseReason::EmptyBalance => {
                             fees.resume_fee
                         }
+                        valence_package::services::rebalancer::PauseReason::NotWhitelistedAccountCodeId(_) => {
+                            fees.resume_fee
+                        }
                         valence_package::services::rebalancer::PauseReason::AccountReason(_) => {
                             Uint128::zero()
                         }

contracts/services/rebalancer/src/rebalance.rs

@@ -17,6 +17,7 @@ use valence_package::{
         ParsedPID, PauseData, RebalanceTrade, RebalancerConfig, SystemRebalanceStatus,
         TargetOverrideStrategy,
     },
+    states::{ACCOUNT_WHITELISTED_CODE_IDS, SERVICES_MANAGER},
     CLOSEST_TO_ONE_POSSIBLE,
 };
 
@@ -111,6 +112,8 @@ pub fn execute_system_rebalance(
     let mut min_amount_limits: Vec<(String, Uint128)> = vec![];
     let mut msgs: Vec<SubMsg> = vec![];
     let mut account_events: Vec<Event> = vec![];
+    let services_manager_addr = SERVICES_MANAGER.load(deps.storage)?;
+    let whitelist = ACCOUNT_WHITELISTED_CODE_IDS.query(&deps.querier, services_manager_addr)?;
 
     for res in configs {
         let Ok((account, config)) = res else {
@@ -119,6 +122,23 @@ pub fn execute_system_rebalance(
 
         last_addr = Some(account.clone());
 
+        // Before rebalancing, verify the account is using a whitelisted code id
+        let account_code_id = deps
+            .querier
+            .query_wasm_contract_info(account.clone())?
+            .code_id;
+
+        if !whitelist.contains(&account_code_id) {
+            // Save to the paused config
+            PAUSED_CONFIGS.save(
+                deps.storage,
+                account.clone(),
+                &PauseData::new_not_whitelisted_account_code_id(env, account_code_id, &config),
+            )?;
+            // remove from active configs
+            CONFIGS.remove(deps.storage, account.clone());
+        }
+
         // Do rebalance for the account, and construct the msg
         let rebalance_res = do_rebalance(
             deps.as_ref(),

contracts/services_manager/Cargo.toml

@@ -25,7 +25,7 @@ serde           = { workspace = true }
 thiserror       = { workspace = true }
 valence-macros  = { workspace = true }
 valence-package = { workspace = true }
-rebalancer = { workspace = true }
+rebalancer      = { workspace = true }
 
 [dev-dependencies]
 cw-multi-test = { workspace = true }

contracts/services_manager/src/contract.rs

@@ -12,12 +12,12 @@ use valence_package::helpers::approve_admin_change;
 use valence_package::msgs::core_execute::ServicesManagerExecuteMsg;
 use valence_package::msgs::core_query::ServicesManagerQueryMsg;
 use valence_package::services::rebalancer::RebalancerConfig;
-use valence_package::states::ADMIN;
+use valence_package::states::{ACCOUNT_WHITELISTED_CODE_IDS, ADMIN};
 
 use crate::error::ContractError;
 use crate::helpers::{get_service_addr, save_service};
 use crate::msg::{InstantiateMsg, MigrateMsg};
-use crate::state::{ACCOUNT_WHITELISTED_CODE_IDS, ADDR_TO_SERVICES, SERVICES_TO_ADDR};
+use crate::state::{ADDR_TO_SERVICES, SERVICES_TO_ADDR};
 
 const CONTRACT_NAME: &str = "crates.io:services-manager";
 const CONTRACT_VERSION: &str = env!("CARGO_PKG_VERSION");
@@ -82,6 +82,16 @@ pub fn execute(
         ServicesManagerExecuteMsg::UpdateService { service_name, data } => {
             let service_addr = get_service_addr(deps.as_ref(), service_name.to_string())?;
 
+            let sender_code_id = deps
+                .querier
+                .query_wasm_contract_info(info.sender.clone())?
+                .code_id;
+            let whitelist = ACCOUNT_WHITELISTED_CODE_IDS.load(deps.storage)?;
+
+            if !whitelist.contains(&sender_code_id) {
+                return Err(ContractError::NotWhitelistedContract(sender_code_id));
+            }
+
             let msg = service_name.get_update_msg(&info, service_addr.as_ref(), data)?;
 
             Ok(Response::default().add_message(msg))
@@ -102,6 +112,17 @@ pub fn execute(
             resume_for,
         } => {
             let service_addr = get_service_addr(deps.as_ref(), service_name.to_string())?;
+
+            let sender_code_id = deps
+                .querier
+                .query_wasm_contract_info(resume_for.clone())?
+                .code_id;
+            let whitelist = ACCOUNT_WHITELISTED_CODE_IDS.load(deps.storage)?;
+
+            if !whitelist.contains(&sender_code_id) {
+                return Err(ContractError::NotWhitelistedContract(sender_code_id));
+            }
+
             let msg = service_name.get_resume_msg(resume_for, &info, service_addr.as_ref())?;
 
             Ok(Response::default().add_message(msg))
@@ -226,6 +247,10 @@ pub fn query(deps: Deps, _env: Env, msg: ServicesManagerQueryMsg) -> StdResult<B
             let is_service = ADDR_TO_SERVICES.has(deps.storage, deps.api.addr_validate(&addr)?);
             to_json_binary(&is_service)
         }
+        ServicesManagerQueryMsg::IsAccountCodeId { code_id } => {
+            let code_ids = ACCOUNT_WHITELISTED_CODE_IDS.load(deps.storage)?;
+            to_json_binary(&code_ids.contains(&code_id))
+        }
         ServicesManagerQueryMsg::GetServiceAddr { service } => {
             let addr = get_service_addr(deps, service.to_string())
                 .map_err(|e| StdError::GenericErr { msg: e.to_string() })?;

contracts/services_manager/src/state.rs

@@ -1,9 +1,5 @@
-use std::collections::HashSet;
-
 use cosmwasm_std::Addr;
-use cw_storage_plus::{Item, Map};
+use cw_storage_plus::Map;
 
 pub(crate) const SERVICES_TO_ADDR: Map<String, Addr> = Map::new("services_to_addr");
 pub(crate) const ADDR_TO_SERVICES: Map<Addr, String> = Map::new("addr_to_services");
-
-pub(crate) const ACCOUNT_WHITELISTED_CODE_IDS: Item<HashSet<u64>> = Item::new("wl_code_ids");

packages/valence-macros/src/lib.rs

@@ -173,6 +173,10 @@ pub fn valence_services_manager_query_msgs(
         IsService {
             addr: String,
         },
+        #[returns(bool)]
+        IsAccountCodeId {
+            code_id: u64,
+        },
         /// Get the address of a service
         #[returns(Addr)]
         GetServiceAddr {

packages/valence-package/src/services/rebalancer.rs

@@ -149,19 +149,39 @@ impl PauseData {
             config: config.clone(),
         }
     }
+
+    pub fn new_not_whitelisted_account_code_id(
+        env: &Env,
+        code_id: u64,
+        config: &RebalancerConfig,
+    ) -> Self {
+        Self {
+            pauser: env.contract.address.clone(),
+            reason: PauseReason::NotWhitelistedAccountCodeId(code_id),
+            config: config.clone(),
+        }
+    }
 }
 
 #[cw_serde]
 pub enum PauseReason {
     /// This reason can only be called if the rebalancer is pausing the account because it
     /// has an empty balance.
     EmptyBalance,
+    NotWhitelistedAccountCodeId(u64),
     /// This reason is given by the user/account, he might forget why he paused the account
     /// this will remind him of it.
     AccountReason(String),
 }
 
 impl PauseReason {
+    pub fn should_pay_fee(&self) -> bool {
+        matches!(
+            self,
+            PauseReason::EmptyBalance | PauseReason::NotWhitelistedAccountCodeId(_)
+        )
+    }
+
     pub fn is_empty_balance(&self) -> bool {
         matches!(self, PauseReason::EmptyBalance)
     }
@@ -395,7 +415,7 @@ impl ServiceFeeConfig {
         let mut msgs: Vec<CosmosMsg> = Vec::with_capacity(1);
 
         if !self.resume_fee.is_zero() {
-            if !reason.is_empty_balance() {
+            if !reason.should_pay_fee() {
                 return Ok(msgs);
             }
 

packages/valence-package/src/states.rs

@@ -1,3 +1,5 @@
+use std::collections::HashSet;
+
 use cosmwasm_schema::cw_serde;
 use cosmwasm_std::Addr;
 use cw_storage_plus::Item;
@@ -14,6 +16,9 @@ pub const ADMIN: Item<Addr> = Item::new("admin");
 /// State for when an admin want to set someone else as admin.
 pub const ADMIN_CHANGE: Item<AdminChange> = Item::new("admin_change");
 
+/// Whitelisted account code ids that is stored on the services manager
+pub const ACCOUNT_WHITELISTED_CODE_IDS: Item<HashSet<u64>> = Item::new("wl_code_ids");
+
 #[cw_serde]
 pub struct AdminChange {
     pub addr: Addr,

tests/rust-tests/src/live_debugging/suite_builder.rs

@@ -68,6 +68,7 @@ impl SuiteBuilder {
                 whitelists.denom_whitelist[1].clone(),
             )),
             account_code_id: builder.account_code_id,
+            rebalancer_code_id: builder.rebalancer_code_id,
             astro_pools: HashMap::new(),
         }
     }

tests/rust-tests/src/suite/contracts.rs

@@ -26,7 +26,8 @@ pub fn rebalancer_contract() -> Box<dyn Contract<Empty>> {
         rebalancer::contract::instantiate,
         rebalancer::contract::query,
     )
-    .with_reply(rebalancer::contract::reply);
+    .with_reply(rebalancer::contract::reply)
+    .with_migrate(rebalancer::contract::migrate);
     Box::new(contract)
 }
 

tests/rust-tests/src/suite/suite.rs

@@ -56,6 +56,7 @@ pub(crate) struct Suite {
 
     // code ids for future use
     pub account_code_id: u64,
+    pub rebalancer_code_id: u64,
 
     // astro
     pub astro_pools: HashMap<(String, String), Addr>,

tests/rust-tests/src/suite/suite_builder.rs

@@ -562,6 +562,7 @@ impl SuiteBuilder {
             account_addrs,
             auction_addrs,
             account_code_id: self.account_code_id,
+            rebalancer_code_id: self.rebalancer_code_id,
             pair: Pair::from((ATOM.to_string(), NTRN.to_string())),
             astro_pools: pools,
         }

tests/rust-tests/src/tests_rebalancer/test_management.rs

@@ -755,3 +755,29 @@ fn test_fee_withdraw() {
         admin_initial_balance.amount + Uint128::new(100_u128)
     );
 }
+
+#[test]
+fn test_not_whitelisted_account_after_register() {
+    let mut suite = Suite::default();
+
+    let account_addr = suite.account_addrs[0].clone();
+
+    // migrate the account to rebalancer code id (not a whitelisted account code id)
+    suite
+        .app
+        .migrate_contract(
+            suite.owner.clone(),
+            account_addr.clone(),
+            &valence_account::msg::MigrateMsg::NoStateChange {},
+            suite.rebalancer_code_id,
+        )
+        .unwrap();
+
+    suite.rebalance(None).unwrap();
+
+    let paused_config = suite.query_rebalancer_paused_config(account_addr).unwrap();
+    assert_eq!(
+        paused_config.reason,
+        PauseReason::NotWhitelistedAccountCodeId(suite.rebalancer_code_id)
+    );
+}