Merge pull request #5 from haiko-xyz/audit/nm-audit-fix

audit: NM audit fixes +`ISolverHooks` refactor

docs/1-technical-architecture.md

@@ -144,11 +144,11 @@ fn deposit(
 ) -> (u256, u256, u256);
 ```
 
-Withdrawals can be made either by calling `withdraw_at_ratio()` to withdraw from a public vault, or `withdraw()` to withdraw an arbitrary amounts from a private vault (available to the vault owner only).
+Withdrawals can be made either by calling `withdraw_public()` to withdraw from a public vault, or `withdraw_private()` to withdraw an arbitrary amounts from a private vault (available to the vault owner only).
 
 ```rust
 // Burn pool shares and withdraw funds from market.
-// Called for public vaults. For private vaults, use `withdraw_amount`.
+// Called for public vaults. For private vaults, use `withdraw_private`.
 //
 // # Arguments
 // * `market_id` - market id
@@ -157,12 +157,12 @@ Withdrawals can be made either by calling `withdraw_at_ratio()` to withdraw from
 // # Returns
 // * `base_amount` - base asset withdrawn
 // * `quote_amount` - quote asset withdrawn
-fn withdraw_at_ratio(
+fn withdraw_public(
     ref self: TContractState, market_id: felt252, shares: u256
 ) -> (u256, u256);
 
 // Withdraw exact token amounts from market.
-// Called for private vaults. For public vaults, use `withdraw_at_ratio`.
+// Called for private vaults. For public vaults, use `withdraw_public`.
 //
 // # Arguments
 // * `market_id` - market id
@@ -172,7 +172,7 @@ fn withdraw_at_ratio(
 // # Returns
 // * `base_amount` - base asset withdrawn
 // * `quote_amount` - quote asset withdrawn
-fn withdraw(
+fn withdraw_private(
     ref self: TContractState, market_id: felt252, base_amount: u256, quote_amount: u256
 ) -> (u256, u256);
 ```

packages/core/src/contracts/mocks/mock_solver.cairo

@@ -169,11 +169,14 @@ pub mod MockSolver {
         }
 
         // Callback function to execute any state updates after a swap is completed.
+        // This fn should only be callable by the solver contract.
         //
         // # Arguments
         // * `market_id` - market id
         // * `swap_params` - swap parameters
-        fn after_swap(ref self: ContractState, market_id: felt252, swap_params: SwapParams) {}
+        fn after_swap(ref self: ContractState, market_id: felt252, swap_params: SwapParams) {
+            assert(self.solver.unlocked.read(), 'NotSolver');
+        }
     }
 
     #[abi(embed_v0)]

packages/core/src/contracts/solver.cairo

@@ -53,6 +53,8 @@ pub mod SolverComponent {
         withdraw_fees: LegacyMap::<ContractAddress, u256>,
         // vault token class hash
         vault_token_class: ClassHash,
+        // reentrancy guard (unlocked for hook calls)
+        unlocked: bool,
     }
 
     ////////////////////////////////
@@ -377,6 +379,7 @@ pub mod SolverComponent {
             // Check params.
             assert(market_info.base_token != contract_address_const::<0x0>(), 'BaseTokenNull');
             assert(market_info.quote_token != contract_address_const::<0x0>(), 'QuoteTokenNull');
+            assert(market_info.base_token != market_info.quote_token, 'SameToken');
             assert(market_info.owner != contract_address_const::<0x0>(), 'OwnerNull');
 
             // Set market info.
@@ -387,7 +390,7 @@ pub mod SolverComponent {
             if market_info.is_public {
                 let vault_token_addr = self._deploy_vault_token(market_info);
                 vault_token = Option::Some(vault_token_addr);
-                let mut state: MarketState = self.market_state.read(market_id);
+                let mut state: MarketState = Default::default();
                 state.vault_token = vault_token_addr;
                 self.market_state.write(market_id, state);
             }
@@ -426,6 +429,12 @@ pub mod SolverComponent {
         fn swap(
             ref self: ComponentState<TContractState>, market_id: felt252, swap_params: SwapParams,
         ) -> (u256, u256) {
+            // Run validity checks.
+            let state: MarketState = self.market_state.read(market_id);
+            let market_info: MarketInfo = self.market_info.read(market_id);
+            assert(market_info.base_token != contract_address_const::<0x0>(), 'MarketNull');
+            assert(!state.is_paused, 'Paused');
+
             // Get amounts.
             let solver_hooks = ISolverHooksDispatcher { contract_address: get_contract_address() };
             let (amount_in, amount_out) = solver_hooks.quote(market_id, swap_params);
@@ -469,7 +478,9 @@ pub mod SolverComponent {
             self.market_state.write(market_id, state);
 
             // Execute after swap hook.
+            self.unlocked.write(true);
             solver_hooks.after_swap(market_id, swap_params);
+            self.unlocked.write(false);
 
             // Emit events.
             self
@@ -754,7 +765,7 @@ pub mod SolverComponent {
         }
 
         // Burn pool shares and withdraw funds from market.
-        // Called for public vaults. For private vaults, use `withdraw_amount`.
+        // Called for public vaults. For private vaults, use `withdraw_private`.
         //
         // # Arguments
         // * `market_id` - market id
@@ -763,7 +774,7 @@ pub mod SolverComponent {
         // # Returns
         // * `base_amount` - base asset withdrawn
         // * `quote_amount` - quote asset withdrawn
-        fn withdraw_at_ratio(
+        fn withdraw_public(
             ref self: ComponentState<TContractState>, market_id: felt252, shares: u256
         ) -> (u256, u256) {
             // Fetch state.
@@ -773,7 +784,7 @@ pub mod SolverComponent {
             // Run checks.
             assert(market_info.base_token != contract_address_const::<0x0>(), 'MarketNull');
             assert(shares != 0, 'SharesZero');
-            assert(market_info.is_public, 'UseWithdraw');
+            assert(market_info.is_public, 'UseWithdrawPrivate');
             let vault_token = ERC20ABIDispatcher { contract_address: state.vault_token };
             let caller = get_caller_address();
             assert(shares <= vault_token.balanceOf(caller), 'InsuffShares');
@@ -795,7 +806,7 @@ pub mod SolverComponent {
         }
 
         // Withdraw exact token amounts from market.
-        // Called for private vaults. For public vaults, use `withdraw_at_ratio`.
+        // Called for private vaults. For public vaults, use `withdraw_public`.
         //
         // # Arguments
         // * `market_id` - market id
@@ -805,7 +816,7 @@ pub mod SolverComponent {
         // # Returns
         // * `base_amount` - base asset withdrawn
         // * `quote_amount` - quote asset withdrawn
-        fn withdraw(
+        fn withdraw_private(
             ref self: ComponentState<TContractState>,
             market_id: felt252,
             base_amount: u256,
@@ -817,7 +828,8 @@ pub mod SolverComponent {
 
             // Run checks.
             assert(market_info.base_token != contract_address_const::<0x0>(), 'MarketNull');
-            assert(!market_info.is_public, 'UseWithdrawAtRatio');
+            assert(!market_info.is_public, 'UseWithdrawPublic');
+            self.assert_market_owner(market_id);
 
             // Cap withdraw amount at available. Commit state changes.
             let base_withdraw = min(base_amount, state.base_reserves);

packages/core/src/interfaces/ISolver.cairo

@@ -178,7 +178,7 @@ pub trait ISolver<TContractState> {
     ) -> (u256, u256, u256);
 
     // Burn pool shares and withdraw funds from market.
-    // Called for public vaults. For private vaults, use `withdraw_amount`.
+    // Called for public vaults. For private vaults, use `withdraw_private`.
     //
     // # Arguments
     // * `market_id` - market id
@@ -187,12 +187,10 @@ pub trait ISolver<TContractState> {
     // # Returns
     // * `base_amount` - base asset withdrawn
     // * `quote_amount` - quote asset withdrawn
-    fn withdraw_at_ratio(
-        ref self: TContractState, market_id: felt252, shares: u256
-    ) -> (u256, u256);
+    fn withdraw_public(ref self: TContractState, market_id: felt252, shares: u256) -> (u256, u256);
 
     // Withdraw exact token amounts from market.
-    // Called for private vaults. For public vaults, use `withdraw_at_ratio`.
+    // Called for private vaults. For public vaults, use `withdraw_public`.
     //
     // # Arguments
     // * `market_id` - market id
@@ -202,7 +200,7 @@ pub trait ISolver<TContractState> {
     // # Returns
     // * `base_amount` - base asset withdrawn
     // * `quote_amount` - quote asset withdrawn
-    fn withdraw(
+    fn withdraw_private(
         ref self: TContractState, market_id: felt252, base_amount: u256, quote_amount: u256
     ) -> (u256, u256);
 
@@ -263,7 +261,7 @@ pub trait ISolver<TContractState> {
     fn upgrade(ref self: TContractState, new_class_hash: ClassHash);
 }
 
-// Solvers must implement the `ISolverHooks` interface, which returns the quote for a swap.
+// Solvers must implement the `ISolverHooks` interface.
 #[starknet::interface]
 pub trait ISolverHooks<TContractState> {
     // Obtain quote for swap through a market.

packages/core/src/tests/libraries/test_erc20_versioned_call.cairo

@@ -23,7 +23,6 @@ fn test_erc20_bytearray() {
     symbol.serialize(ref calldata);
     let contract_address = erc20_class.deploy(@calldata).unwrap();
     let result = erc20_versioned_call::get_symbol(contract_address);
-    println!("result(bytearray): {}", result);
     assert(result == "MOCK", 'Symbol (byte array)');
 }
 
@@ -38,6 +37,5 @@ fn test_erc20_felt252() {
     symbol.serialize(ref calldata);
     let contract_address = erc20_class.deploy(@calldata).unwrap();
     let result = erc20_versioned_call::get_symbol(contract_address);
-    println!("result(felt252): {}", result);
     assert(result == "MOCK", 'Symbol (felt252)');
 }

packages/core/src/tests/solver/test_create_market.cairo

@@ -186,6 +186,25 @@ fn test_create_market_with_null_quote_token_fails() {
     solver.create_market(market_info);
 }
 
+#[test]
+#[should_panic(expected: ('SameToken',))]
+fn test_create_market_with_same_token_fails() {
+    let (base_token, _quote_token, _vault_token_class, solver, _market_id, _vault_token_opt) =
+        before(
+        true
+    );
+
+    // Create market.
+    start_prank(CheatTarget::One(solver.contract_address), owner());
+    let market_info = MarketInfo {
+        base_token: base_token.contract_address,
+        quote_token: base_token.contract_address,
+        owner: alice(),
+        is_public: true,
+    };
+    solver.create_market(market_info);
+}
+
 #[test]
 #[should_panic(expected: ('OwnerNull',))]
 fn test_create_market_with_null_owner_fails() {

packages/core/src/tests/solver/test_pause.cairo

@@ -48,7 +48,7 @@ fn test_pause_allows_withdraws() {
 
     // Withdraw.
     start_prank(CheatTarget::One(solver.contract_address), alice());
-    solver.withdraw_at_ratio(market_id, shares);
+    solver.withdraw_public(market_id, shares);
 }
 
 #[test]

packages/core/src/tests/solver/test_swap.cairo

@@ -428,3 +428,24 @@ fn test_swap_fails_if_swap_sell_below_threshold_amount() {
     };
     solver.swap(market_id, params);
 }
+
+#[test]
+#[should_panic(expected: ('NotSolver',))]
+fn test_after_swap_fails_for_non_solver_caller() {
+    let (_base_token, _quote_token, _vault_token_class, solver, market_id, _vault_token_opt) =
+        before(
+        false
+    );
+
+    // Call after swap.
+    start_prank(CheatTarget::One(solver.contract_address), alice());
+    let solver_hooks = ISolverHooksDispatcher { contract_address: solver.contract_address };
+    let params = SwapParams {
+        is_buy: true,
+        amount: to_e18(10),
+        exact_input: true,
+        threshold_sqrt_price: Option::None(()),
+        threshold_amount: Option::None(()),
+    };
+    solver_hooks.after_swap(market_id, params);
+}