Merge pull request #873 from multiversx/router-further-audit-improvem…

…ents Router further audit improvements
multiversx · Apr 8, 2024 · b8092a3 · b8092a3
2 parents 7787734 + b542ecd
commit b8092a3
Show file tree

Hide file tree

Showing 8 changed files with 117 additions and 80 deletions.
diff --git a/dex/router/src/config.rs b/dex/router/src/config.rs
@@ -0,0 +1,69 @@
+multiversx_sc::imports!();
+multiversx_sc::derive_imports!();
+
+use crate::{enable_swap_by_user::EnableSwapByUserConfig, factory::PairTokens};
+
+#[multiversx_sc::module]
+pub trait ConfigModule {
+    fn is_active(&self) -> bool {
+        self.state().get()
+    }
+
+    fn check_is_pair_sc(&self, pair_address: &ManagedAddress) {
+        require!(
+            self.address_pair_map().contains_key(pair_address),
+            "Not a pair SC"
+        );
+    }
+
+    #[view(getPairCreationEnabled)]
+    #[storage_mapper("pair_creation_enabled")]
+    fn pair_creation_enabled(&self) -> SingleValueMapper<bool>;
+
+    #[view(getState)]
+    #[storage_mapper("state")]
+    fn state(&self) -> SingleValueMapper<bool>;
+
+    #[view(getOwner)]
+    #[storage_mapper("owner")]
+    fn owner(&self) -> SingleValueMapper<ManagedAddress>;
+
+    #[only_owner]
+    #[endpoint(setTemporaryOwnerPeriod)]
+    fn set_temporary_owner_period(&self, period_blocks: u64) {
+        self.temporary_owner_period().set(period_blocks);
+    }
+
+    #[only_owner]
+    #[endpoint(setPairTemplateAddress)]
+    fn set_pair_template_address(&self, address: ManagedAddress) {
+        self.pair_template_address().set(&address);
+    }
+
+    #[storage_mapper("pair_map")]
+    fn pair_map(&self) -> MapMapper<PairTokens<Self::Api>, ManagedAddress>;
+
+    #[storage_mapper("address_pair_map")]
+    fn address_pair_map(&self) -> MapMapper<ManagedAddress, PairTokens<Self::Api>>;
+
+    #[view(getPairTemplateAddress)]
+    #[storage_mapper("pair_template_address")]
+    fn pair_template_address(&self) -> SingleValueMapper<ManagedAddress>;
+
+    #[view(getTemporaryOwnerPeriod)]
+    #[storage_mapper("temporary_owner_period")]
+    fn temporary_owner_period(&self) -> SingleValueMapper<u64>;
+
+    #[storage_mapper("pair_temporary_owner")]
+    fn pair_temporary_owner(&self) -> MapMapper<ManagedAddress, (ManagedAddress, u64)>;
+
+    #[storage_mapper("enableSwapByUserConfig")]
+    fn enable_swap_by_user_config(
+        &self,
+        token_id: &TokenIdentifier,
+    ) -> SingleValueMapper<EnableSwapByUserConfig<Self::Api>>;
+
+    #[view(getCommonTokensForUserPairs)]
+    #[storage_mapper("commonTokensForUserPairs")]
+    fn common_tokens_for_user_pairs(&self) -> UnorderedSetMapper<TokenIdentifier>;
+}
diff --git a/dex/router/src/contract.rs b/dex/router/src/contract.rs
@@ -3,6 +3,7 @@
 multiversx_sc::imports!();
 multiversx_sc::derive_imports!();
 
+pub mod config;
 pub mod enable_swap_by_user;
 mod events;
 pub mod factory;
@@ -24,7 +25,8 @@ const USER_DEFINED_TOTAL_FEE_PERCENT: u64 = 1_000;
 
 #[multiversx_sc::contract]
 pub trait Router:
-    factory::FactoryModule
+    config::ConfigModule
+    + factory::FactoryModule
     + events::EventsModule
     + multi_pair_swap::MultiPairSwap
     + token_send::TokenSendModule
