add protocol fees fixin (#2)

contracts/zero-ex/CHANGELOG.json

@@ -1,4 +1,13 @@
 [
+    {
+        "version": "0.6.0",
+        "changes": [
+            {
+                "note": "Add support for collecting protocol fees in ETH or WETH",
+                "pr": 2
+            }
+        ]
+    },
     {
         "timestamp": 1603851023,
         "version": "0.5.1",

contracts/zero-ex/contracts/src/external/FeeCollector.sol

@@ -0,0 +1,65 @@
+/*
+
+  Copyright 2020 ZeroEx Intl.
+
+  Licensed under the Apache License, Version 2.0 (the "License");
+  you may not use this file except in compliance with the License.
+  You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+  Unless required by applicable law or agreed to in writing, software
+  distributed under the License is distributed on an "AS IS" BASIS,
+  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  See the License for the specific language governing permissions and
+  limitations under the License.
+
+*/
+
+pragma solidity ^0.6.5;
+pragma experimental ABIEncoderV2;
+
+import "@0x/contracts-erc20/contracts/src/v06/IEtherTokenV06.sol";
+import "@0x/contracts-utils/contracts/src/v06/AuthorizableV06.sol";
+import "../vendor/v3/IStaking.sol";
+
+/// @dev The collector contract for protocol fees
+contract FeeCollector is AuthorizableV06 {
+    /// @dev Allow ether transfers to the collector.
+    receive() external payable { }
+
+    constructor() public {
+        _addAuthorizedAddress(msg.sender);
+    }
+
+    /// @dev   Approve the staking contract and join a pool. Only an authority
+    ///        can call this.
+    /// @param weth The WETH contract.
+    /// @param staking The staking contract.
+    /// @param poolId The pool ID this contract is collecting fees for.
+    function initialize(
+        IEtherTokenV06 weth,
+        IStaking staking,
+        bytes32 poolId
+    )
+        external
+        onlyAuthorized
+    {
+        weth.approve(address(staking), type(uint256).max);
+        staking.joinStakingPoolAsMaker(poolId);
+    }
+
+    /// @dev Convert all held ether to WETH. Only an authority can call this.
+    /// @param weth The WETH contract.
+    function convertToWeth(
+        IEtherTokenV06 weth
+    )
+        external
+        onlyAuthorized
+    {
+        // Leave 1 wei behind to avoid expensive zero-->non-zero state change.
+        if (address(this).balance > 1) {
+            weth.deposit{value: address(this).balance - 1}();
+        }
+    }
+}

contracts/zero-ex/contracts/src/features/libs/LibTokenSpender.sol

@@ -27,6 +27,9 @@ import "../ITokenSpenderFeature.sol";
 library LibTokenSpender {
     using LibRichErrorsV06 for bytes;
 
+    // Mask of the lower 20 bytes of a bytes32.
+    uint256 constant private ADDRESS_MASK = 0x000000000000000000000000ffffffffffffffffffffffffffffffffffffffff;
+
     /// @dev Transfers ERC20 tokens from `owner` to `to`.
     /// @param token The token to spend.
     /// @param owner The owner of the tokens.
@@ -50,11 +53,11 @@ library LibTokenSpender {
 
             // selector for transferFrom(address,address,uint256)
             mstore(ptr, 0x23b872dd00000000000000000000000000000000000000000000000000000000)
-            mstore(add(ptr, 0x04), owner)
+            mstore(add(ptr, 0x04), and(owner, ADDRESS_MASK))
-            mstore(add(ptr, 0x24), to)
+            mstore(add(ptr, 0x24), and(to, ADDRESS_MASK))
             mstore(add(ptr, 0x44), amount)
 
-            success := call(gas(), token, 0, ptr, 0x64, 0, 0)
+            success := call(gas(), and(token, ADDRESS_MASK), 0, ptr, 0x64, 0, 0)
 
             let rdsize := returndatasize()
 

contracts/zero-ex/contracts/src/fixins/FixinProtocolFees.sol

@@ -0,0 +1,114 @@
+/*
+
+  Copyright 2020 ZeroEx Intl.
+
+  Licensed under the Apache License, Version 2.0 (the "License");
+  you may not use this file except in compliance with the License.
+  You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+  Unless required by applicable law or agreed to in writing, software
+  distributed under the License is distributed on an "AS IS" BASIS,
+  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  See the License for the specific language governing permissions and
+  limitations under the License.
+
+*/
+
+pragma solidity ^0.6.5;
+pragma experimental ABIEncoderV2;
+
+import "@0x/contracts-erc20/contracts/src/v06/IEtherTokenV06.sol";
+import "../external/FeeCollector.sol";
+import "../features/libs/LibTokenSpender.sol";
+
+/// @dev Helpers for collecting protocol fees.
+abstract contract FixinProtocolFees {
+    bytes32 immutable feeCollectorCodeHash;
+
+    constructor() internal {
+        feeCollectorCodeHash = keccak256(type(FeeCollector).creationCode);
+    }
+
+    /// @dev   Collect the specified protocol fee in either WETH or ETH. If
+    ///        msg.value is non-zero, the fee will be paid in ETH. Otherwise,
+    ///        this function attempts to transfer the fee in WETH. Either way,
+    ///        The fee is stored in a per-pool fee collector contract.
+    /// @param poolId The pool ID for which a fee is being collected.
+    /// @param amount The amount of ETH/WETH to be collected.
+    /// @param weth The WETH token contract.
+    function _collectProtocolFee(
+        bytes32 poolId,
+        uint256 amount,
+        IERC20TokenV06 weth
+    )
+        internal
+    {
+        FeeCollector feeCollector = _getFeeCollector(poolId);
+
+        if (msg.value == 0) {
+            // WETH
+            LibTokenSpender.spendERC20Tokens(weth, msg.sender, address(feeCollector), amount);
+        } else {
+            // ETH
+            (bool success,) = address(feeCollector).call{value: amount}("");
+            require(success, "FixinProtocolFees/ETHER_TRANSFER_FALIED");
+        }
+    }
+
+    /// @dev Transfer fees for a given pool to the staking contract.
+    /// @param poolId Identifies the pool whose fees are being paid.
+    function _transferFeesForPool(
+        bytes32 poolId,
+        IStaking staking,
+        IEtherTokenV06 weth
+    )
+        internal
+    {
+        FeeCollector feeCollector = _getFeeCollector(poolId);
+
+        uint256 codeSize;
+        assembly {
+            codeSize := extcodesize(feeCollector)
+        }
+
+        if (codeSize == 0) {
+            // Create and initialize the contract if necessary.
+            new FeeCollector{salt: poolId}();
+            feeCollector.initialize(weth, staking, poolId);
+        }
+
+        if (address(feeCollector).balance > 1) {
+            feeCollector.convertToWeth(weth);
+        }
+
+        uint256 bal = weth.balanceOf(address(feeCollector));
+        if (bal > 1) {
+            // Leave 1 wei behind to avoid high SSTORE cost of zero-->non-zero.
+            staking.payProtocolFee(
+                address(feeCollector),
+                address(feeCollector),
+                bal - 1);
+        }
+    }
+
+    /// @dev Compute the CREATE2 address for a fee collector.
+    /// @param poolId The fee collector's pool ID.
+    function _getFeeCollector(
+        bytes32 poolId
+    )
+        internal
+        view
+        returns (FeeCollector)
+    {
+        // Compute the CREATE2 address for the fee collector.
+        address payable addr = address(uint256(keccak256(abi.encodePacked(
+            byte(0xff),
+            address(this),
+            poolId, // pool ID is salt
+            feeCollectorCodeHash
+        ))));
+        return FeeCollector(addr);
+    }
+}

contracts/zero-ex/contracts/src/vendor/v3/IStaking.sol

@@ -0,0 +1,24 @@
+/*
+
+  Copyright 2020 ZeroEx Intl.
+
+  Licensed under the Apache License, Version 2.0 (the "License");
+  you may not use this file except in compliance with the License.
+  You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+  Unless required by applicable law or agreed to in writing, software
+  distributed under the License is distributed on an "AS IS" BASIS,
+  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  See the License for the specific language governing permissions and
+  limitations under the License.
+
+*/
+
+pragma solidity ^0.6.5;
+
+interface IStaking {
+    function joinStakingPoolAsMaker(bytes32) external;
+    function payProtocolFee(address, address, uint256) external payable;
+}

contracts/zero-ex/contracts/test/TestProtocolFees.sol

@@ -0,0 +1,55 @@
+/*
+
+  Copyright 2020 ZeroEx Intl.
+
+  Licensed under the Apache License, Version 2.0 (the "License");
+  you may not use this file except in compliance with the License.
+  You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+  Unless required by applicable law or agreed to in writing, software
+  distributed under the License is distributed on an "AS IS" BASIS,
+  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  See the License for the specific language governing permissions and
+  limitations under the License.
+
+*/
+
+pragma solidity ^0.6.5;
+pragma experimental ABIEncoderV2;
+
+import "../src/fixins/FixinProtocolFees.sol";
+
+contract TestProtocolFees is FixinProtocolFees {
+    function collectProtocolFee(
+        bytes32 poolId,
+        uint256 amount,
+        IERC20TokenV06 weth
+    )
+        external
+        payable
+    {
+        _collectProtocolFee(poolId, amount, weth);
+    }
+
+    function transferFeesForPool(
+        bytes32 poolId,
+        IStaking staking,
+        IEtherTokenV06 weth
+    )
+        external
+    {
+        _transferFeesForPool(poolId, staking, weth);
+    }
+
+    function getFeeCollector(
+        bytes32 poolId
+    )
+        external
+        view
+        returns (FeeCollector)
+    {
+        return _getFeeCollector(poolId);
+    }
+}

contracts/zero-ex/contracts/test/TestStaking.sol

@@ -0,0 +1,49 @@
+/*
+
+  Copyright 2020 ZeroEx Intl.
+
+  Licensed under the Apache License, Version 2.0 (the "License");
+  you may not use this file except in compliance with the License.
+  You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+  Unless required by applicable law or agreed to in writing, software
+  distributed under the License is distributed on an "AS IS" BASIS,
+  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  See the License for the specific language governing permissions and
+  limitations under the License.
+
+*/
+
+pragma solidity ^0.6.5;
+pragma experimental ABIEncoderV2;
+
+import "@0x/contracts-erc20/contracts/src/v06/IEtherTokenV06.sol";
+
+contract TestStaking {
+    mapping(address => bytes32) public poolForMaker;
+    mapping(bytes32 => uint256) public balanceForPool;
+
+    IEtherTokenV06 immutable weth;
+
+    constructor(IEtherTokenV06 _weth) public {
+        weth = _weth;
+    }
+
+    function joinStakingPoolAsMaker(bytes32 poolId) external {
+        poolForMaker[msg.sender] = poolId;
+    }
+
+    function payProtocolFee(
+        address makerAddress,
+        address payerAddress,
+        uint256 amount
+    )
+        external
+        payable
+    {
+        require(weth.transferFrom(payerAddress, address(this), amount));
+        balanceForPool[poolForMaker[makerAddress]] += amount;
+    }
+}

contracts/zero-ex/package.json

@@ -41,7 +41,7 @@
     "config": {
         "publicInterfaceContracts": "IZeroEx,ZeroEx,FullMigration,InitialMigration,IFlashWallet,IAllowanceTarget,IERC20Transformer,IOwnableFeature,ISimpleFunctionRegistryFeature,ITokenSpenderFeature,ITransformERC20Feature,FillQuoteTransformer,PayTakerTransformer,WethTransformer,OwnableFeature,SimpleFunctionRegistryFeature,TransformERC20Feature,TokenSpenderFeature,AffiliateFeeTransformer,SignatureValidatorFeature,MetaTransactionsFeature,LogMetadataTransformer,BridgeAdapter,LiquidityProviderFeature",
         "abis:comment": "This list is auto-generated by contracts-gen. Don't edit manually.",
-        "abis": "./test/generated-artifacts/@(AffiliateFeeTransformer|AllowanceTarget|BootstrapFeature|BridgeAdapter|FillQuoteTransformer|FixinCommon|FixinEIP712|FixinReentrancyGuard|FlashWallet|FullMigration|IAllowanceTarget|IBootstrapFeature|IBridgeAdapter|IERC20Bridge|IERC20Transformer|IExchange|IFeature|IFlashWallet|IGasToken|ILiquidityProviderFeature|IMetaTransactionsFeature|IOwnableFeature|ISignatureValidatorFeature|ISimpleFunctionRegistryFeature|ITestSimpleFunctionRegistryFeature|ITokenSpenderFeature|ITransformERC20Feature|IUniswapFeature|IZeroEx|InitialMigration|LibBootstrap|LibCommonRichErrors|LibERC20Transformer|LibLiquidityProviderRichErrors|LibLiquidityProviderStorage|LibMetaTransactionsRichErrors|LibMetaTransactionsStorage|LibMigrate|LibOwnableRichErrors|LibOwnableStorage|LibProxyRichErrors|LibProxyStorage|LibReentrancyGuardStorage|LibSignatureRichErrors|LibSignedCallData|LibSimpleFunctionRegistryRichErrors|LibSimpleFunctionRegistryStorage|LibSpenderRichErrors|LibStorage|LibTokenSpender|LibTokenSpenderStorage|LibTransformERC20RichErrors|LibTransformERC20Storage|LibWalletRichErrors|LiquidityProviderFeature|LogMetadataTransformer|MetaTransactionsFeature|MixinAdapterAddresses|MixinBalancer|MixinCurve|MixinDodo|MixinKyber|MixinMStable|MixinMooniswap|MixinOasis|MixinShell|MixinSushiswap|MixinUniswap|MixinUniswapV2|MixinZeroExBridge|OwnableFeature|PayTakerTransformer|SignatureValidatorFeature|SimpleFunctionRegistryFeature|TestBridge|TestCallTarget|TestDelegateCaller|TestFillQuoteTransformerBridge|TestFillQuoteTransformerExchange|TestFillQuoteTransformerHost|TestFullMigration|TestInitialMigration|TestLibTokenSpender|TestMetaTransactionsTransformERC20Feature|TestMigrator|TestMintTokenERC20Transformer|TestMintableERC20Token|TestSimpleFunctionRegistryFeatureImpl1|TestSimpleFunctionRegistryFeatureImpl2|TestTokenSpender|TestTokenSpenderERC20Token|TestTransformERC20|TestTransformerBase|TestTransformerDeployerTransformer|TestTransformerHost|TestWeth|TestWethTransformerHost|TestZeroExFeature|TokenSpenderFeature|TransformERC20Feature|Transformer|TransformerDeployer|UniswapFeature|WethTransformer|ZeroEx).json"
+        "abis": "./test/generated-artifacts/@(AffiliateFeeTransformer|AllowanceTarget|BootstrapFeature|BridgeAdapter|FeeCollector|FillQuoteTransformer|FixinCommon|FixinEIP712|FixinProtocolFees|FixinReentrancyGuard|FlashWallet|FullMigration|IAllowanceTarget|IBootstrapFeature|IBridgeAdapter|IERC20Bridge|IERC20Transformer|IExchange|IFeature|IFlashWallet|IGasToken|ILiquidityProviderFeature|IMetaTransactionsFeature|IOwnableFeature|ISignatureValidatorFeature|ISimpleFunctionRegistryFeature|IStaking|ITestSimpleFunctionRegistryFeature|ITokenSpenderFeature|ITransformERC20Feature|IUniswapFeature|IZeroEx|InitialMigration|LibBootstrap|LibCommonRichErrors|LibERC20Transformer|LibLiquidityProviderRichErrors|LibLiquidityProviderStorage|LibMetaTransactionsRichErrors|LibMetaTransactionsStorage|LibMigrate|LibOwnableRichErrors|LibOwnableStorage|LibProxyRichErrors|LibProxyStorage|LibReentrancyGuardStorage|LibSignatureRichErrors|LibSignedCallData|LibSimpleFunctionRegistryRichErrors|LibSimpleFunctionRegistryStorage|LibSpenderRichErrors|LibStorage|LibTokenSpender|LibTokenSpenderStorage|LibTransformERC20RichErrors|LibTransformERC20Storage|LibWalletRichErrors|LiquidityProviderFeature|LogMetadataTransformer|MetaTransactionsFeature|MixinAdapterAddresses|MixinBalancer|MixinCurve|MixinDodo|MixinKyber|MixinMStable|MixinMooniswap|MixinOasis|MixinShell|MixinSushiswap|MixinUniswap|MixinUniswapV2|MixinZeroExBridge|OwnableFeature|PayTakerTransformer|SignatureValidatorFeature|SimpleFunctionRegistryFeature|TestBridge|TestCallTarget|TestDelegateCaller|TestFillQuoteTransformerBridge|TestFillQuoteTransformerExchange|TestFillQuoteTransformerHost|TestFullMigration|TestInitialMigration|TestLibTokenSpender|TestMetaTransactionsTransformERC20Feature|TestMigrator|TestMintTokenERC20Transformer|TestMintableERC20Token|TestProtocolFees|TestSimpleFunctionRegistryFeatureImpl1|TestSimpleFunctionRegistryFeatureImpl2|TestStaking|TestTokenSpender|TestTokenSpenderERC20Token|TestTransformERC20|TestTransformerBase|TestTransformerDeployerTransformer|TestTransformerHost|TestWeth|TestWethTransformerHost|TestZeroExFeature|TokenSpenderFeature|TransformERC20Feature|Transformer|TransformerDeployer|UniswapFeature|WethTransformer|ZeroEx).json"
     },
     "repository": {
         "type": "git",

contracts/zero-ex/test/artifacts.ts

@@ -9,9 +9,11 @@ import * as AffiliateFeeTransformer from '../test/generated-artifacts/AffiliateF
 import * as AllowanceTarget from '../test/generated-artifacts/AllowanceTarget.json';
 import * as BootstrapFeature from '../test/generated-artifacts/BootstrapFeature.json';
 import * as BridgeAdapter from '../test/generated-artifacts/BridgeAdapter.json';
+import * as FeeCollector from '../test/generated-artifacts/FeeCollector.json';
 import * as FillQuoteTransformer from '../test/generated-artifacts/FillQuoteTransformer.json';
 import * as FixinCommon from '../test/generated-artifacts/FixinCommon.json';
 import * as FixinEIP712 from '../test/generated-artifacts/FixinEIP712.json';
+import * as FixinProtocolFees from '../test/generated-artifacts/FixinProtocolFees.json';
 import * as FixinReentrancyGuard from '../test/generated-artifacts/FixinReentrancyGuard.json';
 import * as FlashWallet from '../test/generated-artifacts/FlashWallet.json';
 import * as FullMigration from '../test/generated-artifacts/FullMigration.json';
@@ -30,6 +32,7 @@ import * as InitialMigration from '../test/generated-artifacts/InitialMigration.
 import * as IOwnableFeature from '../test/generated-artifacts/IOwnableFeature.json';
 import * as ISignatureValidatorFeature from '../test/generated-artifacts/ISignatureValidatorFeature.json';
 import * as ISimpleFunctionRegistryFeature from '../test/generated-artifacts/ISimpleFunctionRegistryFeature.json';
+import * as IStaking from '../test/generated-artifacts/IStaking.json';
 import * as ITestSimpleFunctionRegistryFeature from '../test/generated-artifacts/ITestSimpleFunctionRegistryFeature.json';
 import * as ITokenSpenderFeature from '../test/generated-artifacts/ITokenSpenderFeature.json';
 import * as ITransformERC20Feature from '../test/generated-artifacts/ITransformERC20Feature.json';
@@ -92,8 +95,10 @@ import * as TestMetaTransactionsTransformERC20Feature from '../test/generated-ar
 import * as TestMigrator from '../test/generated-artifacts/TestMigrator.json';
 import * as TestMintableERC20Token from '../test/generated-artifacts/TestMintableERC20Token.json';
 import * as TestMintTokenERC20Transformer from '../test/generated-artifacts/TestMintTokenERC20Transformer.json';
+import * as TestProtocolFees from '../test/generated-artifacts/TestProtocolFees.json';
 import * as TestSimpleFunctionRegistryFeatureImpl1 from '../test/generated-artifacts/TestSimpleFunctionRegistryFeatureImpl1.json';
 import * as TestSimpleFunctionRegistryFeatureImpl2 from '../test/generated-artifacts/TestSimpleFunctionRegistryFeatureImpl2.json';
+import * as TestStaking from '../test/generated-artifacts/TestStaking.json';
 import * as TestTokenSpender from '../test/generated-artifacts/TestTokenSpender.json';
 import * as TestTokenSpenderERC20Token from '../test/generated-artifacts/TestTokenSpenderERC20Token.json';
 import * as TestTransformerBase from '../test/generated-artifacts/TestTransformerBase.json';
@@ -124,6 +129,7 @@ export const artifacts = {
     LibTransformERC20RichErrors: LibTransformERC20RichErrors as ContractArtifact,
     LibWalletRichErrors: LibWalletRichErrors as ContractArtifact,
     AllowanceTarget: AllowanceTarget as ContractArtifact,
+    FeeCollector: FeeCollector as ContractArtifact,
     FlashWallet: FlashWallet as ContractArtifact,
     IAllowanceTarget: IAllowanceTarget as ContractArtifact,
     IFlashWallet: IFlashWallet as ContractArtifact,
@@ -151,6 +157,7 @@ export const artifacts = {
     LibTokenSpender: LibTokenSpender as ContractArtifact,
     FixinCommon: FixinCommon as ContractArtifact,
     FixinEIP712: FixinEIP712 as ContractArtifact,
+    FixinProtocolFees: FixinProtocolFees as ContractArtifact,
     FixinReentrancyGuard: FixinReentrancyGuard as ContractArtifact,
     FullMigration: FullMigration as ContractArtifact,
     InitialMigration: InitialMigration as ContractArtifact,
@@ -191,6 +198,7 @@ export const artifacts = {
     IERC20Bridge: IERC20Bridge as ContractArtifact,
     IExchange: IExchange as ContractArtifact,
     IGasToken: IGasToken as ContractArtifact,
+    IStaking: IStaking as ContractArtifact,
     ITestSimpleFunctionRegistryFeature: ITestSimpleFunctionRegistryFeature as ContractArtifact,
     TestBridge: TestBridge as ContractArtifact,
     TestCallTarget: TestCallTarget as ContractArtifact,
@@ -205,8 +213,10 @@ export const artifacts = {
     TestMigrator: TestMigrator as ContractArtifact,
     TestMintTokenERC20Transformer: TestMintTokenERC20Transformer as ContractArtifact,
     TestMintableERC20Token: TestMintableERC20Token as ContractArtifact,
+    TestProtocolFees: TestProtocolFees as ContractArtifact,
     TestSimpleFunctionRegistryFeatureImpl1: TestSimpleFunctionRegistryFeatureImpl1 as ContractArtifact,
     TestSimpleFunctionRegistryFeatureImpl2: TestSimpleFunctionRegistryFeatureImpl2 as ContractArtifact,
+    TestStaking: TestStaking as ContractArtifact,
     TestTokenSpender: TestTokenSpender as ContractArtifact,
     TestTokenSpenderERC20Token: TestTokenSpenderERC20Token as ContractArtifact,
     TestTransformERC20: TestTransformERC20 as ContractArtifact,

contracts/zero-ex/test/protocol_fees_test.ts

@@ -0,0 +1,75 @@
+import { blockchainTests, constants, expect } from '@0x/contracts-test-utils';
+import { BigNumber, hexUtils } from '@0x/utils';
+
+import { artifacts } from './artifacts';
+import { TestProtocolFeesContract, TestStakingContract, TestWethContract } from './wrappers';
+
+blockchainTests.resets('ProtocolFees', env => {
+    let payer: string;
+    let protocolFees: TestProtocolFeesContract;
+    let staking: TestStakingContract;
+    let weth: TestWethContract;
+
+    before(async () => {
+        [payer] = await env.getAccountAddressesAsync();
+        protocolFees = await TestProtocolFeesContract.deployFrom0xArtifactAsync(
+            artifacts.TestProtocolFees,
+            env.provider,
+            env.txDefaults,
+            artifacts,
+        );
+        weth = await TestWethContract.deployFrom0xArtifactAsync(
+            artifacts.TestWeth,
+            env.provider,
+            env.txDefaults,
+            artifacts,
+        );
+        staking = await TestStakingContract.deployFrom0xArtifactAsync(
+            artifacts.TestStaking,
+            env.provider,
+            env.txDefaults,
+            artifacts,
+            weth.address,
+        );
+        await weth.mint(payer, constants.ONE_ETHER).awaitTransactionSuccessAsync();
+        await weth.approve(protocolFees.address, constants.ONE_ETHER).awaitTransactionSuccessAsync({ from: payer });
+    });
+
+    describe('_collectProtocolFee()', () => {
+        it('can collect a protocol fee multiple times', async () => {
+            const poolId = hexUtils.random();
+            const amount1 = new BigNumber(123456);
+            const amount2 = new BigNumber(456789);
+
+            // Transfer amount1 via WETH.
+            await protocolFees
+                .collectProtocolFee(poolId, amount1, weth.address)
+                .awaitTransactionSuccessAsync({ from: payer });
+
+            // Send to staking contract.
+            await protocolFees
+                .transferFeesForPool(poolId, staking.address, weth.address)
+                .awaitTransactionSuccessAsync();
+
+            // Transfer amount2 via ETH.
+            await protocolFees
+                .collectProtocolFee(poolId, amount2, weth.address)
+                .awaitTransactionSuccessAsync({ from: payer, value: amount2 });
+
+            // Send to staking contract again.
+            await protocolFees
+                .transferFeesForPool(poolId, staking.address, weth.address)
+                .awaitTransactionSuccessAsync();
+
+            const balance = await staking.balanceForPool(poolId).callAsync();
+            const wethBalance = await weth.balanceOf(staking.address).callAsync();
+
+            // Check that staking accounted for the collected ether properly.
+            expect(balance).to.bignumber.eq(wethBalance);
+
+            // We leave 1 wei behind, of both ETH and WETH, for gas reasons.
+            const total = amount1.plus(amount2).minus(2);
+            return expect(balance).to.bignumber.eq(total);
+        });
+    });
+});

contracts/zero-ex/test/wrappers.ts

@@ -7,9 +7,11 @@ export * from '../test/generated-wrappers/affiliate_fee_transformer';
 export * from '../test/generated-wrappers/allowance_target';
 export * from '../test/generated-wrappers/bootstrap_feature';
 export * from '../test/generated-wrappers/bridge_adapter';
+export * from '../test/generated-wrappers/fee_collector';
 export * from '../test/generated-wrappers/fill_quote_transformer';
 export * from '../test/generated-wrappers/fixin_common';
 export * from '../test/generated-wrappers/fixin_e_i_p712';
+export * from '../test/generated-wrappers/fixin_protocol_fees';
 export * from '../test/generated-wrappers/fixin_reentrancy_guard';
 export * from '../test/generated-wrappers/flash_wallet';
 export * from '../test/generated-wrappers/full_migration';
@@ -27,6 +29,7 @@ export * from '../test/generated-wrappers/i_meta_transactions_feature';
 export * from '../test/generated-wrappers/i_ownable_feature';
 export * from '../test/generated-wrappers/i_signature_validator_feature';
 export * from '../test/generated-wrappers/i_simple_function_registry_feature';
+export * from '../test/generated-wrappers/i_staking';
 export * from '../test/generated-wrappers/i_test_simple_function_registry_feature';
 export * from '../test/generated-wrappers/i_token_spender_feature';
 export * from '../test/generated-wrappers/i_transform_erc20_feature';
@@ -90,8 +93,10 @@ export * from '../test/generated-wrappers/test_meta_transactions_transform_erc20
 export * from '../test/generated-wrappers/test_migrator';
 export * from '../test/generated-wrappers/test_mint_token_erc20_transformer';
 export * from '../test/generated-wrappers/test_mintable_erc20_token';
+export * from '../test/generated-wrappers/test_protocol_fees';
 export * from '../test/generated-wrappers/test_simple_function_registry_feature_impl1';
 export * from '../test/generated-wrappers/test_simple_function_registry_feature_impl2';
+export * from '../test/generated-wrappers/test_staking';
 export * from '../test/generated-wrappers/test_token_spender';
 export * from '../test/generated-wrappers/test_token_spender_erc20_token';
 export * from '../test/generated-wrappers/test_transform_erc20';

contracts/zero-ex/tsconfig.json

@@ -31,9 +31,11 @@
         "test/generated-artifacts/AllowanceTarget.json",
         "test/generated-artifacts/BootstrapFeature.json",
         "test/generated-artifacts/BridgeAdapter.json",
+        "test/generated-artifacts/FeeCollector.json",
         "test/generated-artifacts/FillQuoteTransformer.json",
         "test/generated-artifacts/FixinCommon.json",
         "test/generated-artifacts/FixinEIP712.json",
+        "test/generated-artifacts/FixinProtocolFees.json",
         "test/generated-artifacts/FixinReentrancyGuard.json",
         "test/generated-artifacts/FlashWallet.json",
         "test/generated-artifacts/FullMigration.json",
@@ -51,6 +53,7 @@
         "test/generated-artifacts/IOwnableFeature.json",
         "test/generated-artifacts/ISignatureValidatorFeature.json",
         "test/generated-artifacts/ISimpleFunctionRegistryFeature.json",
+        "test/generated-artifacts/IStaking.json",
         "test/generated-artifacts/ITestSimpleFunctionRegistryFeature.json",
         "test/generated-artifacts/ITokenSpenderFeature.json",
         "test/generated-artifacts/ITransformERC20Feature.json",
@@ -114,8 +117,10 @@
         "test/generated-artifacts/TestMigrator.json",
         "test/generated-artifacts/TestMintTokenERC20Transformer.json",
         "test/generated-artifacts/TestMintableERC20Token.json",
+        "test/generated-artifacts/TestProtocolFees.json",
         "test/generated-artifacts/TestSimpleFunctionRegistryFeatureImpl1.json",
         "test/generated-artifacts/TestSimpleFunctionRegistryFeatureImpl2.json",
+        "test/generated-artifacts/TestStaking.json",
         "test/generated-artifacts/TestTokenSpender.json",
         "test/generated-artifacts/TestTokenSpenderERC20Token.json",
         "test/generated-artifacts/TestTransformERC20.json",