ClaimModule Audit fixes (#97)

* low hanging audit changes

* fix tests

* audit changes

* fix tests

* improve tests

* clean up contracts

* use removeStorage

* review changes

Author: ncitron

contracts/interfaces/IClaimAdapter.sol

@@ -16,6 +16,10 @@
     SPDX-License-Identifier: Apache License, Version 2.0
 */
 
+import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+
+import { ISetToken } from "../interfaces/ISetToken.sol";
+
 pragma solidity 0.6.10;
 
 /**
@@ -25,12 +29,37 @@ pragma solidity 0.6.10;
  */
 interface IClaimAdapter {
 
+    /**
+     * Generates the calldata for claiming tokens from the rewars pool
+     *
+     * @param _setToken     the set token that is owed the tokens
+     * @param _rewardPool   the rewards pool to claim from  
+     *
+     * @return _subject     the rewards pool to call
+     * @return _value       the amount of ether to send in the call
+     * @return _calldata    the calldata to use
+     */
     function getClaimCallData(
-        address _holder,
+        ISetToken _setToken,
         address _rewardPool
     ) external view returns(address _subject, uint256 _value, bytes memory _calldata);
 
-    function getRewards(address _holder, address _rewardPool) external view returns(uint256);
+    /**
+     * Gets the amount of unclaimed rewards
+     *
+     * @param _setToken     the set token that is owed the tokens
+     * @param _rewardPool   the rewards pool to check
+     *
+     * @return uint256      the amount of unclaimed rewards
+     */
+    function getRewardsAmount(ISetToken _setToken, address _rewardPool) external view returns(uint256);
 
-    function getTokenAddress(address _rewardPool) external view returns(address);
+    /**
+     * Gets the rewards token
+     *
+     * @param _rewardPool   the rewards pool to check
+     *
+     * @return IERC20       the reward token
+     */
+    function getTokenAddress(address _rewardPool) external view returns(IERC20);
 }

contracts/mocks/external/ComptrollerMock.sol

@@ -63,4 +63,8 @@ contract ComptrollerMock {
     function setCompAccrued(address _holder, uint _compAmount) external {
         compAccrued[_holder] = _compAmount;
     }
+
+    function getCompAddress() external view returns (address) {
+        return comp;
+    }
 }

contracts/mocks/integrations/ClaimAdapterMock.sol

@@ -19,6 +19,9 @@
 pragma solidity 0.6.10;
 
 import { ERC20 } from "@openzeppelin/contracts/token/ERC20/ERC20.sol";
+import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+
+import { ISetToken } from "../../interfaces/ISetToken.sol";
 
 contract ClaimAdapterMock is ERC20  {
 
@@ -40,7 +43,7 @@ contract ClaimAdapterMock is ERC20  {
     }
 
     function getClaimCallData(
-        address _holder,
+        ISetToken _holder,
         address _rewardPool
     ) external view returns(address _subject, uint256 _value, bytes memory _calldata) {
         // Quell compiler warnings about unused vars
@@ -51,18 +54,18 @@ contract ClaimAdapterMock is ERC20  {
         return (address(this), 0, callData);
     }
 
-    function getRewards(address _holder, address _rewardPool) external view returns(uint256) {
+    function getRewardsAmount(ISetToken _holder, address _rewardPool) external view returns(uint256) {
         // Quell compiler warnings about unused vars
         _holder;
         _rewardPool;
 
         return rewards;
     }
 
-    function getTokenAddress(address _rewardPool) external view returns(address) {
+    function getTokenAddress(address _rewardPool) external view returns(IERC20) {
         // Quell compiler warnings about unused vars
         _rewardPool;
 
-        return address(this);
+        return this;
     }
 }

contracts/protocol/integration/claim/CompClaimAdapter.sol

@@ -20,6 +20,7 @@ pragma solidity 0.6.10;
 pragma experimental "ABIEncoderV2";
 
 import { IComptroller } from "../../../interfaces/external/IComptroller.sol";
+import { ISetToken } from "../../../interfaces/ISetToken.sol";
 
 /**
  * @title CompClaimAdapter
@@ -58,7 +59,7 @@ contract CompClaimAdapter {
      * @return uint256      Unused, since it claims total claimable balance
      * @return bytes        Claim calldata
      */
-    function getClaimCallData(address _setToken, address /* _rewardPool */) external view returns (address, uint256, bytes memory) {
+    function getClaimCallData(ISetToken _setToken, address /* _rewardPool */) external view returns (address, uint256, bytes memory) {
         bytes memory callData = abi.encodeWithSignature("claimComp(address)", _setToken);
 
         return (address(comptroller), 0, callData);
@@ -69,8 +70,8 @@ contract CompClaimAdapter {
      *
      * @return uint256      Claimable COMP balance
      */
-    function getRewards(address _setToken, address /* _rewardPool */) external view returns(uint256) {
-        return comptroller.compAccrued(_setToken);
+    function getRewardsAmount(ISetToken _setToken, address /* _rewardPool */) external view returns(uint256) {
+        return comptroller.compAccrued(address(_setToken));
     }
 
     /**

contracts/protocol/modules/ClaimModule.sol

@@ -15,6 +15,8 @@
 pragma solidity 0.6.10;
 pragma experimental "ABIEncoderV2";
 
+import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+
 import { AddressArrayUtils } from "../../lib/AddressArrayUtils.sol";
 import { IClaimAdapter } from "../../interfaces/IClaimAdapter.sol";
 import { IController } from "../../interfaces/IController.sol";
@@ -52,10 +54,15 @@ contract ClaimModule is ModuleBase {
     event RewardClaimed(
         ISetToken indexed _setToken,
         address indexed _rewardPool,
-        IClaimAdapter _adapter,
+        IClaimAdapter indexed _adapter,
         uint256 _amount
     );
 
+    event AnyoneClaimUpdated(
+        ISetToken indexed _setToken,
+        bool _anyoneClaim
+    );
+
     /* ============ Modifiers ============ */
 
     /**
@@ -71,11 +78,16 @@ contract ClaimModule is ModuleBase {
     // Indicates if any address can call claim or just the manager of the SetToken
     mapping(ISetToken => bool) public anyoneClaim;
 
-    // Array of rewardPool addresses to claim rewards for the SetToken
+    // Map and array of rewardPool addresses to claim rewards for the SetToken
     mapping(ISetToken => address[]) public rewardPoolList;
+    // Map from set tokens to rewards pool address to isAdded boolean. Used to check if a reward pool has been added in O(1) time
+    mapping(ISetToken => mapping(address => bool)) public rewardPoolStatus;
 
-    // Array of adapters associated to the rewardPool for the SetToken
+    // Map and array of adapters associated to the rewardPool for the SetToken
     mapping(ISetToken => mapping(address => address[])) public claimSettings;
+    // Map from set tokens to rewards pool address to claim adapters to isAdded boolean. Used to check if an adapter has been added in O(1) time
+    mapping(ISetToken => mapping(address => mapping(address => bool))) public claimSettingsStatus;
+
 
     /* ============ Constructor ============ */
 
@@ -133,10 +145,10 @@ contract ClaimModule is ModuleBase {
      *
      * @param _setToken             Address of SetToken
      */
-    function updateAnyoneClaim(ISetToken _setToken) external onlyManagerAndValidSet(_setToken) {
-        anyoneClaim[_setToken] = !anyoneClaim[_setToken];
+    function updateAnyoneClaim(ISetToken _setToken, bool _anyoneClaim) external onlyManagerAndValidSet(_setToken) {
+        anyoneClaim[_setToken] = _anyoneClaim;
+        emit AnyoneClaimUpdated(_setToken, _anyoneClaim);
     }
-
     /**
      * SET MANAGER ONLY. Adds a new claim integration for an existent rewardPool. If rewardPool doesn't have existing
      * claims then rewardPool is added to rewardPoolLiost. The claim integration is associated to an adapter that
@@ -252,10 +264,27 @@ contract ClaimModule is ModuleBase {
     function removeModule() external override {
         delete anyoneClaim[ISetToken(msg.sender)];
 
+        // explicitly delete all elements for gas refund
         address[] memory setTokenPoolList = rewardPoolList[ISetToken(msg.sender)];
         for (uint256 i = 0; i < setTokenPoolList.length; i++) {
+
+            address[] storage adapterList = claimSettings[ISetToken(msg.sender)][setTokenPoolList[i]];
+            for (uint256 j = 0; j < adapterList.length; j++) {
+
+                address toRemove = adapterList[j];
+                claimSettingsStatus[ISetToken(msg.sender)][setTokenPoolList[i]][toRemove] = false;
+
+                delete adapterList[j];
+            }
             delete claimSettings[ISetToken(msg.sender)][setTokenPoolList[i]];
         }
+        
+        for (uint256 i = 0; i < rewardPoolList[ISetToken(msg.sender)].length; i++) {
+            address toRemove = rewardPoolList[ISetToken(msg.sender)][i];
+            rewardPoolStatus[ISetToken(msg.sender)][toRemove] = false;
+            
+            delete rewardPoolList[ISetToken(msg.sender)][i];
+        }
         delete rewardPoolList[ISetToken(msg.sender)];
     }
 
@@ -277,7 +306,7 @@ contract ClaimModule is ModuleBase {
      * @return                      Boolean indicating if the rewardPool is in the list for claims.
      */
     function isRewardPool(ISetToken _setToken, address _rewardPool) public view returns (bool) {
-        return rewardPoolList[_setToken].contains(_rewardPool);
+        return rewardPoolStatus[_setToken][_rewardPool];
     }
 
     /**
@@ -309,7 +338,7 @@ contract ClaimModule is ModuleBase {
         returns (bool)
     {
         address adapter = getAndValidateAdapter(_integrationName);
-        return claimSettings[_setToken][_rewardPool].contains(adapter);
+        return claimSettingsStatus[_setToken][_rewardPool][adapter];
     }
 
     /**
@@ -329,9 +358,8 @@ contract ClaimModule is ModuleBase {
         view
         returns (uint256)
     {
-        IClaimAdapter adapter = _getAndValidateIntegrationAdapter(claimSettings[_setToken][_rewardPool], _integrationName);
-        uint256 rewards = adapter.getRewards(address(_setToken), _rewardPool);
-        return rewards;
+        IClaimAdapter adapter = _getAndValidateIntegrationAdapter(_setToken, _rewardPool, _integrationName);
+        return adapter.getRewardsAmount(_setToken, _rewardPool);
     }
 
     /* ============ Internal Functions ============ */
@@ -346,40 +374,45 @@ contract ClaimModule is ModuleBase {
      */
     function _claim(ISetToken _setToken, address _rewardPool, string calldata _integrationName) internal {
         require(isRewardPool(_setToken, _rewardPool), "RewardPool not present");
-        IClaimAdapter adapter = _getAndValidateIntegrationAdapter(claimSettings[_setToken][_rewardPool], _integrationName);
+        IClaimAdapter adapter = _getAndValidateIntegrationAdapter(_setToken, _rewardPool, _integrationName);
 
-        uint256 rewards = adapter.getRewards(address(_setToken), _rewardPool);
+        IERC20 rewardsToken = IERC20(adapter.getTokenAddress(_rewardPool));
+        uint256 initRewardsBalance = rewardsToken.balanceOf(address(_setToken));
 
         (
             address callTarget,
             uint256 callValue,
             bytes memory callByteData
         ) = adapter.getClaimCallData(
-            address(_setToken),
+            _setToken,
             _rewardPool
         );
 
         _setToken.invoke(callTarget, callValue, callByteData);
 
-        emit RewardClaimed(_setToken, _rewardPool, adapter, rewards);
+        uint256 finalRewardsBalance = rewardsToken.balanceOf(address(_setToken));
+
+        emit RewardClaimed(_setToken, _rewardPool, adapter, finalRewardsBalance.sub(initRewardsBalance));
     }
 
     /**
      * Gets the adapter and validate it is associated to the list of claim integration of a rewardPool.
      *
-     * @param _rewardPoolClaimSettings           List of claim integrations associated to a rewardPool
-     * @param _integrationName                   ID of claim module integration (mapping on integration registry)
+     * @param _setToken             Address of SetToken
+     * @param _rewardsPool          Sddress of rewards pool
+     * @param _integrationName      ID of claim module integration (mapping on integration registry)
      */
     function _getAndValidateIntegrationAdapter(
-        address[] memory _rewardPoolClaimSettings,
+        ISetToken _setToken,
+        address _rewardsPool,
         string calldata _integrationName
     )
         internal
         view
         returns (IClaimAdapter)
     {
         address adapter = getAndValidateAdapter(_integrationName);
-        require(_rewardPoolClaimSettings.contains(adapter), "Adapter integration not present");
+        require(claimSettingsStatus[_setToken][_rewardsPool][adapter], "Adapter integration not present");
         return IClaimAdapter(adapter);
     }
 
@@ -395,11 +428,13 @@ contract ClaimModule is ModuleBase {
         address adapter = getAndValidateAdapter(_integrationName);
         address[] storage _rewardPoolClaimSettings = claimSettings[_setToken][_rewardPool];
 
-        require(!_rewardPoolClaimSettings.contains(adapter), "Integration names must be unique");
+        require(!claimSettingsStatus[_setToken][_rewardPool][adapter], "Integration names must be unique");
         _rewardPoolClaimSettings.push(adapter);
+        claimSettingsStatus[_setToken][_rewardPool][adapter] = true;
 
-        if (_rewardPoolClaimSettings.length == 1) {
+        if (!rewardPoolStatus[_setToken][_rewardPool]) {
             rewardPoolList[_setToken].push(_rewardPool);
+            rewardPoolStatus[_setToken][_rewardPool] = true;
         }
     }
 
@@ -428,7 +463,7 @@ contract ClaimModule is ModuleBase {
     }
 
     /**
-     * Validates and store the adapter address used to claim rewards for the passed rewardPool. If no adapters
+     * Validates and stores the adapter address used to claim rewards for the passed rewardPool. If no adapters
      * left after removal then remove rewardPool from rewardPoolList and delete entry in claimSettings. 
      *
      * @param _setToken                 Address of SetToken
@@ -437,14 +472,14 @@ contract ClaimModule is ModuleBase {
      */
     function _removeClaim(ISetToken _setToken, address _rewardPool, string calldata _integrationName) internal {
         address adapter = getAndValidateAdapter(_integrationName);
-        address[] memory _rewardPoolClaimSettings = claimSettings[_setToken][_rewardPool];
 
-        require(_rewardPoolClaimSettings.contains(adapter), "Integration must be added");
-        claimSettings[_setToken][_rewardPool] = _rewardPoolClaimSettings.remove(adapter);
+        require(claimSettingsStatus[_setToken][_rewardPool][adapter], "Integration must be added");
+        claimSettings[_setToken][_rewardPool].removeStorage(adapter);
+        claimSettingsStatus[_setToken][_rewardPool][adapter] = false;
 
         if (claimSettings[_setToken][_rewardPool].length == 0) {
-            rewardPoolList[_setToken] = rewardPoolList[_setToken].remove(_rewardPool);
-            delete claimSettings[_setToken][_rewardPool];
+            rewardPoolList[_setToken].removeStorage(_rewardPool);
+            rewardPoolStatus[_setToken][_rewardPool] = false;
         }
     }
 

test/integration/claim/compClaimAdapter.spec.ts

@@ -67,15 +67,15 @@ describe("CompClaimAdapter", function() {
       });
     });
 
-    describe("#getRewards", async function() {
+    describe("#getRewardsAmount", async function() {
       const rewards: BigNumber = ether(1);
 
       before(async function() {
         await comptroller.mock.compAccrued.returns(rewards);
       });
 
       function subject(): Promise<BigNumber> {
-        return compClaimAdapter.connect(owner.wallet).getRewards(ADDRESS_ZERO, ADDRESS_ZERO);
+        return compClaimAdapter.connect(owner.wallet).getRewardsAmount(ADDRESS_ZERO, ADDRESS_ZERO);
       }
 
       it("should return rewards", async function() {