Change withdrawERC20 => withdrawAsset, reuse transfer logic

6 files changed, 40 insertions, 22 deletions

diff --git a/packages/contracts/src/2.0.0/forwarder/MixinAssets.sol b/packages/contracts/src/2.0.0/forwarder/MixinAssets.sol
index 5cf5f831b..e06f9a8e3 100644
--- a/packages/contracts/src/2.0.0/forwarder/MixinAssets.sol
+++ b/packages/contracts/src/2.0.0/forwarder/MixinAssets.sol
@@ -36,28 +36,25 @@ contract MixinAssets is
 
     bytes4 constant internal ERC20_TRANSFER_SELECTOR = bytes4(keccak256("transfer(address,uint256)"));
 
-    /// @dev Withdraws ERC20 tokens from this contract. The contract requires a ZRX balance in order to 
+    /// @dev Withdraws assets from this contract. The contract requires a ZRX balance in order to 
     ///      function optimally, and this function allows the ZRX to be withdrawn by owner. It may also be
-    ///      used to withdraw tokens that were accidentally sent to this contract.
-    /// @param token Address of ERC20 token to withdraw.
+    ///      used to withdraw assets that were accidentally sent to this contract.
+    /// @param assetData Byte array encoded for the respective asset proxy.
     /// @param amount Amount of ERC20 token to withdraw.
-    function withdrawERC20(
-        address token,
+    function withdrawAsset(
+        bytes assetData,
         uint256 amount
     )
         external
         onlyOwner
     {
-        require(
-            IERC20Token(token).transfer(msg.sender, amount),
-            "WITHDRAWAL_FAILED"
-        );
+        transferAssetToSender(assetData, amount);
     }
 
     /// @dev Transfers given amount of asset to sender.
     /// @param assetData Byte array encoded for the respective asset proxy.
     /// @param amount Amount of asset to transfer to sender.
-    function transferPurchasedAssetToSender(
+    function transferAssetToSender(
         bytes memory assetData,
         uint256 amount
     )
diff --git a/packages/contracts/src/2.0.0/forwarder/MixinForwarderCore.sol b/packages/contracts/src/2.0.0/forwarder/MixinForwarderCore.sol
index 0d313ea91..93cbf79be 100644
--- a/packages/contracts/src/2.0.0/forwarder/MixinForwarderCore.sol
+++ b/packages/contracts/src/2.0.0/forwarder/MixinForwarderCore.sol
@@ -135,7 +135,7 @@ contract MixinForwarderCore is
         );
 
         // Transfer purchased assets to msg.sender.
-        transferPurchasedAssetToSender(orders[0].makerAssetData, makerAssetAmountPurchased);
+        transferAssetToSender(orders[0].makerAssetData, makerAssetAmountPurchased);
     }
 
     /// @dev Attempt to purchase makerAssetFillAmount of makerAsset by selling ETH provided with transaction.
@@ -208,6 +208,6 @@ contract MixinForwarderCore is
         );
 
         // Transfer purchased assets to msg.sender.
-        transferPurchasedAssetToSender(orders[0].makerAssetData, makerAssetAmountPurchased);
+        transferAssetToSender(orders[0].makerAssetData, makerAssetAmountPurchased);
     }
 }