@@ -40,7 +42,9 @@ pub trait Router:
     }
 
     #[endpoint]
-    fn upgrade(&self) {}
+    fn upgrade(&self) {
+        self.state().set(false);
+    }
 
     #[only_owner]
     #[endpoint]
@@ -60,6 +64,10 @@ pub trait Router:
     #[endpoint]
     fn resume(&self, address: ManagedAddress) {
         if address == self.blockchain().get_sc_address() {
+            require!(
+                self.pair_map().len() == self.address_pair_map().len(),
+                "The size of the 2 pair maps is not the same"
+            );
             self.state().set(true);
         } else {
             self.check_is_pair_sc(&address);
@@ -372,11 +380,6 @@ pub trait Router:
         }
     }
 
-    #[inline]
-    fn is_active(&self) -> bool {
-        self.state().get()
-    }
-
     #[only_owner]
     #[endpoint(setPairCreationEnabled)]
     fn set_pair_creation_enabled(&self, enabled: bool) {
@@ -395,17 +398,10 @@ pub trait Router:
         for (pair_tokens, address) in pair_map.iter() {
             address_pair_map.insert(address, pair_tokens);
         }
-    }
 
-    #[view(getPairCreationEnabled)]
-    #[storage_mapper("pair_creation_enabled")]
-    fn pair_creation_enabled(&self) -> SingleValueMapper<bool>;
-
-    #[view(getState)]
-    #[storage_mapper("state")]
-    fn state(&self) -> SingleValueMapper<bool>;
-
-    #[view(getOwner)]
-    #[storage_mapper("owner")]
-    fn owner(&self) -> SingleValueMapper<ManagedAddress>;
+        require!(
+            pair_map.len() == address_pair_map.len(),
+            "The size of the 2 pair maps is not the same"
+        );
+    }
 }
diff --git a/dex/router/src/enable_swap_by_user.rs b/dex/router/src/enable_swap_by_user.rs
@@ -5,7 +5,7 @@ use pair::{config::ProxyTrait as _, pair_actions::views::ProxyTrait as _};
 use pausable::{ProxyTrait as _, State};
 use simple_lock::locked_token::LockedTokenAttributes;
 
-use crate::{DEFAULT_SPECIAL_FEE_PERCENT, USER_DEFINED_TOTAL_FEE_PERCENT};
+use crate::{config, DEFAULT_SPECIAL_FEE_PERCENT, USER_DEFINED_TOTAL_FEE_PERCENT};
 
 static PAIR_LP_TOKEN_ID_STORAGE_KEY: &[u8] = b"lpTokenIdentifier";
 static PAIR_INITIAL_LIQ_ADDER_STORAGE_KEY: &[u8] = b"initial_liquidity_adder";
