reentracy protection draft

madfish-solutions · Jun 16, 2021 · af3f42e · af3f42e
1 parent 9757aeb
commit af3f42e
Show file tree

Hide file tree

Showing 4 changed files with 36 additions and 2 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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,15 +173,15 @@ 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,
                 params.token_b_in,
                 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")