From 5a82105f241eb7f501cf6292f30bd33851541a12 Mon Sep 17 00:00:00 2001
From: Alexey <peretsmobil@gmail.com>
Date: Thu, 11 Mar 2021 21:12:11 +0400
Subject: [PATCH] updates

---
 contracts/TornadoProxy.sol                | 23 +++++++++++++++++------
 contracts/interfaces/ITornadoInstance.sol |  4 ++++
 2 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/contracts/TornadoProxy.sol b/contracts/TornadoProxy.sol
index 264110f..9151376 100644
--- a/contracts/TornadoProxy.sol
+++ b/contracts/TornadoProxy.sol
@@ -17,8 +17,9 @@ contract TornadoProxy {
 
   enum InstanceState { Disabled, Enabled, Mineable }
   struct Instance {
-    address instance;
+    ITornadoInstance instance;
     InstanceState state;
+    bool isERC20;
   }
 
   ITornadoTrees public tornadoTrees;
@@ -39,7 +40,7 @@ contract TornadoProxy {
     governance = _governance;
 
     for (uint256 i = 0; i < _instances.length; i++) {
-      instances[ITornadoInstance(_instances[i].instance)] = _instances[i].state;
+      _updateInstance(_instances[i]);
     }
   }
 
@@ -50,7 +51,9 @@ contract TornadoProxy {
   ) external payable {
     require(instances[_tornado] != InstanceState.Disabled, "The instance is not supported");
 
+    IERC20(_tornado.token()).safeTransferFrom(msg.sender, address(this), _tornado.denomination());
     _tornado.deposit{ value: msg.value }(_commitment);
+
     if (instances[_tornado] == InstanceState.Mineable) {
       tornadoTrees.registerDeposit(address(_tornado), _commitment);
     }
@@ -75,13 +78,13 @@ contract TornadoProxy {
     }
   }
 
-  function updateInstance(ITornadoInstance _instance, InstanceState _state) external onlyGovernance {
-    instances[_instance] = _state;
+  function updateInstance(Instance calldata _instance) external onlyGovernance {
+    _updateInstance(_instance);
     emit InstanceStateUpdate(address(_instance), _state);
   }
 
-  function setTornadoTreesContract(address _instance) external onlyGovernance {
-    tornadoTrees = ITornadoTrees(_instance);
+  function setTornadoTreesContract(address _tornadoTrees) external onlyGovernance {
+    tornadoTrees = ITornadoTrees(_tornadoTrees);
   }
 
   /// @dev Method to claim junk and accidentally sent tokens
@@ -105,4 +108,12 @@ contract TornadoProxy {
       _token.safeTransfer(_to, balance);
     }
   }
+
+  function _updateInstance(Instance memory _tornado) internal {
+    instances[_tornado.instance] = _tornado.state;
+    if (_tornado.isERC20) {
+      IERC20 token = IERC20(_tornado.instance.token());
+      token.safeApprove(address(_tornado.instance), uint256(-1));
+    }
+  }
 }
diff --git a/contracts/interfaces/ITornadoInstance.sol b/contracts/interfaces/ITornadoInstance.sol
index 952586d..b719d30 100644
--- a/contracts/interfaces/ITornadoInstance.sol
+++ b/contracts/interfaces/ITornadoInstance.sol
@@ -3,6 +3,10 @@
 pragma solidity ^0.6.0;
 
 interface ITornadoInstance {
+  function token() external returns (address);
+
+  function denomination() external returns (uint256);
+
   function deposit(bytes32 commitment) external payable;
 
   function withdraw(