feat: add restricted asset management and creator functionality

contracts/oraiswap_factory/src/contract.rs

@@ -12,12 +12,14 @@ use oraiswap::error::ContractError;
 use oraiswap::querier::query_pair_info_from_pair;
 use oraiswap::response::MsgInstantiateContractResponse;
 
-use crate::state::{read_pairs, Config, CONFIG, PAIRS};
+use crate::state::{
+    read_pairs, Config, Creator, RestrictedAssets, CONFIG, CREATOR, PAIRS, RESTRICTED_ASSETS,
+};
 
 use oraiswap::asset::{pair_key, Asset, AssetInfo, PairInfo, PairInfoRaw};
 use oraiswap::factory::{
-    ConfigResponse, ExecuteMsg, InstantiateMsg, MigrateMsg, PairsResponse, ProvideLiquidityParams,
-    QueryMsg,
+    ConfigResponse, CreatorsResponse, ExecuteMsg, InstantiateMsg, MigrateMsg, PairsResponse,
+    ProvideLiquidityParams, QueryMsg, RestrictedAssetResponse,
 };
 use oraiswap::pair::{
     InstantiateMsg as PairInstantiateMsg, DEFAULT_COMMISSION_RATE, DEFAULT_OPERATOR_FEE,
@@ -87,7 +89,69 @@ pub fn execute(
         ExecuteMsg::ProvideLiquidity { assets, receiver } => {
             execute_provide_liquidity(deps, env, info, assets, receiver)
         }
+        ExecuteMsg::RestrictAsset { prefix } => execute_restrict_asset(deps, info, prefix),
+        ExecuteMsg::AddCreator { address } => add_creator(deps, info, address),
+        ExecuteMsg::RemoveCreator { address } => remove_creator(deps, info, address),
+    }
+}
+
+pub fn add_creator(
+    deps: DepsMut,
+    info: MessageInfo,
+    address: Addr,
+) -> Result<Response, ContractError> {
+    let config: Config = CONFIG.load(deps.storage)?;
+    // permission check
+    if deps.api.addr_canonicalize(info.sender.as_str())? != config.owner {
+        return Err(ContractError::Unauthorized {});
+    }
+
+    let mut creators = CREATOR.may_load(deps.storage)?.unwrap_or(Creator {
+        whitelist_addresses: vec![],
+    });
+    if creators.whitelist_addresses.contains(&address) {
+        return Err(ContractError::CreatorAlreadyExists {});
     }
+    creators.whitelist_addresses.push(address.clone());
+
+    CREATOR.save(deps.storage, &creators)?;
+
+    let res = Response::new()
+        .add_attribute("method", "add_creator")
+        .add_attribute("creator", address.to_string());
+
+    Ok(res)
+}
+
+pub fn remove_creator(
+    deps: DepsMut,
+    info: MessageInfo,
+    address: Addr,
+) -> Result<Response, ContractError> {
+    let config: Config = CONFIG.load(deps.storage)?;
+    // permission check
+    if deps.api.addr_canonicalize(info.sender.as_str())? != config.owner {
+        return Err(ContractError::Unauthorized {});
+    }
+
+    let mut creators = CREATOR.load(deps.storage)?;
+    if let Some(pos) = creators
+        .whitelist_addresses
+        .iter()
+        .position(|x| x == &address)
+    {
+        creators.whitelist_addresses.remove(pos);
+    } else {
+        return Err(ContractError::CreatorNotFound {});
+    }
+
+    CREATOR.save(deps.storage, &creators)?;
+
+    let res = Response::new()
+        .add_attribute("method", "remove_creator")
+        .add_attribute("creator", address.to_string());
+
+    Ok(res)
 }
 
 pub fn migrate_pair(
@@ -155,7 +219,7 @@ pub fn execute_create_pair(
     info: MessageInfo,
     asset_infos: [AssetInfo; 2],
     pair_admin: Option<String>,
-    operator: Option<String>,
+    _operator: Option<String>,
     provide_liquidity: Option<ProvideLiquidityParams>,
 ) -> Result<Response, ContractError> {
     let config: Config = CONFIG.load(deps.storage)?;
@@ -164,6 +228,28 @@ pub fn execute_create_pair(
         asset_infos[1].to_raw(deps.api)?,
     ];
 
+    let restricted_list = RESTRICTED_ASSETS
+        .may_load(deps.storage)?
+        .unwrap_or(RestrictedAssets { assets: Vec::new() });
+
+    let creators = CREATOR.may_load(deps.storage)?.unwrap_or(Creator {
+        whitelist_addresses: vec![],
+    });
+
+    for asset in asset_infos.as_ref().into_iter() {
+        if let AssetInfo::NativeToken { denom, .. } = &asset {
+            if denom.contains("factory/orai1") {
+                let parts: Vec<&str> = denom.split('/').collect();
+                if parts.len() > 2 && restricted_list.assets.contains(&parts[0..2].join("/")) {
+                    // permission check
+                    if !creators.whitelist_addresses.contains(&info.sender) {
+                        return Err(ContractError::Unauthorized {});
+                    }
+                }
+            }
+        }
+    }
+
     let pair_key = pair_key(&raw_infos);
 
     // can not update pair once updated
@@ -185,8 +271,6 @@ pub fn execute_create_pair(
     )?;
     let pair_admin = pair_admin.unwrap_or(env.contract.address.to_string());
 
-    let operator_addr = operator.map(|op| deps.api.addr_validate(&op)).transpose()?;
-
     // if provide_liquidity is not None, transfer all cw20 tokens to this contract
     let mut messages: Vec<CosmosMsg> = vec![];
 
@@ -348,6 +432,35 @@ pub fn execute_provide_liquidity(
         .add_message(provide_msg))
 }
 
+pub fn execute_restrict_asset(
+    deps: DepsMut,
+    info: MessageInfo,
+    prefix: String,
+) -> Result<Response, ContractError> {
+    let config = CONFIG.load(deps.storage)?;
+
+    // permission check
+    if deps.api.addr_canonicalize(info.sender.as_str())? != config.owner {
+        return Err(ContractError::Unauthorized {});
+    }
+
+    let mut restrict_list = RESTRICTED_ASSETS
+        .may_load(deps.storage)?
+        .unwrap_or(RestrictedAssets { assets: vec![] });
+    if restrict_list.assets.contains(&prefix) {
+        return Err(ContractError::RestrictPrefixExisted {});
+    }
+    restrict_list.assets.push(prefix.clone());
+
+    RESTRICTED_ASSETS.save(deps.storage, &restrict_list)?;
+
+    let res = Response::new()
+        .add_attribute("method", "restrict_asset")
+        .add_attribute("restrict_asset", prefix.to_string());
+
+    Ok(res)
+}
+
 /// This just stores the result for future query
 #[cfg_attr(not(feature = "library"), entry_point)]
 pub fn reply(deps: DepsMut, _env: Env, msg: Reply) -> Result<Response, ContractError> {
@@ -390,6 +503,8 @@ pub fn query(deps: Deps, _env: Env, msg: QueryMsg) -> StdResult<Binary> {
         QueryMsg::Pairs { start_after, limit } => {
             to_json_binary(&query_pairs(deps, start_after, limit)?)
         }
+        QueryMsg::RestrictedAssets {} => to_json_binary(&query_restricted_assets(deps)?),
+        QueryMsg::GetCreators {} => to_json_binary(&get_creators(deps)?),
     }
 }
 
@@ -437,6 +552,25 @@ pub fn query_pairs(
     Ok(resp)
 }
 
+pub fn query_restricted_assets(deps: Deps) -> StdResult<RestrictedAssetResponse> {
+    let restricted_list = RESTRICTED_ASSETS
+        .may_load(deps.storage)?
+        .unwrap_or(RestrictedAssets { assets: Vec::new() });
+
+    Ok(RestrictedAssetResponse {
+        prefixes: restricted_list.assets,
+    })
+}
+
+fn get_creators(deps: Deps) -> StdResult<CreatorsResponse> {
+    let creators = CREATOR.may_load(deps.storage)?.unwrap_or(Creator {
+        whitelist_addresses: vec![],
+    });
+    Ok(CreatorsResponse {
+        creators: creators.whitelist_addresses,
+    })
+}
+
 #[cfg_attr(not(feature = "library"), entry_point)]
 pub fn migrate(deps: DepsMut, _env: Env, msg: MigrateMsg) -> StdResult<Response> {
     let config = Config {

contracts/oraiswap_factory/src/state.rs

@@ -1,6 +1,6 @@
 use cosmwasm_schema::cw_serde;
 
-use cosmwasm_std::{Api, CanonicalAddr, Order, StdResult, Storage};
+use cosmwasm_std::{Addr, Api, CanonicalAddr, Order, StdResult, Storage};
 use cw_storage_plus::{Bound, Item, Map};
 use oraiswap::asset::{AssetInfoRaw, PairInfo, PairInfoRaw};
 
@@ -15,12 +15,26 @@ pub struct Config {
     pub operator: CanonicalAddr,
 }
 
+#[cw_serde]
+pub struct RestrictedAssets {
+    // token factory contract address (token-bindings), only store the prefix "factory/orai1"
+    pub assets: Vec<String>,
+}
+
+#[cw_serde]
+pub struct Creator {
+    pub whitelist_addresses: Vec<Addr>,
+}
+
 // put the length bytes at the first for compatibility with legacy singleton store
 pub const CONFIG: Item<Config> = Item::new("\u{0}\u{6}config");
 
 // store temporary pair info while waiting for deployment
 pub const PAIRS: Map<&[u8], PairInfoRaw> = Map::new("pairs");
 
+pub const RESTRICTED_ASSETS: Item<RestrictedAssets> = Item::new("restricted_assets");
+pub const CREATOR: Item<Creator> = Item::new("creator");
+
 // settings for pagination
 const MAX_LIMIT: u32 = 30;
 const DEFAULT_LIMIT: u32 = 10;

contracts/oraiswap_factory/src/testing.rs

@@ -1,11 +1,13 @@
-use cosmwasm_std::{to_json_binary, Addr};
+use std::str::FromStr;
+
+use cosmwasm_std::{coin, to_json_binary, Addr, Coin};
 use oraiswap::asset::{AssetInfo, PairInfo};
 
 use oraiswap::create_entry_points_testing;
 use oraiswap::factory::ConfigResponse;
 use oraiswap::pair::{PairResponse, DEFAULT_COMMISSION_RATE, DEFAULT_OPERATOR_FEE};
 use oraiswap::querier::query_pair_info_from_pair;
-use oraiswap::testing::MockApp;
+use oraiswap::testing::{MockApp, APP_OWNER};
 
 #[test]
 fn create_pair() {
@@ -82,6 +84,97 @@ fn create_pair() {
     );
 }
 
+#[test]
+fn create_pair_restricted() {
+    let denom_1 = "factory/orai1token/token1";
+    let denom_2 = "factory/orai1hehe/token2";
+
+    let init_balance: &[(&str, &[Coin])] = &[
+        (
+            APP_OWNER,
+            &[coin(1000000000000000000u128, "factory/orai1token1")],
+        ),
+        (
+            APP_OWNER,
+            &[coin(1000000000000000000u128, "factory/orai1token2")],
+        ),
+    ];
+    let mut app = MockApp::new(&init_balance);
+    app.set_token_contract(Box::new(create_entry_points_testing!(oraiswap_token)));
+    app.set_oracle_contract(Box::new(create_entry_points_testing!(oraiswap_oracle)));
+
+    app.set_factory_and_pair_contract(
+        Box::new(create_entry_points_testing!(crate).with_reply_empty(crate::contract::reply)),
+        Box::new(
+            create_entry_points_testing!(oraiswap_pair)
+                .with_reply_empty(oraiswap_pair::contract::reply),
+        ),
+    );
+
+    let asset_infos = [
+        AssetInfo::NativeToken {
+            denom: String::from_str(&denom_1).unwrap(),
+        },
+        AssetInfo::NativeToken {
+            denom: String::from_str(&denom_2).unwrap(),
+        },
+    ];
+
+    // restrict
+    app.restrict_asset("factory/orai1token".to_string())
+        .unwrap();
+    let restrict_prefix = app.query_restrict_denom().unwrap();
+    assert_eq!(restrict_prefix.prefixes.len(), 1);
+
+    // create pair failed
+    app.create_pair_by(asset_infos.clone(), "user1".to_string())
+        .unwrap_err();
+
+    // add creator
+    app.add_creator(APP_OWNER.to_string()).unwrap();
+    let creators = app.query_creators().unwrap();
+    assert_eq!(creators.creators.len(), 1);
+    assert_eq!(creators.creators[0].to_string(), APP_OWNER.to_string());
+
+    let contract_addr = app.create_pair(asset_infos.clone()).unwrap();
+
+    // query pair info
+    let pair_info =
+        query_pair_info_from_pair(&app.as_querier().into_empty(), contract_addr.clone()).unwrap();
+
+    // get config
+    let config: String = app
+        .as_querier()
+        .query_wasm_smart(
+            contract_addr.clone(),
+            &oraiswap::pair::QueryMsg::Operator {},
+        )
+        .unwrap();
+
+    let factory_config: ConfigResponse = app
+        .query(
+            app.factory_addr.clone(),
+            &oraiswap::factory::QueryMsg::Config {},
+        )
+        .unwrap();
+
+    assert_eq!(config, factory_config.operator);
+
+    // should never change commission rate once deployed
+    let pair_res = app.query_pair(asset_infos.clone()).unwrap();
+    assert_eq!(
+        pair_res,
+        PairInfo {
+            oracle_addr: app.oracle_addr,
+            liquidity_token: pair_info.liquidity_token,
+            contract_addr,
+            asset_infos,
+            commission_rate: DEFAULT_COMMISSION_RATE.into(),
+            operator_fee: DEFAULT_OPERATOR_FEE.to_string()
+        }
+    );
+}
+
 #[test]
 fn add_pair() {
     let mut app = MockApp::new(&[]);

packages/oraiswap/src/error.rs

@@ -113,4 +113,13 @@ pub enum ContractError {
 
     #[error("Contract paused")]
     Paused {},
+
+    #[error("Restricted prefix existed")]
+    RestrictPrefixExisted {},
+
+    #[error("Creator is whitelisted already")]
+    CreatorAlreadyExists {},
+
+    #[error("Not found this creator")]
+    CreatorNotFound {}
 }