@@ -27,7 +27,7 @@ pub struct SafePriceResult<M: ManagedTypeApi> {
 
 #[multiversx_sc::module]
 pub trait EnableSwapByUserModule:
-    crate::factory::FactoryModule + crate::events::EventsModule
+    config::ConfigModule + crate::factory::FactoryModule + crate::events::EventsModule
 {
     #[only_owner]
     #[endpoint(configEnableByUserParameters)]
@@ -250,14 +250,4 @@ pub trait EnableSwapByUserModule:
 
     #[proxy]
     fn user_pair_proxy(&self, to: ManagedAddress) -> pair::Proxy<Self::Api>;
-
-    #[storage_mapper("enableSwapByUserConfig")]
-    fn enable_swap_by_user_config(
-        &self,
-        token_id: &TokenIdentifier,
-    ) -> SingleValueMapper<EnableSwapByUserConfig<Self::Api>>;
-
-    #[view(getCommonTokensForUserPairs)]
-    #[storage_mapper("commonTokensForUserPairs")]
-    fn common_tokens_for_user_pairs(&self) -> UnorderedSetMapper<TokenIdentifier>;
 }
diff --git a/dex/router/src/factory.rs b/dex/router/src/factory.rs
@@ -1,6 +1,8 @@
 multiversx_sc::imports!();
 multiversx_sc::derive_imports!();
 
+use crate::config;
+
 const TEMPORARY_OWNER_PERIOD_BLOCKS: u64 = 50;
 
 #[derive(TopEncode, TopDecode, NestedEncode, NestedDecode, PartialEq, TypeAbi)]
@@ -17,7 +19,7 @@ pub struct PairContractMetadata<M: ManagedTypeApi> {
 }
 
 #[multiversx_sc::module]
-pub trait FactoryModule {
+pub trait FactoryModule: config::ConfigModule {
     #[proxy]
     fn pair_contract_deploy_proxy(&self) -> pair::Proxy<Self::Api>;
 
@@ -172,11 +174,11 @@ pub trait FactoryModule {
         address
     }
 
-    fn check_is_pair_sc(&self, pair_address: &ManagedAddress) {
-        require!(
-            self.address_pair_map().contains_key(pair_address),
-            "Not a pair SC"
-        );
+    #[view(getPairTokens)]
+    fn get_pair_tokens(&self, pair_address: ManagedAddress) -> PairTokens<Self::Api> {
+        let pair_tokens_opt = self.address_pair_map().get(&pair_address);
+        require!(pair_tokens_opt.is_some(), "Pair address not found");
+        pair_tokens_opt.unwrap()
     }
 
     fn get_pair_temporary_owner(&self, pair_address: &ManagedAddress) -> Option<ManagedAddress> {
@@ -204,33 +206,4 @@ pub trait FactoryModule {
         self.pair_temporary_owner().clear();
         size
     }
-
-    #[only_owner]
-    #[endpoint(setTemporaryOwnerPeriod)]
-    fn set_temporary_owner_period(&self, period_blocks: u64) {
-        self.temporary_owner_period().set(period_blocks);
-    }
-
-    #[only_owner]
-    #[endpoint(setPairTemplateAddress)]
-    fn set_pair_template_address(&self, address: ManagedAddress) {
-        self.pair_template_address().set(&address);
-    }
-
-    #[storage_mapper("pair_map")]
-    fn pair_map(&self) -> MapMapper<PairTokens<Self::Api>, ManagedAddress>;
-
-    #[storage_mapper("address_pair_map")]
-    fn address_pair_map(&self) -> MapMapper<ManagedAddress, PairTokens<Self::Api>>;
-
-    #[view(getPairTemplateAddress)]
-    #[storage_mapper("pair_template_address")]
-    fn pair_template_address(&self) -> SingleValueMapper<ManagedAddress>;
-
-    #[view(getTemporaryOwnerPeriod)]
-    #[storage_mapper("temporary_owner_period")]
-    fn temporary_owner_period(&self) -> SingleValueMapper<u64>;
-
-    #[storage_mapper("pair_temporary_owner")]
-    fn pair_temporary_owner(&self) -> MapMapper<ManagedAddress, (ManagedAddress, u64)>;
 }
diff --git a/dex/router/src/multi_pair_swap.rs b/dex/router/src/multi_pair_swap.rs
@@ -3,6 +3,8 @@ multiversx_sc::derive_imports!();
 
 use pair::pair_actions::swap::ProxyTrait as _;
 
+use crate::config;
+
 use super::factory;
 
 type SwapOperationType<M> =
@@ -12,10 +14,14 @@ pub const SWAP_TOKENS_FIXED_INPUT_FUNC_NAME: &[u8] = b"swapTokensFixedInput";
 pub const SWAP_TOKENS_FIXED_OUTPUT_FUNC_NAME: &[u8] = b"swapTokensFixedOutput";
 
 #[multiversx_sc::module]
-pub trait MultiPairSwap: factory::FactoryModule + token_send::TokenSendModule {
+pub trait MultiPairSwap:
+    config::ConfigModule + factory::FactoryModule + token_send::TokenSendModule
+{
     #[payable("*")]
     #[endpoint(multiPairSwap)]
     fn multi_pair_swap(&self, swap_operations: MultiValueEncoded<SwapOperationType<Self::Api>>) -> ManagedVec<EsdtTokenPayment> {
+        require!(self.is_active(), "Not active");
+
         let (token_id, nonce, amount) = self.call_value().single_esdt().into_tuple();
         require!(nonce == 0, "Invalid nonce. Should be zero");
         require!(amount > 0u64, "Invalid amount. Should not be zero");

diff --git a/dex/router/tests/router_setup/mod.rs b/dex/router/tests/router_setup/mod.rs
@@ -27,10 +27,11 @@ pub const MIN_LOCKED_PERIOD_EPOCHS: u64 = 100;
 pub const USER_CUSTOM_TOKEN_BALANCE: u64 = 1_000_000_000;
 pub const USER_USDC_BALANCE: u64 = 1_000_000;
 
-use pair::config::*;
+use pair::config::ConfigModule as PairConfigModule;
 use pair::pair_actions::add_liq::AddLiquidityModule;
 use pair::*;
 use pausable::{PausableModule, State};
+use router::config::ConfigModule;
 use router::factory::*;
 use router::multi_pair_swap::*;
 use router::*;

diff --git a/dex/router/tests/router_test.rs b/dex/router/tests/router_test.rs
@@ -9,13 +9,14 @@ use multiversx_sc::{
         MultiValueEncoded,
     },
 };
-use pair::{config::ConfigModule, pair_actions::initial_liq::InitialLiquidityModule, Pair};
+use pair::{
+    config::ConfigModule as PairConfigModule, pair_actions::initial_liq::InitialLiquidityModule,
+    Pair,
+};
 use pausable::{PausableModule, State};
 use router::{
-    enable_swap_by_user::EnableSwapByUserModule,
-    factory::{FactoryModule, PairTokens},
-    multi_pair_swap::SWAP_TOKENS_FIXED_INPUT_FUNC_NAME,
-    Router,
+    config::ConfigModule, enable_swap_by_user::EnableSwapByUserModule, factory::PairTokens,
+    multi_pair_swap::SWAP_TOKENS_FIXED_INPUT_FUNC_NAME, Router,
 };
 use router_setup::*;
 

diff --git a/dex/router/wasm/src/lib.rs b/dex/router/wasm/src/lib.rs
@@ -5,9 +5,9 @@
 ////////////////////////////////////////////////////
 
 // Init:                                 1
-// Endpoints:                           31
+// Endpoints:                           33
 // Async Callback:                       1
-// Total number of exported functions:  33
+// Total number of exported functions:  35
 
 #![no_std]
 #![allow(internal_features)]
@@ -35,22 +35,23 @@ multiversx_sc_wasm_adapter::endpoints! {
         getPairCreationEnabled => pair_creation_enabled
         getState => state
         getOwner => owner
+        setTemporaryOwnerPeriod => set_temporary_owner_period
+        setPairTemplateAddress => set_pair_template_address
+        getPairTemplateAddress => pair_template_address
+        getTemporaryOwnerPeriod => temporary_owner_period
+        getCommonTokensForUserPairs => common_tokens_for_user_pairs
         getAllPairsManagedAddresses => get_all_pairs_addresses
         getAllPairTokens => get_all_token_pairs
         getAllPairContractMetadata => get_all_pair_contract_metadata
         getPair => get_pair
+        getPairTokens => get_pair_tokens
         clearPairTemporaryOwnerStorage => clear_pair_temporary_owner_storage
-        setTemporaryOwnerPeriod => set_temporary_owner_period
-        setPairTemplateAddress => set_pair_template_address
-        getPairTemplateAddress => pair_template_address
-        getTemporaryOwnerPeriod => temporary_owner_period
         multiPairSwap => multi_pair_swap
         configEnableByUserParameters => config_enable_by_user_parameters
         addCommonTokensForUserPairs => add_common_tokens_for_user_pairs
         removeCommonTokensForUserPairs => remove_common_tokens_for_user_pairs
         setSwapEnabledByUser => set_swap_enabled_by_user
         getEnableSwapByUserConfig => try_get_config
-        getCommonTokensForUserPairs => common_tokens_for_user_pairs
     )
 }