diff --git a/packages/contracts/src/2.0.0/forwarder/interfaces/IAssets.sol b/packages/contracts/src/2.0.0/forwarder/interfaces/IAssets.sol
index 9b0d995eb..1e034c003 100644
--- a/packages/contracts/src/2.0.0/forwarder/interfaces/IAssets.sol
+++ b/packages/contracts/src/2.0.0/forwarder/interfaces/IAssets.sol
@@ -21,13 +21,13 @@ pragma solidity 0.4.24;
 
 contract IAssets {
 
-    /// @dev Withdraws ERC20 tokens from this contract. The contract requires a ZRX balance in order to 
+    /// @dev Withdraws assets from this contract. The contract requires a ZRX balance in order to 
     ///      function optimally, and this function allows the ZRX to be withdrawn by owner. It may also be
-    ///      used to withdraw tokens that were accidentally sent to this contract.
-    /// @param token Address of ERC20 token to withdraw.
+    ///      used to withdraw assets that were accidentally sent to this contract.
+    /// @param assetData Byte array encoded for the respective asset proxy.
     /// @param amount Amount of ERC20 token to withdraw.
-    function withdrawERC20(
-        address token,
+    function withdrawAsset(
+        bytes assetData,
         uint256 amount
     )
         external;
diff --git a/packages/contracts/src/2.0.0/forwarder/mixins/MAssets.sol b/packages/contracts/src/2.0.0/forwarder/mixins/MAssets.sol
index 340ee0bcb..83636432a 100644
--- a/packages/contracts/src/2.0.0/forwarder/mixins/MAssets.sol
+++ b/packages/contracts/src/2.0.0/forwarder/mixins/MAssets.sol
@@ -28,7 +28,7 @@ contract MAssets is
     /// @dev Transfers given amount of asset to sender.
     /// @param assetData Byte array encoded for the respective asset proxy.
     /// @param amount Amount of asset to transfer to sender.
-    function transferPurchasedAssetToSender(
+    function transferAssetToSender(
         bytes memory assetData,
         uint256 amount
     )
diff --git a/packages/contracts/test/forwarder/forwarder.ts b/packages/contracts/test/forwarder/forwarder.ts
index cd7ae59c2..28ffdeabe 100644
--- a/packages/contracts/test/forwarder/forwarder.ts
+++ b/packages/contracts/test/forwarder/forwarder.ts
@@ -36,6 +36,7 @@ describe(ContractName.Forwarder, () => {
     let feeRecipientAddress: string;
     let otherAddress: string;
     let defaultMakerAssetAddress: string;
+    let zrxAssetData: string;
 
     let weth: DummyERC20TokenContract;
     let zrxToken: DummyERC20TokenContract;
@@ -90,7 +91,7 @@ describe(ContractName.Forwarder, () => {
         erc20Wrapper.addDummyTokenContract(weth);
 
         const wethAssetData = assetDataUtils.encodeERC20AssetData(wethContract.address);
-        const zrxAssetData = assetDataUtils.encodeERC20AssetData(zrxToken.address);
+        zrxAssetData = assetDataUtils.encodeERC20AssetData(zrxToken.address);
         const exchangeInstance = await ExchangeContract.deployFrom0xArtifactAsync(
             artifacts.Exchange,
             provider,
@@ -970,6 +971,26 @@ describe(ContractName.Forwarder, () => {
             );
         });
     });
+    describe('withdrawAsset', () => {
+        it('should allow owner to withdraw ERC20 tokens', async () => {
+            const zrxWithdrawAmount = erc20Balances[forwarderContract.address][zrxToken.address];
+            await forwarderWrapper.withdrawAssetAsync(zrxAssetData, zrxWithdrawAmount, { from: owner });
+            const newBalances = await erc20Wrapper.getBalancesAsync();
+            expect(newBalances[owner][zrxToken.address]).to.be.bignumber.equal(
+                erc20Balances[owner][zrxToken.address].plus(zrxWithdrawAmount),
+            );
+            expect(newBalances[forwarderContract.address][zrxToken.address]).to.be.bignumber.equal(
+                erc20Balances[forwarderContract.address][zrxToken.address].minus(zrxWithdrawAmount),
+            );
+        });
+        it('should revert if not called by owner', async () => {
+            const zrxWithdrawAmount = erc20Balances[forwarderContract.address][zrxToken.address];
+            await expectTransactionFailedAsync(
+                forwarderWrapper.withdrawAssetAsync(zrxAssetData, zrxWithdrawAmount, { from: makerAddress }),
+                RevertReason.OnlyContractOwner,
+            );
+        });
+    });
 });
 // tslint:disable:max-file-line-count
 // tslint:enable:no-unnecessary-type-assertion
diff --git a/packages/contracts/test/utils/forwarder_wrapper.ts b/packages/contracts/test/utils/forwarder_wrapper.ts
index ef7476e36..5b9a63ddf 100644
--- a/packages/contracts/test/utils/forwarder_wrapper.ts
+++ b/packages/contracts/test/utils/forwarder_wrapper.ts
@@ -106,12 +106,12 @@ export class ForwarderWrapper {
         const tx = await this._logDecoder.getTxWithDecodedLogsAsync(txHash);
         return tx;
     }
-    public async withdrawERC20Async(
-        tokenAddress: string,
+    public async withdrawAssetAsync(
+        assetData: string,
         amount: BigNumber,
         txData: TxDataPayable,
     ): Promise<TransactionReceiptWithDecodedLogs> {
-        const txHash = await this._forwarderContract.withdrawERC20.sendTransactionAsync(tokenAddress, amount, txData);
+        const txHash = await this._forwarderContract.withdrawAsset.sendTransactionAsync(assetData, amount, txData);
         const tx = await this._logDecoder.getTxWithDecodedLogsAsync(txHash);
         return tx;
     }