add source chainid for proxy; double check for onTokenBridged; remove account

contracts/CrossChainUpgradeableProxy.sol

@@ -7,6 +7,8 @@ import "@openzeppelin/contracts/contracts/proxy/TransparentUpgradeableProxy.sol"
 
 interface IAMB {
   function messageSender() external view returns (address);
+
+  function messageSourceChainId() external view returns (bytes32);
 }
 
 interface IOmniBridge {
@@ -18,6 +20,7 @@ interface IOmniBridge {
  */
 contract CrossChainUpgradeableProxy is TransparentUpgradeableProxy {
   IOmniBridge public immutable omniBridge;
+  bytes32 public immutable adminChainId;
 
   /**
    * @dev Initializes an upgradeable proxy backed by the implementation at `_logic`.
@@ -26,16 +29,22 @@ contract CrossChainUpgradeableProxy is TransparentUpgradeableProxy {
     address _logic,
     address _admin,
     bytes memory _data,
-    IOmniBridge _omniBridge
+    IOmniBridge _omniBridge,
+    uint256 _adminChainId
   ) TransparentUpgradeableProxy(_logic, _admin, _data) {
     omniBridge = _omniBridge;
+    adminChainId = bytes32(uint256(_adminChainId));
   }
 
   /**
    * @dev Modifier used internally that will delegate the call to the implementation unless the sender is the cross chain admin.
    */
   modifier ifAdmin() override {
-    if (msg.sender == address(omniBridge) && omniBridge.bridgeContract().messageSender() == _admin()) {
+    if (
+      msg.sender == address(omniBridge) &&
+      omniBridge.bridgeContract().messageSourceChainId() == adminChainId &&
+      omniBridge.bridgeContract().messageSender() == _admin()
+    ) {
       _;
     } else {
       _fallback();

contracts/Mocks/MockAMB.sol

@@ -5,9 +5,11 @@ import { IAMB } from "../CrossChainUpgradeableProxy.sol";
 
 contract MockAMB is IAMB {
   address public xDomainMessageSender;
+  bytes32 public xDomainMessageChainId;
 
-  constructor(address _xDomainMessageSender) {
+  constructor(address _xDomainMessageSender, uint256 _xDomainMessageChainId) {
     xDomainMessageSender = _xDomainMessageSender;
+    xDomainMessageChainId = bytes32(uint256(_xDomainMessageChainId));
   }
 
   function setMessageSender(address _sender) external {
@@ -17,4 +19,8 @@ contract MockAMB is IAMB {
   function messageSender() external view override returns (address) {
     return xDomainMessageSender;
   }
+
+  function messageSourceChainId() external view override returns (bytes32) {
+    return xDomainMessageChainId;
+  }
 }

contracts/TornadoPool.sol

@@ -16,8 +16,6 @@ pragma experimental ABIEncoderV2;
 import "@openzeppelin/contracts/contracts/token/ERC20/IERC20.sol";
 import "./MerkleTreeWithHistory.sol";
 
-import "hardhat/console.sol";
-
 interface IERC6777 is IERC20 {
   function transferAndCall(
     address,
@@ -44,13 +42,15 @@ contract TornadoPool is MerkleTreeWithHistory, IERC20Receiver {
   int256 public constant MAX_EXT_AMOUNT = 2**248;
   uint256 public constant MAX_FEE = 2**248;
 
-  mapping(bytes32 => bool) public nullifierHashes;
   IVerifier public immutable verifier2;
   IVerifier public immutable verifier16;
   IERC6777 public immutable token;
   address public immutable omniBridge;
   address public immutable l1Unwrapper;
 
+  uint256 public totalDeposited;
+  mapping(bytes32 => bool) public nullifierHashes;
+
   struct ExtData {
     address recipient;
     int256 extAmount;
@@ -70,11 +70,6 @@ contract TornadoPool is MerkleTreeWithHistory, IERC20Receiver {
     bytes32 extDataHash;
   }
 
-  struct Register {
-    bytes pubKey;
-    bytes account;
-  }
-
   event NewCommitment(bytes32 commitment, uint256 index, bytes encryptedOutput);
   event NewNullifier(bytes32 nullifier);
   event PublicKey(address indexed owner, bytes key);
@@ -105,6 +100,7 @@ contract TornadoPool is MerkleTreeWithHistory, IERC20Receiver {
     if (_extData.extAmount > 0) {
       // for deposits from L2
       token.transferFrom(msg.sender, address(this), uint256(_extData.extAmount));
+      totalDeposited += uint256(_extData.extAmount);
     }
 
     _transact(_args, _extData);
@@ -130,6 +126,7 @@ contract TornadoPool is MerkleTreeWithHistory, IERC20Receiver {
       } else {
         token.transfer(_extData.recipient, uint256(-_extData.extAmount));
       }
+      totalDeposited -= uint256(-_extData.extAmount);
     }
     if (_extData.fee > 0) {
       token.transfer(_extData.relayer, _extData.fee);
@@ -203,32 +200,34 @@ contract TornadoPool is MerkleTreeWithHistory, IERC20Receiver {
     }
   }
 
-  function register(Register memory args) public {
+  function register(bytes memory _publicKey) public {
-    emit PublicKey(msg.sender, args.pubKey);
+    emit PublicKey(msg.sender, _publicKey);
-    emit EncryptedAccount(msg.sender, args.account);
   }
 
   function registerAndTransact(
-    Register memory _registerArgs,
+    bytes memory _publicKey,
     Proof memory _proofArgs,
     ExtData memory _extData
   ) public {
-    register(_registerArgs);
+    register(_publicKey);
     transact(_proofArgs, _extData);
   }
 
-  /// TOTHINK security. should we track all incoming trasfers so we can to double check the bridge actually sent tokens to this contract?
   function onTokenBridged(
     IERC6777 _token,
-    uint256,
+    uint256 _amount,
     bytes calldata _data
   ) external override {
+    (bytes memory _publicKey, Proof memory _args, ExtData memory _extData) = abi.decode(_data, (bytes, Proof, ExtData));
     require(_token == token, "provided token is not supported");
-    require(msg.sender == omniBridge, "only omni bridge"); // we can also get real msg.sender from L1, but it does not matter
+    require(msg.sender == omniBridge, "only omni bridge");
+    require(_amount == uint256(_extData.extAmount), "amount from bridge is incorrect");
+    require(uint256(_extData.extAmount) + totalDeposited >= token.balanceOf(address(this)), "bridge did not send enough tokens");
 
-    (Register memory _registerArgs, Proof memory _args, ExtData memory _extData) = abi.decode(_data, (Register, Proof, ExtData));
+    totalDeposited += uint256(_extData.extAmount);
-    if (_registerArgs.pubKey.length != 0 && _registerArgs.account.length != 0) {
+
-      register(_registerArgs);
+    if (_publicKey.length != 0) {
+      register(_publicKey);
     }
     _transact(_args, _extData);
   }

contracts/bridge/BridgeHelper.sol

@@ -0,0 +1,737 @@
+/**
+ *Submitted for verification at BscScan.com on 2021-03-09
+ */
+
+pragma solidity 0.7.6;
+
+/**
+ * @dev Wrappers over Solidity's arithmetic operations with added overflow
+ * checks.
+ *
+ * Arithmetic operations in Solidity wrap on overflow. This can easily result
+ * in bugs, because programmers usually assume that an overflow raises an
+ * error, which is the standard behavior in high level programming languages.
+ * `SafeMath` restores this intuition by reverting the transaction when an
+ * operation overflows.
+ *
+ * Using this library instead of the unchecked operations eliminates an entire
+ * class of bugs, so it's recommended to use it always.
+ */
+library SafeMath {
+  /**
+   * @dev Returns the addition of two unsigned integers, reverting on
+   * overflow.
+   *
+   * Counterpart to Solidity's `+` operator.
+   *
+   * Requirements:
+   *
+   * - Addition cannot overflow.
+   */
+  function add(uint256 a, uint256 b) internal pure returns (uint256) {
+    uint256 c = a + b;
+    require(c >= a, "SafeMath: addition overflow");
+
+    return c;
+  }
+
+  /**
+   * @dev Returns the subtraction of two unsigned integers, reverting on
+   * overflow (when the result is negative).
+   *
+   * Counterpart to Solidity's `-` operator.
+   *
+   * Requirements:
+   *
+   * - Subtraction cannot overflow.
+   */
+  function sub(uint256 a, uint256 b) internal pure returns (uint256) {
+    return sub(a, b, "SafeMath: subtraction overflow");
+  }
+
+  /**
+   * @dev Returns the subtraction of two unsigned integers, reverting with custom message on
+   * overflow (when the result is negative).
+   *
+   * Counterpart to Solidity's `-` operator.
+   *
+   * Requirements:
+   *
+   * - Subtraction cannot overflow.
+   */
+  function sub(
+    uint256 a,
+    uint256 b,
+    string memory errorMessage
+  ) internal pure returns (uint256) {
+    require(b <= a, errorMessage);
+    uint256 c = a - b;
+
+    return c;
+  }
+
+  /**
+   * @dev Returns the multiplication of two unsigned integers, reverting on
+   * overflow.
+   *
+   * Counterpart to Solidity's `*` operator.
+   *
+   * Requirements:
+   *
+   * - Multiplication cannot overflow.
+   */
+  function mul(uint256 a, uint256 b) internal pure returns (uint256) {
+    // Gas optimization: this is cheaper than requiring 'a' not being zero, but the
+    // benefit is lost if 'b' is also tested.
+    // See: https://github.com/OpenZeppelin/openzeppelin-contracts/pull/522
+    if (a == 0) {
+      return 0;
+    }
+
+    uint256 c = a * b;
+    require(c / a == b, "SafeMath: multiplication overflow");
+
+    return c;
+  }
+
+  /**
+   * @dev Returns the integer division of two unsigned integers. Reverts on
+   * division by zero. The result is rounded towards zero.
+   *
+   * Counterpart to Solidity's `/` operator. Note: this function uses a
+   * `revert` opcode (which leaves remaining gas untouched) while Solidity
+   * uses an invalid opcode to revert (consuming all remaining gas).
+   *
+   * Requirements:
+   *
+   * - The divisor cannot be zero.
+   */
+  function div(uint256 a, uint256 b) internal pure returns (uint256) {
+    return div(a, b, "SafeMath: division by zero");
+  }
+
+  /**
+   * @dev Returns the integer division of two unsigned integers. Reverts with custom message on
+   * division by zero. The result is rounded towards zero.
+   *
+   * Counterpart to Solidity's `/` operator. Note: this function uses a
+   * `revert` opcode (which leaves remaining gas untouched) while Solidity
+   * uses an invalid opcode to revert (consuming all remaining gas).
+   *
+   * Requirements:
+   *
+   * - The divisor cannot be zero.
+   */
+  function div(
+    uint256 a,
+    uint256 b,
+    string memory errorMessage
+  ) internal pure returns (uint256) {
+    require(b > 0, errorMessage);
+    uint256 c = a / b;
+    // assert(a == b * c + a % b); // There is no case in which this doesn't hold
+
+    return c;
+  }
+
+  /**
+   * @dev Returns the remainder of dividing two unsigned integers. (unsigned integer modulo),
+   * Reverts when dividing by zero.
+   *
+   * Counterpart to Solidity's `%` operator. This function uses a `revert`
+   * opcode (which leaves remaining gas untouched) while Solidity uses an
+   * invalid opcode to revert (consuming all remaining gas).
+   *
+   * Requirements:
+   *
+   * - The divisor cannot be zero.
+   */
+  function mod(uint256 a, uint256 b) internal pure returns (uint256) {
+    return mod(a, b, "SafeMath: modulo by zero");
+  }
+
+  /**
+   * @dev Returns the remainder of dividing two unsigned integers. (unsigned integer modulo),
+   * Reverts with custom message when dividing by zero.
+   *
+   * Counterpart to Solidity's `%` operator. This function uses a `revert`
+   * opcode (which leaves remaining gas untouched) while Solidity uses an
+   * invalid opcode to revert (consuming all remaining gas).
+   *
+   * Requirements:
+   *
+   * - The divisor cannot be zero.
+   */
+  function mod(
+    uint256 a,
+    uint256 b,
+    string memory errorMessage
+  ) internal pure returns (uint256) {
+    require(b != 0, errorMessage);
+    return a % b;
+  }
+}
+
+interface IOmnibridge {
+  function relayTokens(
+    address _token,
+    address _receiver,
+    uint256 _value
+  ) external;
+
+  /**
+   * @dev Initiate the bridge operation for some amount of tokens from msg.sender.
+   * The user should first call Approve method of the ERC677 token.
+   * @param token bridged token contract address.
+   * @param _receiver address that will receive the native tokens on the other network.
+   * @param _value amount of tokens to be transferred to the other network.
+   * @param _data additional transfer data to be used on the other side.
+   */
+  function relayTokensAndCall(
+    address token,
+    address _receiver,
+    uint256 _value,
+    bytes memory _data
+  ) external;
+}
+
+interface IWETH {
+  function deposit() external payable;
+
+  function withdraw(uint256 _value) external;
+
+  function approve(address _to, uint256 _value) external;
+}
+
+contract Sacrifice {
+  constructor(address payable _recipient) payable {
+    selfdestruct(_recipient);
+  }
+}
+
+/**
+ * @title AddressHelper
+ * @dev Helper methods for Address type.
+ */
+library AddressHelper {
+  /**
+   * @dev Try to send native tokens to the address. If it fails, it will force the transfer by creating a selfdestruct contract
+   * @param _receiver address that will receive the native tokens
+   * @param _value the amount of native tokens to send
+   */
+  function safeSendValue(address payable _receiver, uint256 _value) internal {
+    if (!(_receiver).send(_value)) {
+      new Sacrifice{ value: _value }(_receiver);
+    }
+  }
+}
+
+/**
+ * @dev Collection of functions related to the address type
+ */
+library Address {
+  /**
+   * @dev Returns true if `account` is a contract.
+   *
+   * [IMPORTANT]
+   * ====
+   * It is unsafe to assume that an address for which this function returns
+   * false is an externally-owned account (EOA) and not a contract.
+   *
+   * Among others, `isContract` will return false for the following
+   * types of addresses:
+   *
+   *  - an externally-owned account
+   *  - a contract in construction
+   *  - an address where a contract will be created
+   *  - an address where a contract lived, but was destroyed
+   * ====
+   */
+  function isContract(address account) internal view returns (bool) {
+    // According to EIP-1052, 0x0 is the value returned for not-yet created accounts
+    // and 0xc5d2460186f7233c927e7db2dcc703c0e500b653ca82273b7bfad8045d85a470 is returned
+    // for accounts without code, i.e. `keccak256('')`
+    bytes32 codehash;
+    bytes32 accountHash = 0xc5d2460186f7233c927e7db2dcc703c0e500b653ca82273b7bfad8045d85a470;
+    // solhint-disable-next-line no-inline-assembly
+    assembly {
+      codehash := extcodehash(account)
+    }
+    return (codehash != accountHash && codehash != 0x0);
+  }
+
+  /**
+   * @dev Replacement for Solidity's `transfer`: sends `amount` wei to
+   * `recipient`, forwarding all available gas and reverting on errors.
+   *
+   * https://eips.ethereum.org/EIPS/eip-1884[EIP1884] increases the gas cost
+   * of certain opcodes, possibly making contracts go over the 2300 gas limit
+   * imposed by `transfer`, making them unable to receive funds via
+   * `transfer`. {sendValue} removes this limitation.
+   *
+   * https://diligence.consensys.net/posts/2019/09/stop-using-soliditys-transfer-now/[Learn more].
+   *
+   * IMPORTANT: because control is transferred to `recipient`, care must be
+   * taken to not create reentrancy vulnerabilities. Consider using
+   * {ReentrancyGuard} or the
+   * https://solidity.readthedocs.io/en/v0.5.11/security-considerations.html#use-the-checks-effects-interactions-pattern[checks-effects-interactions pattern].
+   */
+  function sendValue(address payable recipient, uint256 amount) internal {
+    require(address(this).balance >= amount, "Address: insufficient balance");
+
+    // solhint-disable-next-line avoid-low-level-calls, avoid-call-value
+    (bool success, ) = recipient.call{ value: amount }("");
+    require(success, "Address: unable to send value, recipient may have reverted");
+  }
+
+  /**
+   * @dev Performs a Solidity function call using a low level `call`. A
+   * plain`call` is an unsafe replacement for a function call: use this
+   * function instead.
+   *
+   * If `target` reverts with a revert reason, it is bubbled up by this
+   * function (like regular Solidity function calls).
+   *
+   * Returns the raw returned data. To convert to the expected return value,
+   * use https://solidity.readthedocs.io/en/latest/units-and-global-variables.html?highlight=abi.decode#abi-encoding-and-decoding-functions[`abi.decode`].
+   *
+   * Requirements:
+   *
+   * - `target` must be a contract.
+   * - calling `target` with `data` must not revert.
+   *
+   * _Available since v3.1._
+   */
+  function functionCall(address target, bytes memory data) internal returns (bytes memory) {
+    return functionCall(target, data, "Address: low-level call failed");
+  }
+
+  /**
+   * @dev Same as {xref-Address-functionCall-address-bytes-}[`functionCall`], but with
+   * `errorMessage` as a fallback revert reason when `target` reverts.
+   *
+   * _Available since v3.1._
+   */
+  function functionCall(
+    address target,
+    bytes memory data,
+    string memory errorMessage
+  ) internal returns (bytes memory) {
+    return _functionCallWithValue(target, data, 0, errorMessage);
+  }
+
+  /**
+   * @dev Same as {xref-Address-functionCall-address-bytes-}[`functionCall`],
+   * but also transferring `value` wei to `target`.
+   *
+   * Requirements:
+   *
+   * - the calling contract must have an ETH balance of at least `value`.
+   * - the called Solidity function must be `payable`.
+   *
+   * _Available since v3.1._
+   */
+  function functionCallWithValue(
+    address target,
+    bytes memory data,
+    uint256 value
+  ) internal returns (bytes memory) {
+    return functionCallWithValue(target, data, value, "Address: low-level call with value failed");
+  }
+
+  /**
+   * @dev Same as {xref-Address-functionCallWithValue-address-bytes-uint256-}[`functionCallWithValue`], but
+   * with `errorMessage` as a fallback revert reason when `target` reverts.
+   *
+   * _Available since v3.1._
+   */
+  function functionCallWithValue(
+    address target,
+    bytes memory data,
+    uint256 value,
+    string memory errorMessage
+  ) internal returns (bytes memory) {
+    require(address(this).balance >= value, "Address: insufficient balance for call");
+    return _functionCallWithValue(target, data, value, errorMessage);
+  }
+
+  function _functionCallWithValue(
+    address target,
+    bytes memory data,
+    uint256 weiValue,
+    string memory errorMessage
+  ) private returns (bytes memory) {
+    require(isContract(target), "Address: call to non-contract");
+
+    // solhint-disable-next-line avoid-low-level-calls
+    (bool success, bytes memory returndata) = target.call{ value: weiValue }(data);
+    if (success) {
+      return returndata;
+    } else {
+      // Look for revert reason and bubble it up if present
+      if (returndata.length > 0) {
+        // The easiest way to bubble the revert reason is using memory via assembly
+
+        // solhint-disable-next-line no-inline-assembly
+        assembly {
+          let returndata_size := mload(returndata)
+          revert(add(32, returndata), returndata_size)
+        }
+      } else {
+        revert(errorMessage);
+      }
+    }
+  }
+}
+
+/**
+ * @dev Interface of the ERC20 standard as defined in the EIP.
+ */
+interface IERC20 {
+  /**
+   * @dev Returns the amount of tokens in existence.
+   */
+  function totalSupply() external view returns (uint256);
+
+  /**
+   * @dev Returns the amount of tokens owned by `account`.
+   */
+  function balanceOf(address account) external view returns (uint256);
+
+  /**
+   * @dev Moves `amount` tokens from the caller's account to `recipient`.
+   *
+   * Returns a boolean value indicating whether the operation succeeded.
+   *
+   * Emits a {Transfer} event.
+   */
+  function transfer(address recipient, uint256 amount) external returns (bool);
+
+  /**
+   * @dev Returns the remaining number of tokens that `spender` will be
+   * allowed to spend on behalf of `owner` through {transferFrom}. This is
+   * zero by default.
+   *
+   * This value changes when {approve} or {transferFrom} are called.
+   */
+  function allowance(address owner, address spender) external view returns (uint256);
+
+  /**
+   * @dev Sets `amount` as the allowance of `spender` over the caller's tokens.
+   *
+   * Returns a boolean value indicating whether the operation succeeded.
+   *
+   * IMPORTANT: Beware that changing an allowance with this method brings the risk
+   * that someone may use both the old and the new allowance by unfortunate
+   * transaction ordering. One possible solution to mitigate this race
+   * condition is to first reduce the spender's allowance to 0 and set the
+   * desired value afterwards:
+   * https://github.com/ethereum/EIPs/issues/20#issuecomment-263524729
+   *
+   * Emits an {Approval} event.
+   */
+  function approve(address spender, uint256 amount) external returns (bool);
+
+  /**
+   * @dev Moves `amount` tokens from `sender` to `recipient` using the
+   * allowance mechanism. `amount` is then deducted from the caller's
+   * allowance.
+   *
+   * Returns a boolean value indicating whether the operation succeeded.
+   *
+   * Emits a {Transfer} event.
+   */
+  function transferFrom(
+    address sender,
+    address recipient,
+    uint256 amount
+  ) external returns (bool);
+
+  /**
+   * @dev Emitted when `value` tokens are moved from one account (`from`) to
+   * another (`to`).
+   *
+   * Note that `value` may be zero.
+   */
+  event Transfer(address indexed from, address indexed to, uint256 value);
+
+  /**
+   * @dev Emitted when the allowance of a `spender` for an `owner` is set by
+   * a call to {approve}. `value` is the new allowance.
+   */
+  event Approval(address indexed owner, address indexed spender, uint256 value);
+}
+
+/**
+ * @title SafeERC20
+ * @dev Wrappers around ERC20 operations that throw on failure (when the token
+ * contract returns false). Tokens that return no value (and instead revert or
+ * throw on failure) are also supported, non-reverting calls are assumed to be
+ * successful.
+ * To use this library you can add a `using SafeERC20 for IERC20;` statement to your contract,
+ * which allows you to call the safe operations as `token.safeTransfer(...)`, etc.
+ */
+library SafeERC20 {
+  using SafeMath for uint256;
+  using Address for address;
+
+  function safeTransfer(
+    IERC20 token,
+    address to,
+    uint256 value
+  ) internal {
+    _callOptionalReturn(token, abi.encodeWithSelector(token.transfer.selector, to, value));
+  }
+
+  function safeTransferFrom(
+    IERC20 token,
+    address from,
+    address to,
+    uint256 value
+  ) internal {
+    _callOptionalReturn(token, abi.encodeWithSelector(token.transferFrom.selector, from, to, value));
+  }
+
+  /**
+   * @dev Deprecated. This function has issues similar to the ones found in
+   * {IERC20-approve}, and its usage is discouraged.
+   *
+   * Whenever possible, use {safeIncreaseAllowance} and
+   * {safeDecreaseAllowance} instead.
+   */
+  function safeApprove(
+    IERC20 token,
+    address spender,
+    uint256 value
+  ) internal {
+    // safeApprove should only be called when setting an initial allowance,
+    // or when resetting it to zero. To increase and decrease it, use
+    // 'safeIncreaseAllowance' and 'safeDecreaseAllowance'
+    // solhint-disable-next-line max-line-length
+    require(
+      (value == 0) || (token.allowance(address(this), spender) == 0),
+      "SafeERC20: approve from non-zero to non-zero allowance"
+    );
+    _callOptionalReturn(token, abi.encodeWithSelector(token.approve.selector, spender, value));
+  }
+
+  function safeIncreaseAllowance(
+    IERC20 token,
+    address spender,
+    uint256 value
+  ) internal {
+    uint256 newAllowance = token.allowance(address(this), spender).add(value);
+    _callOptionalReturn(token, abi.encodeWithSelector(token.approve.selector, spender, newAllowance));
+  }
+
+  function safeDecreaseAllowance(
+    IERC20 token,
+    address spender,
+    uint256 value
+  ) internal {
+    uint256 newAllowance = token.allowance(address(this), spender).sub(value, "SafeERC20: decreased allowance below zero");
+    _callOptionalReturn(token, abi.encodeWithSelector(token.approve.selector, spender, newAllowance));
+  }
+
+  /**
+   * @dev Imitates a Solidity high-level call (i.e. a regular function call to a contract), relaxing the requirement
+   * on the return value: the return value is optional (but if data is returned, it must not be false).
+   * @param token The token targeted by the call.
+   * @param data The call data (encoded using abi.encode or one of its variants).
+   */
+  function _callOptionalReturn(IERC20 token, bytes memory data) private {
+    // We need to perform a low level call here, to bypass Solidity's return data size checking mechanism, since
+    // we're implementing it ourselves. We use {Address.functionCall} to perform this call, which verifies that
+    // the target address contains contract code and also asserts for success in the low-level call.
+
+    bytes memory returndata = address(token).functionCall(data, "SafeERC20: low-level call failed");
+    if (returndata.length > 0) {
+      // Return data is optional
+      // solhint-disable-next-line max-line-length
+      require(abi.decode(returndata, (bool)), "SafeERC20: ERC20 operation did not succeed");
+    }
+  }
+}
+
+/**
+ * @title OwnableModule
+ * @dev Common functionality for multi-token extension non-upgradeable module.
+ */
+contract OwnableModule {
+  address public owner;
+
+  /**
+   * @dev Initializes this contract.
+   * @param _owner address of the owner that is allowed to perform additional actions on the particular module.
+   */
+  constructor(address _owner) {
+    owner = _owner;
+  }
+
+  /**
+   * @dev Throws if sender is not the owner of this contract.
+   */
+  modifier onlyOwner() {
+    require(msg.sender == owner);
+    _;
+  }
+
+  /**
+   * @dev Changes the owner of this contract.
+   * @param _newOwner address of the new owner.
+   */
+  function transferOwnership(address _newOwner) external onlyOwner {
+    owner = _newOwner;
+  }
+}
+
+/**
+ * @title Claimable
+ * @dev Implementation of the claiming utils that can be useful for withdrawing accidentally sent tokens that are not used in bridge operations.
+ */
+contract Claimable {
+  using SafeERC20 for IERC20;
+
+  /**
+   * Throws if a given address is equal to address(0)
+   */
+  modifier validAddress(address _to) {
+    require(_to != address(0));
+    _;
+  }
+
+  /**
+   * @dev Withdraws the erc20 tokens or native coins from this contract.
+   * Caller should additionally check that the claimed token is not a part of bridge operations (i.e. that token != erc20token()).
+   * @param _token address of the claimed token or address(0) for native coins.
+   * @param _to address of the tokens/coins receiver.
+   */
+  function claimValues(address _token, address _to) internal validAddress(_to) {
+    if (_token == address(0)) {
+      claimNativeCoins(_to);
+    } else {
+      claimErc20Tokens(_token, _to);
+    }
+  }
+
+  /**
+   * @dev Internal function for withdrawing all native coins from the contract.
+   * @param _to address of the coins receiver.
+   */
+  function claimNativeCoins(address _to) internal {
+    uint256 value = address(this).balance;
+    AddressHelper.safeSendValue(payable(_to), value);
+  }
+
+  /**
+   * @dev Internal function for withdrawing all tokens of ssome particular ERC20 contract from this contract.
+   * @param _token address of the claimed ERC20 token.
+   * @param _to address of the tokens receiver.
+   */
+  function claimErc20Tokens(address _token, address _to) internal {
+    IERC20 token = IERC20(_token);
+    uint256 balance = token.balanceOf(address(this));
+    token.safeTransfer(_to, balance);
+  }
+}
+
+/**
+ * @title WETHOmnibridgeRouter
+ * @dev Omnibridge extension for processing native and wrapped native assets.
+ * Intended to work with WETH/WBNB/WXDAI tokens, see:
+ *   https://etherscan.io/address/0xc02aaa39b223fe8d0a0e5c4f27ead9083c756cc2
+ *   https://bscscan.com/address/0xbb4CdB9CBd36B01bD1cBaEBF2De08d9173bc095c
+ *   https://blockscout.com/poa/xdai/address/0xe91D153E0b41518A2Ce8Dd3D7944Fa863463a97d
+ */
+contract WETHOmnibridgeRouter is OwnableModule, Claimable {
+  IOmnibridge public immutable bridge;
+  IWETH public immutable WETH;
+
+  /**
+   * @dev Initializes this contract.
+   * @param _bridge address of the HomeOmnibridge/ForeignOmnibridge contract.
+   * @param _weth address of the WETH token used for wrapping/unwrapping native coins (e.g. WETH/WBNB/WXDAI).
+   * @param _owner address of the contract owner.
+   */
+  constructor(
+    IOmnibridge _bridge,
+    IWETH _weth,
+    address _owner
+  ) OwnableModule(_owner) {
+    bridge = _bridge;
+    WETH = _weth;
+    _weth.approve(address(_bridge), uint256(-1));
+  }
+
+  /**
+   * @dev Wraps native assets and relays wrapped ERC20 tokens to the other chain.
+   * Call msg.sender will receive assets on the other side of the bridge.
+   */
+  function wrapAndRelayTokens() external payable {
+    wrapAndRelayTokens(msg.sender);
+  }
+
+  /**
+   * @dev Wraps native assets and relays wrapped ERC20 tokens to the other chain.
+   * @param _receiver bridged assets receiver on the other side of the bridge.
+   */
+  function wrapAndRelayTokens(address _receiver) public payable {
+    WETH.deposit{ value: msg.value }();
+    bridge.relayTokens(address(WETH), _receiver, msg.value);
+  }
+
+  /**
+   * @dev Wraps native assets and relays wrapped ERC20 tokens to the other chain.
+   * It also calls receiver on other side with the _data provided.
+   * @param _receiver bridged assets receiver on the other side of the bridge.
+   * @param _data data for the call of receiver on other side.
+   */
+  function wrapAndRelayTokensWithData(address _receiver, bytes memory _data) public payable {
+    WETH.deposit{ value: msg.value }();
+    bridge.relayTokensAndCall(address(WETH), _receiver, msg.value, _data);
+  }
+
+  /**
+   * @dev Bridged callback function used for unwrapping received tokens.
+   * Can only be called by the associated Omnibridge contract.
+   * @param _token bridged token contract address, should be WETH.
+   * @param _value amount of bridged/received tokens.
+   * @param _data extra data passed alongside with relayTokensAndCall on the other side of the bridge.
+   * Should contain coins receiver address.
+   */
+  function onTokenBridged(
+    address _token,
+    uint256 _value,
+    bytes calldata _data
+  ) external {
+    require(_token == address(WETH));
+    require(msg.sender == address(bridge));
+    require(_data.length == 20);
+
+    WETH.withdraw(_value);
+
+    address payable receiver;
+    assembly {
+      receiver := calldataload(120)
+    }
+    AddressHelper.safeSendValue(receiver, _value);
+  }
+
+  /**
+   * @dev Claims stuck coins/tokens.
+   * Only contract owner can call this method.
+   * @param _token address of claimed token contract, address(0) for native coins.
+   * @param _to address of tokens receiver
+   */
+  function claimTokens(address _token, address _to) external onlyOwner {
+    claimValues(_token, _to);
+  }
+
+  /**
+   * @dev Ether receive function.
+   * Should be only called from the WETH contract when withdrawing native coins. Will revert otherwise.
+   */
+  receive() external payable {
+    require(msg.sender == address(WETH));
+  }
+}

src/index.js

@@ -149,18 +149,13 @@ async function transaction({ tornadoPool, ...rest }) {
   return await receipt.wait()
 }
 
-async function registerAndTransact({ tornadoPool, packedPrivateKeyData, poolAddress, ...rest }) {
+async function registerAndTransact({ tornadoPool, poolAddress, ...rest }) {
   const { args, extData } = await prepareTransaction({
     tornadoPool,
     ...rest,
   })
 
-  const params = {
+  const receipt = await tornadoPool.registerAndTransact(poolAddress, args, extData, {
-    pubKey: poolAddress,
-    account: packedPrivateKeyData,
-  }
-
-  const receipt = await tornadoPool.registerAndTransact(params, args, extData, {
     gasLimit: 2e6,
   })
   await receipt.wait()

test/full.test.js

@@ -29,7 +29,7 @@ describe('TornadoPool', function () {
     const token = await deploy('PermittableToken', 'Wrapped ETH', 'WETH', 18, 1)
     await token.mint(sender.address, utils.parseEther('10000'))
 
-    const amb = await deploy('MockAMB', gov.address)
+    const amb = await deploy('MockAMB', gov.address, 1)
     const omniBridge = await deploy('MockOmniBridge', amb.address)
 
     /** @type {TornadoPool} */
@@ -58,6 +58,7 @@ describe('TornadoPool', function () {
       gov.address,
       [],
       omniBridge.address,
+      1,
     )
 
     /** @type {TornadoPool} */
@@ -110,15 +111,9 @@ describe('TornadoPool', function () {
     const aliceDepositAmount = 1e7
     const aliceDepositUtxo = new Utxo({ amount: aliceDepositAmount })
 
-    const backupAccount = new Keypair()
-
-    const bufferPrivateKey = Buffer.from(aliceDepositUtxo.keypair.privkey)
-    const packedPrivateKeyData = backupAccount.encrypt(bufferPrivateKey)
-
     tornadoPool = tornadoPool.connect(sender)
     await registerAndTransact({
       tornadoPool,
-      packedPrivateKeyData,
       outputs: [aliceDepositUtxo],
       poolAddress: aliceDepositUtxo.keypair.address(),
     })
@@ -151,16 +146,6 @@ describe('TornadoPool', function () {
     const [registerEvent] = registerEvents.sort((a, b) => a.blockNumber - b.blockNumber).slice(-1)
 
     expect(registerEvent.args.key).to.be.equal(aliceDepositUtxo.keypair.address())
-
-    const accountFilter = tornadoPool.filters.EncryptedAccount(sender.address)
-    const accountFromBlock = await ethers.provider.getBlock()
-    const accountEvents = await tornadoPool.queryFilter(accountFilter, accountFromBlock.number)
-
-    const [accountEvent] = accountEvents.sort((a, b) => a.blockNumber - b.blockNumber).slice(-1)
-
-    const privateKey = backupAccount.decrypt(accountEvent.args.account)
-
-    expect(bufferPrivateKey.toString('hex')).to.be.equal(privateKey.toString('hex'))
   })
 
   it('should deposit, transact and withdraw', async function () {
@@ -223,11 +208,7 @@ describe('TornadoPool', function () {
       tornadoPool,
       outputs: [aliceDepositUtxo],
     })
-    const transactTx = await tornadoPool.populateTransaction.registerAndTransact(
+    const transactTx = await tornadoPool.populateTransaction.registerAndTransact([], args, extData)
-      { pubKey: [], account: [] },
-      args,
-      extData,
-    )
     const onTokenBridgedData = '0x' + transactTx.data.slice(10)
     const onTokenBridgedTx = await tornadoPool.populateTransaction.onTokenBridged(
       token.address,