From 68748a066f87364c5b6b4dd20784cb7714618c54 Mon Sep 17 00:00:00 2001 From: Claudiu Lataretu Date: Fri, 9 Jun 2023 16:27:29 +0300 Subject: [PATCH] add back locked_token_id field from EnableSwapByUserConfig - check payment token id against locked_token_id in setSwapEnabledByUser - add unit test fro locked_token_id check against forged payment Signed-off-by: Claudiu Lataretu --- dex/router/src/enable_swap_by_user.rs | 20 ++- dex/router/tests/router_test.rs | 179 +++++++++++++++++++++++++- 2 files changed, 194 insertions(+), 5 deletions(-) diff --git a/dex/router/src/enable_swap_by_user.rs b/dex/router/src/enable_swap_by_user.rs index ba2ddb280..eec3f4d02 100644 --- a/dex/router/src/enable_swap_by_user.rs +++ b/dex/router/src/enable_swap_by_user.rs @@ -13,6 +13,7 @@ static PAIR_STATE_STORAGE_KEY: &[u8] = b"state"; #[derive(TypeAbi, TopEncode, TopDecode)] pub struct EnableSwapByUserConfig { + pub locked_token_id: TokenIdentifier, pub min_locked_token_value: BigUint, pub min_lock_period_epochs: u64, } @@ -33,22 +34,31 @@ pub trait EnableSwapByUserModule: fn config_enable_by_user_parameters( &self, common_token_id: TokenIdentifier, + locked_token_id: TokenIdentifier, min_locked_token_value: BigUint, min_lock_period_epochs: u64, - common_tokens_for_user_pairs: MultiValueEncoded, ) { require!( common_token_id.is_valid_esdt_identifier(), "Invalid locked token ID" ); + require!( + locked_token_id.is_valid_esdt_identifier(), + "Invalid locked token ID" + ); + + let whitelist = self.common_tokens_for_user_pairs(); + require!( + whitelist.contains(&common_token_id), + "Common token not whitelisted" + ); self.enable_swap_by_user_config(&common_token_id) .set(&EnableSwapByUserConfig { + locked_token_id, min_locked_token_value, min_lock_period_epochs, }); - - self.add_common_tokens_for_user_pairs(common_tokens_for_user_pairs); } #[only_owner] @@ -97,6 +107,10 @@ pub trait EnableSwapByUserModule: let lp_token_safe_price_result = self.get_lp_token_value(pair_address.clone(), locked_lp_token_amount); let config = self.try_get_config(&lp_token_safe_price_result.common_token_id); + require!( + payment.token_identifier == config.locked_token_id, + "Invalid locked token" + ); require!( lp_token_safe_price_result.safe_price_in_common_token >= config.min_locked_token_value, "Not enough value locked" diff --git a/dex/router/tests/router_test.rs b/dex/router/tests/router_test.rs index 7f6bbdeb9..5c3e50812 100644 --- a/dex/router/tests/router_test.rs +++ b/dex/router/tests/router_test.rs @@ -2,7 +2,9 @@ mod router_setup; use multiversx_sc::{ codec::multi_types::OptionalValue, storage::mappers::StorageTokenWrapper, - types::{EsdtLocalRole, ManagedAddress, ManagedVec, MultiValueEncoded}, + types::{ + EgldOrEsdtTokenIdentifier, EsdtLocalRole, ManagedAddress, ManagedVec, MultiValueEncoded, + }, }; use pair::{config::ConfigModule, Pair}; use pausable::{PausableModule, State}; @@ -151,11 +153,15 @@ fn user_enable_pair_swaps_through_router_test() { managed_address!(pair_wrapper.address_ref()), ); + sc.add_common_tokens_for_user_pairs(MultiValueEncoded::from(ManagedVec::from(vec![ + managed_token_id!(USDC_TOKEN_ID), + ]))); + sc.config_enable_by_user_parameters( managed_token_id!(USDC_TOKEN_ID), + managed_token_id!(LOCKED_TOKEN_ID), managed_biguint!(MIN_LOCKED_TOKEN_VALUE), MIN_LOCKED_PERIOD_EPOCHS, - ManagedVec::from_single_item(managed_token_id!(USDC_TOKEN_ID)).into(), ) }) .assert_ok(); @@ -288,3 +294,172 @@ fn user_enable_pair_swaps_through_router_test() { }), ); } + +#[test] +fn user_enable_pair_swaps_fail_test() { + let rust_zero = rust_biguint!(0u64); + let mut b_mock = BlockchainStateWrapper::new(); + let owner = b_mock.create_user_account(&rust_zero); + let user = b_mock.create_user_account(&rust_zero); + + let current_epoch = 5; + b_mock.set_block_epoch(current_epoch); + + b_mock.set_esdt_balance( + &user, + CUSTOM_TOKEN_ID, + &rust_biguint!(USER_CUSTOM_TOKEN_BALANCE), + ); + b_mock.set_esdt_balance(&user, USDC_TOKEN_ID, &rust_biguint!(USER_USDC_BALANCE)); + + let router_wrapper = b_mock.create_sc_account( + &rust_zero, + Some(&owner), + router::contract_obj, + ROUTER_WASM_PATH, + ); + let pair_wrapper = b_mock.create_sc_account( + &rust_zero, + Some(router_wrapper.address_ref()), + pair::contract_obj, + PAIR_WASM_PATH, + ); + + // setup router + b_mock + .execute_tx(&owner, &router_wrapper, &rust_zero, |sc| { + sc.init(OptionalValue::None); + + sc.pair_map().insert( + PairTokens { + first_token_id: managed_token_id!(CUSTOM_TOKEN_ID), + second_token_id: managed_token_id!(USDC_TOKEN_ID), + }, + managed_address!(pair_wrapper.address_ref()), + ); + + sc.add_common_tokens_for_user_pairs(MultiValueEncoded::from(ManagedVec::from(vec![ + managed_token_id!(USDC_TOKEN_ID), + ]))); + + sc.config_enable_by_user_parameters( + managed_token_id!(USDC_TOKEN_ID), + managed_token_id!(LOCKED_TOKEN_ID), + managed_biguint!(MIN_LOCKED_TOKEN_VALUE), + MIN_LOCKED_PERIOD_EPOCHS, + ) + }) + .assert_ok(); + + // setup pair + b_mock + .execute_tx(&owner, &pair_wrapper, &rust_zero, |sc| { + let first_token_id = managed_token_id!(CUSTOM_TOKEN_ID); + let second_token_id = managed_token_id!(USDC_TOKEN_ID); + let router_address = managed_address!(router_wrapper.address_ref()); + let router_owner_address = managed_address!(&owner); + + sc.init( + first_token_id, + second_token_id, + router_address, + router_owner_address, + 0, + 0, + managed_address!(&user), + MultiValueEncoded::>::new(), + ); + + assert_eq!(sc.state().get(), State::Inactive); + + sc.lp_token_identifier() + .set(&managed_token_id!(LPUSDC_TOKEN_ID)); + }) + .assert_ok(); + + b_mock.set_esdt_local_roles( + pair_wrapper.address_ref(), + LPUSDC_TOKEN_ID, + &[EsdtLocalRole::Mint, EsdtLocalRole::Burn], + ); + + // add liquidity + let payments = vec![ + TxTokenTransfer { + token_identifier: CUSTOM_TOKEN_ID.to_vec(), + nonce: 0, + value: rust_biguint!(USER_CUSTOM_TOKEN_BALANCE), + }, + TxTokenTransfer { + token_identifier: USDC_TOKEN_ID.to_vec(), + nonce: 0, + value: rust_biguint!(USER_USDC_BALANCE), + }, + ]; + + let user_lp_tokens_balance = 999_000u64; + b_mock + .execute_esdt_multi_transfer(&user, &pair_wrapper, &payments, |sc| { + let (lp_tokens_received, _, _) = sc.add_initial_liquidity().into_tuple(); + assert_eq!( + lp_tokens_received.token_identifier, + managed_token_id!(LPUSDC_TOKEN_ID) + ); + assert_eq!( + lp_tokens_received.amount, + managed_biguint!(user_lp_tokens_balance) + ); + }) + .assert_ok(); + + let custom_locked_token = b"LTOK2-123456"; + let _ = DebugApi::dummy(); + b_mock.set_nft_balance( + &user, + custom_locked_token, + 1, + &rust_biguint!(user_lp_tokens_balance), + &LockedTokenAttributes:: { + original_token_id: EgldOrEsdtTokenIdentifier::esdt(managed_token_id!(LPUSDC_TOKEN_ID)), + original_token_nonce: 0, + unlock_epoch: current_epoch + MIN_LOCKED_PERIOD_EPOCHS, + }, + ); + + // pass blocks time to update safe price + b_mock.set_block_nonce(1_000_000); + + // activate swaps through router + b_mock + .execute_esdt_transfer( + &user, + &router_wrapper, + custom_locked_token, + 1, + &rust_biguint!(user_lp_tokens_balance), + |sc| { + sc.set_swap_enabled_by_user(managed_address!(pair_wrapper.address_ref())); + }, + ) + .assert_user_error("Invalid locked token"); + + // check pair state is active + b_mock + .execute_query(&pair_wrapper, |sc| { + assert_eq!(sc.state().get(), State::PartialActive); + }) + .assert_ok(); + + // check user received the locked tokens back + b_mock.check_nft_balance( + &user, + custom_locked_token, + 1, + &rust_biguint!(user_lp_tokens_balance), + Some(&LockedTokenAttributes:: { + original_token_id: managed_token_id_wrapped!(LPUSDC_TOKEN_ID), + original_token_nonce: 0, + unlock_epoch: current_epoch + MIN_LOCKED_PERIOD_EPOCHS, + }), + ); +}