add client echo limit

ref-exchange/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "ref-exchange"
-version = "1.9.4"
+version = "1.9.5"
 authors = ["Illia Polosukhin <[email protected]>"]
 edition = "2018"
 publish = false

ref-exchange/release_notes.md

@@ -1,5 +1,11 @@
 # Release Notes
 
+### Version 1.9.5
+```
+3hF1UzsT5mzxbJLMA8BD7gmCH1BL8cWjvFmSZYREpZXK
+```
+1. add client echo limit
+
 ### Version 1.9.4
 ```
 DBz69SAuDcvGWrEraoKjNEMiiDxE3PagejSfhYw3SfqH

ref-exchange/src/client_echo_limit.rs

@@ -0,0 +1,99 @@
+use crate::*;
+
+pub fn read_ce_tw_from_storage() -> UnorderedSet<AccountId> {
+    if let Some(content) = env::storage_read(CLIENT_ECHO_TOKEN_ID_WHITHELIST.as_bytes()) {
+        UnorderedSet::try_from_slice(&content).expect("deserialize client echo token id whitelist failed.")
+    } else {
+        UnorderedSet::new(StorageKey::ClientEchoTokenIdWhitelistItem)
+    }
+}
+
+pub fn write_ce_tw_to_storage(client_echo_token_id_whitelist: UnorderedSet<AccountId>) {
+    env::storage_write(
+        CLIENT_ECHO_TOKEN_ID_WHITHELIST.as_bytes(), 
+        &client_echo_token_id_whitelist.try_to_vec().unwrap(),
+    );
+}
+
+pub fn read_ce_sw_from_storage() -> UnorderedSet<AccountId> {
+    if let Some(content) = env::storage_read(CLIENT_ECHO_SENDER_ID_WHITHELIST.as_bytes()) {
+        UnorderedSet::try_from_slice(&content).expect("deserialize client echo sender id whitelist failed.")
+    } else {
+        UnorderedSet::new(StorageKey::ClientEchoSenderIdWhitelistItem)
+    }
+}
+
+pub fn write_ce_sw_to_storage(client_echo_sender_id_whitelist: UnorderedSet<AccountId>) {
+    env::storage_write(
+        CLIENT_ECHO_SENDER_ID_WHITHELIST.as_bytes(), 
+        &client_echo_sender_id_whitelist.try_to_vec().unwrap(),
+    );
+}
+
+pub fn assert_client_echo_valid(token_id: &AccountId, sender_id: &AccountId) {
+    let client_echo_token_id_whitelist = read_ce_tw_from_storage();
+    let client_echo_sender_id_whitelist = read_ce_sw_from_storage();
+    assert!(client_echo_token_id_whitelist.contains(token_id), "Invalid client echo token id");
+    assert!(client_echo_sender_id_whitelist.contains(sender_id), "Invalid client echo sender id");
+}
+
+#[near_bindgen]
+impl Contract {
+    #[payable]
+    pub fn extend_client_echo_token_id_whitelist(&mut self, token_ids: Vec<ValidAccountId>) {
+        assert_one_yocto();
+        assert!(self.is_owner_or_guardians(), "{}", ERR100_NOT_ALLOWED);
+        let mut client_echo_token_id_whitelist = read_ce_tw_from_storage();
+        for token_id in token_ids {
+            let is_success = client_echo_token_id_whitelist.insert(token_id.as_ref());
+            assert!(is_success, "Token id already exist");
+        }
+        write_ce_tw_to_storage(client_echo_token_id_whitelist);
+    }
+
+    #[payable]
+    pub fn remove_client_echo_token_id_whitelist(&mut self, token_ids: Vec<ValidAccountId>) {
+        assert_one_yocto();
+        self.assert_owner();
+        let mut client_echo_token_id_whitelist = read_ce_tw_from_storage();
+        for token_id in token_ids {
+            let is_success = client_echo_token_id_whitelist.remove(token_id.as_ref());
+            assert!(is_success, "Invalid token id");
+        }
+        write_ce_tw_to_storage(client_echo_token_id_whitelist);
+    }
+
+    pub fn get_client_echo_token_id_whitelist(&self) -> Vec<AccountId> {
+        let client_echo_token_id_whitelist = read_ce_tw_from_storage();
+        client_echo_token_id_whitelist.to_vec()
+    }
+
+    #[payable]
+    pub fn extend_client_echo_sender_id_whitelist(&mut self, sender_ids: Vec<ValidAccountId>) {
+        assert_one_yocto();
+        assert!(self.is_owner_or_guardians(), "{}", ERR100_NOT_ALLOWED);
+        let mut client_echo_sender_id_whitelist = read_ce_sw_from_storage();
+        for sender_id in sender_ids {
+            let is_success = client_echo_sender_id_whitelist.insert(sender_id.as_ref());
+            assert!(is_success, "Sender id already exist");
+        }
+        write_ce_sw_to_storage(client_echo_sender_id_whitelist);
+    }
+
+    #[payable]
+    pub fn remove_client_echo_sender_id_whitelist(&mut self, sender_ids: Vec<ValidAccountId>) {
+        assert_one_yocto();
+        self.assert_owner();
+        let mut client_echo_sender_id_whitelist = read_ce_sw_from_storage();
+        for sender_id in sender_ids {
+            let is_success = client_echo_sender_id_whitelist.remove(sender_id.as_ref());
+            assert!(is_success, "Invalid sender id");
+        }
+        write_ce_sw_to_storage(client_echo_sender_id_whitelist);
+    }
+
+    pub fn get_client_echo_sender_id_whitelist(&self) -> Vec<AccountId> {
+        let client_echo_sender_id_whitelist = read_ce_sw_from_storage();
+        client_echo_sender_id_whitelist.to_vec()
+    }
+}

ref-exchange/src/custom_keys.rs

@@ -6,4 +6,8 @@ pub const DEGEN_STORAGE_KEY: &str = "custom_degen_key";
 pub const DEGEN_ORACLE_CONFIG_STORAGE_KEY: &str = "custom_degen_oracle_config_key";
 
 // Key for pool limit
-pub const POOL_LIMIT: &str = "pl";
+pub const POOL_LIMIT: &str = "pl";
+
+// Key for client echo limit
+pub const CLIENT_ECHO_TOKEN_ID_WHITHELIST: &str = "ce_tw";
+pub const CLIENT_ECHO_SENDER_ID_WHITHELIST: &str = "ce_sw";

ref-exchange/src/lib.rs

@@ -34,6 +34,7 @@ pub use crate::unit_lpt_cumulative_infos::*;
 pub use crate::oracle::*;
 pub use crate::degen_swap::*;
 pub use crate::pool_limit_info::*;
+pub use crate::client_echo_limit::*;
 
 mod account_deposit;
 mod action;
@@ -56,6 +57,7 @@ mod custom_keys;
 mod shadow_actions;
 mod unit_lpt_cumulative_infos;
 mod pool_limit_info;
+mod client_echo_limit;
 
 near_sdk::setup_alloc!();
 
@@ -72,6 +74,8 @@ pub(crate) enum StorageKey {
     ShadowRecord {account_id: AccountId},
     UnitShareCumulativeInfo,
     PoolLimit,
+    ClientEchoTokenIdWhitelistItem,
+    ClientEchoSenderIdWhitelistItem,
 }
 
 #[derive(BorshDeserialize, BorshSerialize, Serialize, Deserialize, Eq, PartialEq, Clone)]

ref-exchange/src/token_receiver.rs

@@ -113,6 +113,9 @@ impl FungibleTokenReceiver for Contract {
                 } => {
                     assert!(!(swap_out_recipient.is_some() && client_echo.is_some()), "client_echo and swap_out_recipient cannot have value at the same time");
                     assert_ne!(actions.len(), 0, "{}", ERR72_AT_LEAST_ONE_SWAP);
+                    if client_echo.is_some() {
+                        assert_client_echo_valid(&token_in, sender_id.as_ref());
+                    }
                     let referral_id = referral_id.map(|x| x.to_string());
                     let out_amounts = self.internal_direct_actions(
                         token_in,

ref-exchange/tests/test_client_echo.rs

@@ -0,0 +1,194 @@
+use crate::common::utils::*;
+pub mod common;
+
+use test_token::ContractContract as TestToken;
+use mock_boost_farming::{ContractContract as MockBoostFarming};
+
+use near_sdk::{json_types::U128, serde_json::Value, AccountId};
+use near_sdk_sim::{call, deploy, view, to_yocto, ContractAccount, ExecutionResult, UserAccount};
+
+near_sdk_sim::lazy_static_include::lazy_static_include_bytes! {
+    MOCK_BOOST_FARMING_WASM_BYTES => "../res/mock_boost_farming.wasm",
+}
+
+fn do_swap(
+    user: &UserAccount,
+    contract: &ContractAccount<TestToken>,
+    actions: Vec<String>,
+    amount: u128,
+    client_echo: Option<String>,
+    swap_out_recipient: Option<AccountId>,
+) -> ExecutionResult {
+    let client_echo = if let Some(client_echo) = client_echo {
+        format!(",\"client_echo\":\"{}\"", client_echo)
+    } else {
+        "".to_string()
+    };
+    let swap_out_recipient = if let Some(swap_out_recipient) = swap_out_recipient {
+        format!(",\"swap_out_recipient\":\"{}\"", swap_out_recipient)
+    } else {
+        "".to_string()
+    };
+    let actions_str = actions.join(", ");
+    let msg_str = format!(
+        "{{\"actions\": [{}]{}{}}}",
+        actions_str, client_echo, swap_out_recipient
+    );
+    call!(
+        user,
+        contract.ft_transfer_call(to_va(swap()), amount.into(), None, msg_str),
+        deposit = 1
+    )
+}
+
+fn pack_action(
+    pool_id: u32,
+    token_in: &str,
+    token_out: &str,
+    amount_in: Option<u128>,
+    min_amount_out: u128,
+) -> String {
+    if let Some(amount_in) = amount_in {
+        format!(
+            "{{\"pool_id\": {}, \"token_in\": \"{}\", \"amount_in\": \"{}\", \"token_out\": \"{}\", \"min_amount_out\": \"{}\"}}",
+            pool_id, token_in, amount_in, token_out, min_amount_out
+        )
+    } else {
+        format!(
+            "{{\"pool_id\": {}, \"token_in\": \"{}\", \"token_out\": \"{}\", \"min_amount_out\": \"{}\"}}",
+            pool_id, token_in, token_out, min_amount_out
+        )
+    }
+}
+
+fn boost_farming() -> AccountId {
+    "boost_farming".to_string()
+}
+
+#[test]
+fn test_client_echo() {
+    let (root, owner, pool, token1, token2, _) = setup_pool_with_liquidity();
+    let new_user = root.create_user("new_user".to_string(), to_yocto("100"));
+    let mock_boost_farming = deploy!(
+        contract: MockBoostFarming,
+        contract_id: boost_farming(),
+        bytes: &MOCK_BOOST_FARMING_WASM_BYTES,
+        signer_account: root,
+        init_method: new(root.account_id())
+    );
+    let outcome = call!(
+        root,
+        mock_boost_farming.create_seed(token2.account_id(),24, Some(U128(0)), Some(0)),
+        deposit = 1
+    );
+    outcome.assert_success();
+    call!(
+        pool.user_account,
+        mock_boost_farming.storage_deposit(None, None),
+        deposit = to_yocto("1")
+    )
+    .assert_success();
+
+    call!(
+        mock_boost_farming.user_account,
+        token1.mint(to_va(mock_boost_farming.user_account.account_id.clone()), U128(to_yocto("10")))
+    )
+    .assert_success();
+    call!(
+        mock_boost_farming.user_account,
+        token2.storage_deposit(None, None),
+        deposit = to_yocto("1")
+    )
+    .assert_success();
+    call!(
+        new_user,
+        token2.storage_deposit(None, None),
+        deposit = to_yocto("1")
+    )
+    .assert_success();
+    assert_eq!(balance_of(&token1, &mock_boost_farming.user_account.account_id), to_yocto("10"));
+
+    let action = pack_action(0, &token1.account_id(), &token2.account_id(), None, 0);
+    let out_come = do_swap(
+        &mock_boost_farming.user_account,
+        &token1,
+        vec![action.clone()],
+        to_yocto("1"),
+        Some("Hi".to_string()),
+        Some(mock_boost_farming.user_account.account_id.to_string()),
+    );
+    out_come.assert_success();
+    assert_eq!(get_error_count(&out_come), 1);
+    assert!(get_error_status(&out_come)
+        .contains("client_echo and swap_out_recipient cannot have value at the same time"));
+
+    let out_come = do_swap(
+        &mock_boost_farming.user_account,
+        &token1,
+        vec![action.clone()],
+        to_yocto("1"),
+        Some("Hi".to_string()),
+        None,
+    );
+    out_come.assert_success();
+    assert_eq!(get_error_count(&out_come), 1);
+    assert!(get_error_status(&out_come).contains("Invalid client echo token id"));
+
+    call!(
+        owner,
+        pool.extend_client_echo_token_id_whitelist(vec![to_va(token1.account_id())]),
+        deposit = 1
+    )
+    .assert_success();
+
+    let out_come = do_swap(
+        &mock_boost_farming.user_account,
+        &token1,
+        vec![action.clone()],
+        to_yocto("1"),
+        Some("Hi".to_string()),
+        None,
+    );
+    out_come.assert_success();
+    assert_eq!(get_error_count(&out_come), 1);
+    assert!(get_error_status(&out_come).contains("Invalid client echo sender id"));
+
+    assert_eq!(balance_of(&token1, &mock_boost_farming.user_account.account_id), to_yocto("10"));
+    assert_eq!(balance_of(&token2, &mock_boost_farming.user_account.account_id), to_yocto("0"));
+
+    call!(
+        owner,
+        pool.extend_client_echo_sender_id_whitelist(vec![to_va(mock_boost_farming.account_id())]),
+        deposit = 1
+    )
+    .assert_success();
+
+    let out_come = do_swap(
+        &mock_boost_farming.user_account,
+        &token1,
+        vec![action.clone()],
+        to_yocto("1"),
+        Some("\\\"Free\\\"".to_string()),
+        None,
+    );
+    out_come.assert_success();
+
+    println!(" {:?}", view!(mock_boost_farming.get_seed(token2.account_id())).unwrap_json::<Value>());
+
+    assert_eq!(balance_of(&token1, &mock_boost_farming.user_account.account_id), to_yocto("9"));
+    assert_eq!(balance_of(&token2, &mock_boost_farming.user_account.account_id), 1814048647419868151852693);
+
+    let out_come = do_swap(
+        &mock_boost_farming.user_account,
+        &token1,
+        vec![action.clone()],
+        to_yocto("1"),
+        None,
+        Some(new_user.account_id.clone()),
+    );
+    out_come.assert_success();
+
+    assert_eq!(balance_of(&token1, &mock_boost_farming.user_account.account_id), to_yocto("8"));
+    assert_eq!(balance_of(&token2, &mock_boost_farming.user_account.account_id), 1814048647419868151852693);
+    assert_eq!(balance_of(&token2, &new_user.account_id), 1512022210810475642302724);
+}

ref-exchange/tests/test_migrate.rs

@@ -51,7 +51,7 @@ fn test_upgrade() {
     .assert_success();
     let metadata = get_metadata(&pool);
     // println!("{:#?}", metadata);
-    assert_eq!(metadata.version, "1.9.4".to_string());
+    assert_eq!(metadata.version, "1.9.5".to_string());
     assert_eq!(metadata.admin_fee_bps, 5);
     assert_eq!(metadata.boost_farm_id, "boost_farm".to_string());
     assert_eq!(metadata.burrowland_id, "burrowland".to_string());