From 30fbbc358f0ae4807ffd5b5441dbd7c68f8e8e32 Mon Sep 17 00:00:00 2001 From: ncitron Date: Wed, 9 Jun 2021 15:13:37 -0400 Subject: [PATCH 1/8] low hanging audit changes --- contracts/interfaces/IClaimAdapter.sol | 35 +++++++++++++++++-- .../mocks/integrations/ClaimAdapterMock.sol | 11 +++--- .../integration/claim/CompClaimAdapter.sol | 7 ++-- contracts/protocol/modules/ClaimModule.sol | 22 +++++++----- test/protocol/modules/claimModule.spec.ts | 7 ++-- 5 files changed, 61 insertions(+), 21 deletions(-) diff --git a/contracts/interfaces/IClaimAdapter.sol b/contracts/interfaces/IClaimAdapter.sol index 892477c1a..dc5dbac0a 100644 --- a/contracts/interfaces/IClaimAdapter.sol +++ b/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); } diff --git a/contracts/mocks/integrations/ClaimAdapterMock.sol b/contracts/mocks/integrations/ClaimAdapterMock.sol index f451913fa..ff0383a51 100644 --- a/contracts/mocks/integrations/ClaimAdapterMock.sol +++ b/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,7 +54,7 @@ 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; @@ -59,10 +62,10 @@ contract ClaimAdapterMock is ERC20 { 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; } } diff --git a/contracts/protocol/integration/claim/CompClaimAdapter.sol b/contracts/protocol/integration/claim/CompClaimAdapter.sol index 14460cd2d..e00118ae4 100644 --- a/contracts/protocol/integration/claim/CompClaimAdapter.sol +++ b/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)); } /** diff --git a/contracts/protocol/modules/ClaimModule.sol b/contracts/protocol/modules/ClaimModule.sol index 88a3d67aa..6dc831537 100644 --- a/contracts/protocol/modules/ClaimModule.sol +++ b/contracts/protocol/modules/ClaimModule.sol @@ -52,10 +52,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 ============ */ /** @@ -133,10 +138,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 @@ -330,8 +335,7 @@ contract ClaimModule is ModuleBase { returns (uint256) { IClaimAdapter adapter = _getAndValidateIntegrationAdapter(claimSettings[_setToken][_rewardPool], _integrationName); - uint256 rewards = adapter.getRewards(address(_setToken), _rewardPool); - return rewards; + return adapter.getRewardsAmount(_setToken, _rewardPool); } /* ============ Internal Functions ============ */ @@ -348,14 +352,14 @@ contract ClaimModule is ModuleBase { require(isRewardPool(_setToken, _rewardPool), "RewardPool not present"); IClaimAdapter adapter = _getAndValidateIntegrationAdapter(claimSettings[_setToken][_rewardPool], _integrationName); - uint256 rewards = adapter.getRewards(address(_setToken), _rewardPool); + uint256 rewards = adapter.getRewardsAmount(_setToken, _rewardPool); ( address callTarget, uint256 callValue, bytes memory callByteData ) = adapter.getClaimCallData( - address(_setToken), + _setToken, _rewardPool ); @@ -428,7 +432,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 diff --git a/test/protocol/modules/claimModule.spec.ts b/test/protocol/modules/claimModule.spec.ts index b38b761c7..12948706c 100644 --- a/test/protocol/modules/claimModule.spec.ts +++ b/test/protocol/modules/claimModule.spec.ts @@ -229,6 +229,7 @@ describe("ClaimModule", () => { let subjectSetToken: Address; let subjectCaller: Account; + let subjectAnyoneClaim: boolean; before(async () => { isInitialized = true; @@ -250,18 +251,20 @@ describe("ClaimModule", () => { }); async function subject(): Promise { - return claimModule.connect(subjectCaller.wallet).updateAnyoneClaim(subjectSetToken); + return claimModule.connect(subjectCaller.wallet).updateAnyoneClaim(subjectSetToken, subjectAnyoneClaim); } - it("should flip the anyoneClaim indicator", async () => { + it("should change the anyoneClaim indicator", async () => { const anyoneClaimBefore = await claimModule.anyoneClaim(subjectSetToken); expect(anyoneClaimBefore).to.eq(true); + subjectAnyoneClaim = false; await subject(); const anyoneClaim = await claimModule.anyoneClaim(subjectSetToken); expect(anyoneClaim).to.eq(false); + subjectAnyoneClaim = true; await subject(); const anyoneClaimAfter = await claimModule.anyoneClaim(subjectSetToken); From 2296353be7aad6a8ed615e45fd5b2d4493a37c50 Mon Sep 17 00:00:00 2001 From: ncitron Date: Wed, 9 Jun 2021 15:36:17 -0400 Subject: [PATCH 2/8] fix tests --- test/integration/claim/compClaimAdapter.spec.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/integration/claim/compClaimAdapter.spec.ts b/test/integration/claim/compClaimAdapter.spec.ts index 72c660f7c..fe9a5faf6 100644 --- a/test/integration/claim/compClaimAdapter.spec.ts +++ b/test/integration/claim/compClaimAdapter.spec.ts @@ -67,7 +67,7 @@ describe("CompClaimAdapter", function() { }); }); - describe("#getRewards", async function() { + describe("#getRewardsAmount", async function() { const rewards: BigNumber = ether(1); before(async function() { @@ -75,7 +75,7 @@ describe("CompClaimAdapter", function() { }); function subject(): Promise { - 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() { From 15c3be744e8a3ce34c1be8bd8cc83b1a490fad23 Mon Sep 17 00:00:00 2001 From: ncitron Date: Thu, 10 Jun 2021 10:58:54 -0400 Subject: [PATCH 3/8] audit changes --- contracts/protocol/modules/ClaimModule.sol | 58 +++++++++++++++++----- 1 file changed, 45 insertions(+), 13 deletions(-) diff --git a/contracts/protocol/modules/ClaimModule.sol b/contracts/protocol/modules/ClaimModule.sol index 6dc831537..069756cc4 100644 --- a/contracts/protocol/modules/ClaimModule.sol +++ b/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"; @@ -76,11 +78,14 @@ 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; + 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; + mapping(ISetToken => mapping(address => mapping(address => bool))) public claimSettingsStatus; + /* ============ Constructor ============ */ @@ -259,8 +264,21 @@ contract ClaimModule is ModuleBase { 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++) { + delete rewardPoolList[ISetToken(msg.sender)][i]; + } delete rewardPoolList[ISetToken(msg.sender)]; } @@ -282,7 +300,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]; } /** @@ -314,7 +332,7 @@ contract ClaimModule is ModuleBase { returns (bool) { address adapter = getAndValidateAdapter(_integrationName); - return claimSettings[_setToken][_rewardPool].contains(adapter); + return claimSettingsStatus[_setToken][_rewardPool][adapter]; } /** @@ -334,7 +352,7 @@ contract ClaimModule is ModuleBase { view returns (uint256) { - IClaimAdapter adapter = _getAndValidateIntegrationAdapter(claimSettings[_setToken][_rewardPool], _integrationName); + IClaimAdapter adapter = _getAndValidateIntegrationAdapter(_setToken, _rewardPool, _integrationName); return adapter.getRewardsAmount(_setToken, _rewardPool); } @@ -350,9 +368,10 @@ 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.getRewardsAmount(_setToken, _rewardPool); + IERC20 rewardsToken = IERC20(adapter.getTokenAddress(_rewardPool)); + uint256 initRewardsBalance = rewardsToken.balanceOf(address(_setToken)); ( address callTarget, @@ -365,17 +384,22 @@ contract ClaimModule is ModuleBase { _setToken.invoke(callTarget, callValue, callByteData); + uint256 finalRewardsBalance = rewardsToken.balanceOf(address(_setToken)); + uint256 rewards = finalRewardsBalance.sub(initRewardsBalance); + emit RewardClaimed(_setToken, _rewardPool, adapter, rewards); } /** * 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 @@ -383,7 +407,7 @@ contract ClaimModule is ModuleBase { 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); } @@ -399,11 +423,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) { rewardPoolList[_setToken].push(_rewardPool); + rewardPoolStatus[_setToken][_rewardPool] = true; } } @@ -443,11 +469,17 @@ contract ClaimModule is ModuleBase { address adapter = getAndValidateAdapter(_integrationName); address[] memory _rewardPoolClaimSettings = claimSettings[_setToken][_rewardPool]; - require(_rewardPoolClaimSettings.contains(adapter), "Integration must be added"); + require(claimSettingsStatus[_setToken][_rewardPool][adapter], "Integration must be added"); claimSettings[_setToken][_rewardPool] = _rewardPoolClaimSettings.remove(adapter); + claimSettingsStatus[_setToken][_rewardPool][adapter] = false; if (claimSettings[_setToken][_rewardPool].length == 0) { rewardPoolList[_setToken] = rewardPoolList[_setToken].remove(_rewardPool); + rewardPoolStatus[_setToken][_rewardPool] = false; + + for (uint256 i = 0; i < claimSettings[_setToken][_rewardPool].length; i++) { + delete claimSettings[_setToken][_rewardPool][i]; + } delete claimSettings[_setToken][_rewardPool]; } } From 65aeb0d785bd3e47485f446bddc17fbbf3dc683c Mon Sep 17 00:00:00 2001 From: ncitron Date: Thu, 10 Jun 2021 12:07:16 -0400 Subject: [PATCH 4/8] fix tests --- contracts/mocks/external/ComptrollerMock.sol | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/contracts/mocks/external/ComptrollerMock.sol b/contracts/mocks/external/ComptrollerMock.sol index bc9f17874..18c2326e3 100644 --- a/contracts/mocks/external/ComptrollerMock.sol +++ b/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; + } } \ No newline at end of file From 3c905dcfc448ba45dc660246d682ad4cd2d81893 Mon Sep 17 00:00:00 2001 From: ncitron Date: Mon, 14 Jun 2021 10:31:09 -0400 Subject: [PATCH 5/8] improve tests --- contracts/protocol/modules/ClaimModule.sol | 3 + test/protocol/modules/claimModule.spec.ts | 64 ++++++++++++++++++++-- 2 files changed, 63 insertions(+), 4 deletions(-) diff --git a/contracts/protocol/modules/ClaimModule.sol b/contracts/protocol/modules/ClaimModule.sol index 069756cc4..837d69f6e 100644 --- a/contracts/protocol/modules/ClaimModule.sol +++ b/contracts/protocol/modules/ClaimModule.sol @@ -277,6 +277,9 @@ contract ClaimModule is ModuleBase { } 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)]; diff --git a/test/protocol/modules/claimModule.spec.ts b/test/protocol/modules/claimModule.spec.ts index 12948706c..da08ee282 100644 --- a/test/protocol/modules/claimModule.spec.ts +++ b/test/protocol/modules/claimModule.spec.ts @@ -209,15 +209,23 @@ describe("ClaimModule", () => { it("should properly remove the module and settings", async () => { const rewardPoolsBefore = await claimModule.getRewardPools(setToken.address); const rewardPoolClaimsBefore = await claimModule.getRewardPoolClaims(setToken.address, rewardPool); + const isPoolAddedBefore = await claimModule.rewardPoolStatus(setToken.address, rewardPool); + const isAdapterAddedBefore = await claimModule.claimSettingsStatus(setToken.address, rewardPool, claimAdapterMock.address); expect(rewardPoolsBefore.length).to.eq(1); expect(rewardPoolClaimsBefore.length).to.eq(1); + expect(isPoolAddedBefore).to.be.true; + expect(isAdapterAddedBefore).to.be.true; await subject(); const rewardPools = await claimModule.getRewardPools(setToken.address); const rewardPoolClaims = await claimModule.getRewardPoolClaims(setToken.address, rewardPool); + const isPoolAdded = await claimModule.rewardPoolStatus(setToken.address, rewardPool); + const isAdapterAdded = await claimModule.claimSettingsStatus(setToken.address, rewardPool, claimAdapterMock.address); expect(rewardPools.length).to.eq(0); expect(rewardPoolClaims.length).to.eq(0); + expect(isPoolAdded).to.be.false; + expect(isAdapterAdded).to.be.false; const isModuleEnabled = await setToken.isInitializedModule(subjectModule); expect(isModuleEnabled).to.eq(false); }); @@ -348,13 +356,21 @@ describe("ClaimModule", () => { } it("should add the rewardPools to the rewardPoolList", async () => { + const isFirstAddedBefore = await claimModule.rewardPoolStatus(subjectSetToken, subjectRewardPools[0]); + const isSecondAddedBefore = await claimModule.rewardPoolStatus(subjectSetToken, subjectRewardPools[2]); expect((await claimModule.getRewardPools(subjectSetToken)).length).to.eq(1); + expect(isFirstAddedBefore).to.be.false; + expect(isSecondAddedBefore).to.be.false; await subject(); const rewardPools = await claimModule.getRewardPools(subjectSetToken); + const isFirstAdded = await claimModule.rewardPoolStatus(subjectSetToken, subjectRewardPools[0]); + const isSecondAdded = await claimModule.rewardPoolStatus(subjectSetToken, subjectRewardPools[2]); expect(rewardPools[1]).to.eq(subjectRewardPools[0]); expect(rewardPools[2]).to.eq(subjectRewardPools[2]); + expect(isFirstAdded).to.be.true; + expect(isSecondAdded).to.be.true; }); it("should add all new integrations for the rewardPools", async () => { @@ -362,9 +378,15 @@ describe("ClaimModule", () => { const rewardPoolOneClaims = await claimModule.getRewardPoolClaims(setToken.address, subjectRewardPools[0]); const rewardPoolTwoClaims = await claimModule.getRewardPoolClaims(setToken.address, subjectRewardPools[2]); + const isFirstIntegrationAddedPool1 = await claimModule.claimSettingsStatus(setToken.address, subjectRewardPools[0], claimAdapterMock.address); + const isSecondIntegrationAddedPool1 = await claimModule.claimSettingsStatus(setToken.address, subjectRewardPools[1], claimAdapterMock2.address); + const isIntegrationAddedPool2 = await claimModule.claimSettingsStatus(setToken.address, subjectRewardPools[0], claimAdapterMock.address); expect(rewardPoolOneClaims[0]).to.eq(claimAdapterMock.address); expect(rewardPoolOneClaims[1]).to.eq(claimAdapterMock2.address); expect(rewardPoolTwoClaims[0]).to.eq(claimAdapterMock.address); + expect(isFirstIntegrationAddedPool1).to.be.true; + expect(isSecondIntegrationAddedPool1).to.be.true; + expect(isIntegrationAddedPool2).to.be.true; }); describe("when passed arrays are different length", async () => { @@ -474,23 +496,28 @@ describe("ClaimModule", () => { return claimModule.connect(subjectCaller.wallet).addClaim(subjectSetToken, subjectRewardPool, subjectIntegration); } - it("should add the rewardPool to the rewardPoolList", async () => { + it("should add the rewardPool to the rewardPoolList and rewardPoolStatus", async () => { expect(await claimModule.isRewardPool(subjectSetToken, subjectRewardPool)).to.be.false; await subject(); expect(await claimModule.isRewardPool(subjectSetToken, subjectRewardPool)).to.be.true; + expect(await claimModule.rewardPoolList(subjectSetToken, 1)).to.eq(subjectRewardPool); }); it("should add new integration for the rewardPool", async () => { const rewardPoolClaimsBefore = await claimModule.getRewardPoolClaims(setToken.address, subjectRewardPool); + const isIntegrationAddedBefore = await claimModule.claimSettingsStatus(setToken.address, subjectRewardPool, claimAdapterMock2.address); expect(rewardPoolClaimsBefore.length).to.eq(0); + expect(isIntegrationAddedBefore).to.be.false; await subject(); const rewardPoolClaims = await claimModule.getRewardPoolClaims(setToken.address, subjectRewardPool); + const isIntegrationAdded = await claimModule.claimSettingsStatus(setToken.address, subjectRewardPool, claimAdapterMock2.address); expect(rewardPoolClaims.length).to.eq(1); expect(rewardPoolClaims[0]).to.eq(claimAdapterMock2.address); + expect(isIntegrationAdded).to.be.true; }); describe("when new claim is being added to existing rewardPool", async () => { @@ -500,13 +527,17 @@ describe("ClaimModule", () => { it("should add new integration for the rewardPool", async () => { const rewardPoolClaimsBefore = await claimModule.getRewardPoolClaims(setToken.address, subjectRewardPool); + const isIntegrationAddedBefore = await claimModule.claimSettingsStatus(setToken.address, subjectRewardPool, claimAdapterMock2.address); expect(rewardPoolClaimsBefore.length).to.eq(1); + expect(isIntegrationAddedBefore).to.be.false; await subject(); + const isIntegrationAdded = await claimModule.claimSettingsStatus(setToken.address, subjectRewardPool, claimAdapterMock2.address); const rewardPoolClaims = await claimModule.getRewardPoolClaims(setToken.address, subjectRewardPool); expect(rewardPoolClaims.length).to.eq(2); expect(rewardPoolClaims[1]).to.eq(claimAdapterMock2.address); + expect(isIntegrationAdded).to.be.true; }); it("should not add the rewardPool again", async () => { @@ -609,18 +640,35 @@ describe("ClaimModule", () => { it("should remove the adapter associated to the reward pool", async () => { const rewardPoolOneClaimsBefore = await claimModule.getRewardPoolClaims(setToken.address, subjectRewardPools[0]); const rewardPoolTwoClaimsBefore = await claimModule.getRewardPoolClaims(setToken.address, subjectRewardPools[1]); + const isRewardPoolOneAdapterOneBefore = await claimModule.claimSettingsStatus( + setToken.address, + subjectRewardPools[0], + claimAdapterMock.address + ); + const isRewardPoolTwoAdapterOneBefore = await claimModule.claimSettingsStatus( + setToken.address, + subjectRewardPools[0], + claimAdapterMock.address + ); expect(rewardPoolOneClaimsBefore.length).to.eq(1); expect(rewardPoolTwoClaimsBefore.length).to.eq(1); + expect(isRewardPoolOneAdapterOneBefore).to.be.true; + expect(isRewardPoolTwoAdapterOneBefore).to.be.true; await subject(); const rewardPoolOneClaims = await claimModule.getRewardPoolClaims(setToken.address, subjectRewardPools[0]); const rewardPoolTwoClaims = await claimModule.getRewardPoolClaims(setToken.address, subjectRewardPools[1]); + const isRewardPoolOneAdapterOne = await claimModule.claimSettingsStatus(setToken.address, subjectRewardPools[0], claimAdapterMock.address); + const isRewardPoolTwoAdapterOne = await claimModule.claimSettingsStatus(setToken.address, subjectRewardPools[0], claimAdapterMock.address); expect(rewardPoolOneClaims.length).to.eq(0); expect(rewardPoolTwoClaims.length).to.eq(0); + expect(isRewardPoolOneAdapterOne).to.be.false; + expect(isRewardPoolTwoAdapterOne).to.be.false; + }); - it("should remove the rewardPool from the rewardPoolList", async () => { + it("should remove the rewardPool from the rewardPoolStatus", async () => { expect(await claimModule.isRewardPool(subjectSetToken, subjectRewardPools[0])).to.be.true; expect(await claimModule.isRewardPool(subjectSetToken, subjectRewardPools[1])).to.be.true; @@ -739,15 +787,19 @@ describe("ClaimModule", () => { it("should remove the adapter associated to the reward pool", async () => { const rewardPoolClaimsBefore = await claimModule.getRewardPoolClaims(setToken.address, subjectRewardPool); + const isAdapterAddedBefore = await claimModule.claimSettingsStatus(setToken.address, subjectRewardPool, claimAdapterMock.address); expect(rewardPoolClaimsBefore.length).to.eq(1); + expect(isAdapterAddedBefore).to.be.true; await subject(); const rewardPoolClaims = await claimModule.getRewardPoolClaims(setToken.address, subjectRewardPool); + const isAdapterAdded = await claimModule.claimSettingsStatus(setToken.address, subjectRewardPool, claimAdapterMock.address); expect(rewardPoolClaims.length).to.eq(0); + expect(isAdapterAdded).to.be.false; }); - it("should remove the rewardPool from the rewardPoolList", async () => { + it("should remove the rewardPool from the rewardPoolStatus", async () => { expect(await claimModule.isRewardPool(subjectSetToken, subjectRewardPool)).to.be.true; await subject(); @@ -762,16 +814,20 @@ describe("ClaimModule", () => { it("should remove the adapter associated to the reward pool", async () => { const rewardPoolClaimsBefore = await claimModule.getRewardPoolClaims(setToken.address, subjectRewardPool); + const isAdapterAddedBefore = await claimModule.claimSettingsStatus(setToken.address, subjectRewardPool, claimAdapterMock.address); expect(rewardPoolClaimsBefore.length).to.eq(2); + expect(isAdapterAddedBefore).to.be.true; await subject(); const rewardPoolClaims = await claimModule.getRewardPoolClaims(setToken.address, subjectRewardPool); + const isAdapterAdded = await claimModule.claimSettingsStatus(setToken.address, subjectRewardPool, claimAdapterMock.address); expect(rewardPoolClaims.length).to.eq(1); expect(rewardPoolClaims[0]).to.eq(claimAdapterMock2.address); + expect(isAdapterAdded).to.be.false; }); - it("should not remove the rewardPool from the rewardPoolList", async () => { + it("should not remove the rewardPool from the rewardPoolStatus", async () => { expect(await claimModule.isRewardPool(subjectSetToken, subjectRewardPool)).to.be.true; await subject(); From 883c557106e7bd0756439b77bf1641db58347430 Mon Sep 17 00:00:00 2001 From: ncitron Date: Mon, 14 Jun 2021 10:45:50 -0400 Subject: [PATCH 6/8] clean up contracts --- contracts/protocol/modules/ClaimModule.sol | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/contracts/protocol/modules/ClaimModule.sol b/contracts/protocol/modules/ClaimModule.sol index 837d69f6e..0ccd9b112 100644 --- a/contracts/protocol/modules/ClaimModule.sol +++ b/contracts/protocol/modules/ClaimModule.sol @@ -262,6 +262,7 @@ 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++) { @@ -430,7 +431,7 @@ contract ClaimModule is ModuleBase { _rewardPoolClaimSettings.push(adapter); claimSettingsStatus[_setToken][_rewardPool][adapter] = true; - if (_rewardPoolClaimSettings.length == 1) { + if (!rewardPoolStatus[_setToken][_rewardPool]) { rewardPoolList[_setToken].push(_rewardPool); rewardPoolStatus[_setToken][_rewardPool] = true; } @@ -479,11 +480,6 @@ contract ClaimModule is ModuleBase { if (claimSettings[_setToken][_rewardPool].length == 0) { rewardPoolList[_setToken] = rewardPoolList[_setToken].remove(_rewardPool); rewardPoolStatus[_setToken][_rewardPool] = false; - - for (uint256 i = 0; i < claimSettings[_setToken][_rewardPool].length; i++) { - delete claimSettings[_setToken][_rewardPool][i]; - } - delete claimSettings[_setToken][_rewardPool]; } } From 3e2236a48406de03364317cb249d81d31446e96a Mon Sep 17 00:00:00 2001 From: ncitron Date: Mon, 14 Jun 2021 11:11:13 -0400 Subject: [PATCH 7/8] use removeStorage --- contracts/protocol/modules/ClaimModule.sol | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/contracts/protocol/modules/ClaimModule.sol b/contracts/protocol/modules/ClaimModule.sol index 0ccd9b112..6c2780e3f 100644 --- a/contracts/protocol/modules/ClaimModule.sol +++ b/contracts/protocol/modules/ClaimModule.sol @@ -471,14 +471,13 @@ 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(claimSettingsStatus[_setToken][_rewardPool][adapter], "Integration must be added"); - claimSettings[_setToken][_rewardPool] = _rewardPoolClaimSettings.remove(adapter); + claimSettings[_setToken][_rewardPool].removeStorage(adapter); claimSettingsStatus[_setToken][_rewardPool][adapter] = false; if (claimSettings[_setToken][_rewardPool].length == 0) { - rewardPoolList[_setToken] = rewardPoolList[_setToken].remove(_rewardPool); + rewardPoolList[_setToken].removeStorage(_rewardPool); rewardPoolStatus[_setToken][_rewardPool] = false; } } From 6266aea9b0ee876b863ff4877ebf4aeae9233aab Mon Sep 17 00:00:00 2001 From: ncitron Date: Fri, 2 Jul 2021 12:29:04 -0400 Subject: [PATCH 8/8] review changes --- contracts/protocol/modules/ClaimModule.sol | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/contracts/protocol/modules/ClaimModule.sol b/contracts/protocol/modules/ClaimModule.sol index 6c2780e3f..67dbf5134 100644 --- a/contracts/protocol/modules/ClaimModule.sol +++ b/contracts/protocol/modules/ClaimModule.sol @@ -80,10 +80,12 @@ contract ClaimModule is ModuleBase { // 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; // 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; @@ -389,9 +391,8 @@ contract ClaimModule is ModuleBase { _setToken.invoke(callTarget, callValue, callByteData); uint256 finalRewardsBalance = rewardsToken.balanceOf(address(_setToken)); - uint256 rewards = finalRewardsBalance.sub(initRewardsBalance); - emit RewardClaimed(_setToken, _rewardPool, adapter, rewards); + emit RewardClaimed(_setToken, _rewardPool, adapter, finalRewardsBalance.sub(initRewardsBalance)); } /**