From af3f42ebe56ec6564509ce2d46bb8012861978ec Mon Sep 17 00:00:00 2001 From: kstasi Date: Wed, 16 Jun 2021 16:17:04 +0300 Subject: [PATCH] reentracy protection draft --- contracts/main/TTDex.ligo | 1 + contracts/partials/ITTDex.ligo | 2 ++ contracts/partials/TTDex.ligo | 11 +++++++++++ contracts/partials/TTMethodDex.ligo | 24 ++++++++++++++++++++++-- 4 files changed, 36 insertions(+), 2 deletions(-) diff --git a/contracts/main/TTDex.ligo b/contracts/main/TTDex.ligo index 90f2116f..691ff562 100644 --- a/contracts/main/TTDex.ligo +++ b/contracts/main/TTDex.ligo @@ -13,6 +13,7 @@ function main (const p : full_action; const s : full_dex_storage) : full_return | Balance_of(params) -> call_token(IBalance_of(params), this, 2n, s) | Update_operators(params) -> call_token(IUpdate_operators(params), this, 1n, s) | Get_reserves(params) -> get_reserves(params, s) + | Close -> ((nil:list(operation)), close(s)) | SetDexFunction(params) -> ((nil:list(operation)), if params.index > 4n then (failwith("Dex/wrong-index") : full_dex_storage) else set_dex_function(params.index, params.func, s)) | SetTokenFunction(params) -> ((nil:list(operation)), if params.index > 2n then (failwith("Dex/wrong-index") : full_dex_storage) else set_token_function(params.index, params.func, s)) end diff --git a/contracts/partials/ITTDex.ligo b/contracts/partials/ITTDex.ligo index 3433abc9..0eefe89b 100644 --- a/contracts/partials/ITTDex.ligo +++ b/contracts/partials/ITTDex.ligo @@ -41,6 +41,7 @@ type token_pair is bytes (* record for the dex storage *) type dex_storage is record [ + entered : bool; pairs_count : nat; (* total shares count *) tokens : big_map(nat, tokens_info); (* all the tokens list *) token_to_id : big_map(token_pair, nat); (* all the tokens list *) @@ -144,6 +145,7 @@ type full_action is | Balance_of of balance_params | Update_operators of update_operator_params | Get_reserves of get_reserves_params +| Close of unit | SetDexFunction of set_dex_function_params (* sets the dex specific function. Is used before the whole system is launched *) | SetTokenFunction of set_token_function_params (* sets the FA function, is used before the whole system is launched *) diff --git a/contracts/partials/TTDex.ligo b/contracts/partials/TTDex.ligo index 03325058..1c74ef70 100644 --- a/contracts/partials/TTDex.ligo +++ b/contracts/partials/TTDex.ligo @@ -47,6 +47,17 @@ block { s.storage := res.1; } with (res.0, s) +[@inline] function close (const s : full_dex_storage) : full_dex_storage is +block { + if not s.storage.entered then + failwith("Dex/not-entered") + else skip; + if Tezos.sender =/= Tezos.self_address then + failwith("Dex/not-self") + else skip; + s.storage.entered := False; +} with s + (* Return the reserves to the contracts. *) [@inline] function get_reserves (const params : get_reserves_params; const s : full_dex_storage) : full_return is block { diff --git a/contracts/partials/TTMethodDex.ligo b/contracts/partials/TTMethodDex.ligo index 335429f2..4cfe66a3 100644 --- a/contracts/partials/TTMethodDex.ligo +++ b/contracts/partials/TTMethodDex.ligo @@ -94,6 +94,10 @@ function initialize_exchange (const p : dex_action ; const s : dex_storage ; con var operations : list(operation) := list[]; case p of | InitializeExchange(params) -> { + if s.entered then + failwith("Dex/reentrancy") + else s.entered := True; + (* check preconditions *) if params.pair.token_a_address = params.pair.token_b_address and params.pair.token_a_id >= params.pair.token_b_id then failwith("Dex/wrong-token-id") @@ -169,7 +173,7 @@ function initialize_exchange (const p : dex_action ; const s : dex_storage ; con (* prepare operations to get initial liquidity *) case params.pair.token_b_type of | Fa12 -> { - operations := + operations := transfer_fa12( Tezos.sender, this, @@ -177,7 +181,7 @@ function initialize_exchange (const p : dex_action ; const s : dex_storage ; con params.pair.token_b_address) # operations; } | Fa2 -> { - operations := + operations := transfer_fa2( Tezos.sender, this, @@ -203,6 +207,10 @@ function token_to_token (const p : dex_action; const s : dex_storage; const this | InitializeExchange(n) -> skip | TokenToTokenRoutePayment(n) -> skip | TokenToTokenPayment(params) -> { + if s.entered then + failwith("Dex/reentrancy") + else s.entered := True; + (* check preconditions *) if params.pair.token_a_address = params.pair.token_b_address and params.pair.token_a_id >= params.pair.token_b_id then failwith("Dex/wrong-token-id") @@ -492,6 +500,10 @@ function token_to_token_route (const p : dex_action; const s : dex_storage; cons | InitializeExchange(n) -> skip | TokenToTokenPayment(n) -> skip | TokenToTokenRoutePayment(params) -> { + if s.entered then + failwith("Dex/reentrancy") + else s.entered := True; + if List.size(params.swaps) > 1n (* non-zero amount of tokens exchanged *) then skip else failwith ("Dex/too-few-swaps"); @@ -600,6 +612,10 @@ function invest_liquidity (const p : dex_action; const s : dex_storage; const th | TokenToTokenRoutePayment(n) -> skip | TokenToTokenPayment(n) -> skip | InvestLiquidity(params) -> { + if s.entered then + failwith("Dex/reentrancy") + else s.entered := True; + (* check preconditions *) if params.pair.token_a_address = params.pair.token_b_address and params.pair.token_a_id >= params.pair.token_b_id then failwith("Dex/wrong-token-id") @@ -718,6 +734,10 @@ function divest_liquidity (const p : dex_action; const s : dex_storage; const th | TokenToTokenRoutePayment(n) -> skip | InvestLiquidity(n) -> skip | DivestLiquidity(params) -> { + if s.entered then + failwith("Dex/reentrancy") + else s.entered := True; + (* check preconditions *) if params.pair.token_a_address = params.pair.token_b_address and params.pair.token_a_id >= params.pair.token_b_id then failwith("Dex/wrong-token-id")