Skip to content

Commit 2741184

Browse files
ncitronrichardliang
ncitron
authored and
richardliang
committed
ClaimModule Audit fixes (#97)
* low hanging audit changes * fix tests * audit changes * fix tests * improve tests * clean up contracts * use removeStorage * review changes
1 parent 4f89e12 commit 2741184

File tree

7 files changed

+176
-45
lines changed

7 files changed

+176
-45
lines changed

‎contracts/interfaces/IClaimAdapter.sol

+32-3
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,10 @@
1616
SPDX-License-Identifier: Apache License, Version 2.0
1717
*/
1818

19+
import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
20+
21+
import { ISetToken } from "../interfaces/ISetToken.sol";
22+
1923
pragma solidity 0.6.10;
2024

2125
/**
@@ -25,12 +29,37 @@ pragma solidity 0.6.10;
2529
*/
2630
interface IClaimAdapter {
2731

32+
/**
33+
* Generates the calldata for claiming tokens from the rewars pool
34+
*
35+
* @param _setToken the set token that is owed the tokens
36+
* @param _rewardPool the rewards pool to claim from
37+
*
38+
* @return _subject the rewards pool to call
39+
* @return _value the amount of ether to send in the call
40+
* @return _calldata the calldata to use
41+
*/
2842
function getClaimCallData(
29-
address _holder,
43+
ISetToken _setToken,
3044
address _rewardPool
3145
) external view returns(address _subject, uint256 _value, bytes memory _calldata);
3246

33-
function getRewards(address _holder, address _rewardPool) external view returns(uint256);
47+
/**
48+
* Gets the amount of unclaimed rewards
49+
*
50+
* @param _setToken the set token that is owed the tokens
51+
* @param _rewardPool the rewards pool to check
52+
*
53+
* @return uint256 the amount of unclaimed rewards
54+
*/
55+
function getRewardsAmount(ISetToken _setToken, address _rewardPool) external view returns(uint256);
3456

35-
function getTokenAddress(address _rewardPool) external view returns(address);
57+
/**
58+
* Gets the rewards token
59+
*
60+
* @param _rewardPool the rewards pool to check
61+
*
62+
* @return IERC20 the reward token
63+
*/
64+
function getTokenAddress(address _rewardPool) external view returns(IERC20);
3665
}

‎contracts/mocks/external/ComptrollerMock.sol

+4
Original file line numberDiff line numberDiff line change
@@ -63,4 +63,8 @@ contract ComptrollerMock {
6363
function setCompAccrued(address _holder, uint _compAmount) external {
6464
compAccrued[_holder] = _compAmount;
6565
}
66+
67+
function getCompAddress() external view returns (address) {
68+
return comp;
69+
}
6670
}

‎contracts/mocks/integrations/ClaimAdapterMock.sol

+7-4
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,9 @@
1919
pragma solidity 0.6.10;
2020

2121
import { ERC20 } from "@openzeppelin/contracts/token/ERC20/ERC20.sol";
22+
import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
23+
24+
import { ISetToken } from "../../interfaces/ISetToken.sol";
2225

2326
contract ClaimAdapterMock is ERC20 {
2427

@@ -40,7 +43,7 @@ contract ClaimAdapterMock is ERC20 {
4043
}
4144

4245
function getClaimCallData(
43-
address _holder,
46+
ISetToken _holder,
4447
address _rewardPool
4548
) external view returns(address _subject, uint256 _value, bytes memory _calldata) {
4649
// Quell compiler warnings about unused vars
@@ -51,18 +54,18 @@ contract ClaimAdapterMock is ERC20 {
5154
return (address(this), 0, callData);
5255
}
5356

54-
function getRewards(address _holder, address _rewardPool) external view returns(uint256) {
57+
function getRewardsAmount(ISetToken _holder, address _rewardPool) external view returns(uint256) {
5558
// Quell compiler warnings about unused vars
5659
_holder;
5760
_rewardPool;
5861

5962
return rewards;
6063
}
6164

62-
function getTokenAddress(address _rewardPool) external view returns(address) {
65+
function getTokenAddress(address _rewardPool) external view returns(IERC20) {
6366
// Quell compiler warnings about unused vars
6467
_rewardPool;
6568

66-
return address(this);
69+
return this;
6770
}
6871
}

‎contracts/protocol/integration/claim/CompClaimAdapter.sol

+4-3
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ pragma solidity 0.6.10;
2020
pragma experimental "ABIEncoderV2";
2121

2222
import { IComptroller } from "../../../interfaces/external/IComptroller.sol";
23+
import { ISetToken } from "../../../interfaces/ISetToken.sol";
2324

2425
/**
2526
* @title CompClaimAdapter
@@ -58,7 +59,7 @@ contract CompClaimAdapter {
5859
* @return uint256 Unused, since it claims total claimable balance
5960
* @return bytes Claim calldata
6061
*/
61-
function getClaimCallData(address _setToken, address /* _rewardPool */) external view returns (address, uint256, bytes memory) {
62+
function getClaimCallData(ISetToken _setToken, address /* _rewardPool */) external view returns (address, uint256, bytes memory) {
6263
bytes memory callData = abi.encodeWithSignature("claimComp(address)", _setToken);
6364

6465
return (address(comptroller), 0, callData);
@@ -69,8 +70,8 @@ contract CompClaimAdapter {
6970
*
7071
* @return uint256 Claimable COMP balance
7172
*/
72-
function getRewards(address _setToken, address /* _rewardPool */) external view returns(uint256) {
73-
return comptroller.compAccrued(_setToken);
73+
function getRewardsAmount(ISetToken _setToken, address /* _rewardPool */) external view returns(uint256) {
74+
return comptroller.compAccrued(address(_setToken));
7475
}
7576

7677
/**

‎contracts/protocol/modules/ClaimModule.sol

+62-27
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,8 @@
1515
pragma solidity 0.6.10;
1616
pragma experimental "ABIEncoderV2";
1717

18+
import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
19+
1820
import { AddressArrayUtils } from "../../lib/AddressArrayUtils.sol";
1921
import { IClaimAdapter } from "../../interfaces/IClaimAdapter.sol";
2022
import { IController } from "../../interfaces/IController.sol";
@@ -52,10 +54,15 @@ contract ClaimModule is ModuleBase {
5254
event RewardClaimed(
5355
ISetToken indexed _setToken,
5456
address indexed _rewardPool,
55-
IClaimAdapter _adapter,
57+
IClaimAdapter indexed _adapter,
5658
uint256 _amount
5759
);
5860

61+
event AnyoneClaimUpdated(
62+
ISetToken indexed _setToken,
63+
bool _anyoneClaim
64+
);
65+
5966
/* ============ Modifiers ============ */
6067

6168
/**
@@ -71,11 +78,16 @@ contract ClaimModule is ModuleBase {
7178
// Indicates if any address can call claim or just the manager of the SetToken
7279
mapping(ISetToken => bool) public anyoneClaim;
7380

74-
// Array of rewardPool addresses to claim rewards for the SetToken
81+
// Map and array of rewardPool addresses to claim rewards for the SetToken
7582
mapping(ISetToken => address[]) public rewardPoolList;
83+
// 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
84+
mapping(ISetToken => mapping(address => bool)) public rewardPoolStatus;
7685

77-
// Array of adapters associated to the rewardPool for the SetToken
86+
// Map and array of adapters associated to the rewardPool for the SetToken
7887
mapping(ISetToken => mapping(address => address[])) public claimSettings;
88+
// 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
89+
mapping(ISetToken => mapping(address => mapping(address => bool))) public claimSettingsStatus;
90+
7991

8092
/* ============ Constructor ============ */
8193

@@ -133,10 +145,10 @@ contract ClaimModule is ModuleBase {
133145
*
134146
* @param _setToken Address of SetToken
135147
*/
136-
function updateAnyoneClaim(ISetToken _setToken) external onlyManagerAndValidSet(_setToken) {
137-
anyoneClaim[_setToken] = !anyoneClaim[_setToken];
148+
function updateAnyoneClaim(ISetToken _setToken, bool _anyoneClaim) external onlyManagerAndValidSet(_setToken) {
149+
anyoneClaim[_setToken] = _anyoneClaim;
150+
emit AnyoneClaimUpdated(_setToken, _anyoneClaim);
138151
}
139-
140152
/**
141153
* SET MANAGER ONLY. Adds a new claim integration for an existent rewardPool. If rewardPool doesn't have existing
142154
* claims then rewardPool is added to rewardPoolLiost. The claim integration is associated to an adapter that
@@ -252,10 +264,27 @@ contract ClaimModule is ModuleBase {
252264
function removeModule() external override {
253265
delete anyoneClaim[ISetToken(msg.sender)];
254266

267+
// explicitly delete all elements for gas refund
255268
address[] memory setTokenPoolList = rewardPoolList[ISetToken(msg.sender)];
256269
for (uint256 i = 0; i < setTokenPoolList.length; i++) {
270+
271+
address[] storage adapterList = claimSettings[ISetToken(msg.sender)][setTokenPoolList[i]];
272+
for (uint256 j = 0; j < adapterList.length; j++) {
273+
274+
address toRemove = adapterList[j];
275+
claimSettingsStatus[ISetToken(msg.sender)][setTokenPoolList[i]][toRemove] = false;
276+
277+
delete adapterList[j];
278+
}
257279
delete claimSettings[ISetToken(msg.sender)][setTokenPoolList[i]];
258280
}
281+
282+
for (uint256 i = 0; i < rewardPoolList[ISetToken(msg.sender)].length; i++) {
283+
address toRemove = rewardPoolList[ISetToken(msg.sender)][i];
284+
rewardPoolStatus[ISetToken(msg.sender)][toRemove] = false;
285+
286+
delete rewardPoolList[ISetToken(msg.sender)][i];
287+
}
259288
delete rewardPoolList[ISetToken(msg.sender)];
260289
}
261290

@@ -277,7 +306,7 @@ contract ClaimModule is ModuleBase {
277306
* @return Boolean indicating if the rewardPool is in the list for claims.
278307
*/
279308
function isRewardPool(ISetToken _setToken, address _rewardPool) public view returns (bool) {
280-
return rewardPoolList[_setToken].contains(_rewardPool);
309+
return rewardPoolStatus[_setToken][_rewardPool];
281310
}
282311

283312
/**
@@ -309,7 +338,7 @@ contract ClaimModule is ModuleBase {
309338
returns (bool)
310339
{
311340
address adapter = getAndValidateAdapter(_integrationName);
312-
return claimSettings[_setToken][_rewardPool].contains(adapter);
341+
return claimSettingsStatus[_setToken][_rewardPool][adapter];
313342
}
314343

315344
/**
@@ -329,9 +358,8 @@ contract ClaimModule is ModuleBase {
329358
view
330359
returns (uint256)
331360
{
332-
IClaimAdapter adapter = _getAndValidateIntegrationAdapter(claimSettings[_setToken][_rewardPool], _integrationName);
333-
uint256 rewards = adapter.getRewards(address(_setToken), _rewardPool);
334-
return rewards;
361+
IClaimAdapter adapter = _getAndValidateIntegrationAdapter(_setToken, _rewardPool, _integrationName);
362+
return adapter.getRewardsAmount(_setToken, _rewardPool);
335363
}
336364

337365
/* ============ Internal Functions ============ */
@@ -346,40 +374,45 @@ contract ClaimModule is ModuleBase {
346374
*/
347375
function _claim(ISetToken _setToken, address _rewardPool, string calldata _integrationName) internal {
348376
require(isRewardPool(_setToken, _rewardPool), "RewardPool not present");
349-
IClaimAdapter adapter = _getAndValidateIntegrationAdapter(claimSettings[_setToken][_rewardPool], _integrationName);
377+
IClaimAdapter adapter = _getAndValidateIntegrationAdapter(_setToken, _rewardPool, _integrationName);
350378

351-
uint256 rewards = adapter.getRewards(address(_setToken), _rewardPool);
379+
IERC20 rewardsToken = IERC20(adapter.getTokenAddress(_rewardPool));
380+
uint256 initRewardsBalance = rewardsToken.balanceOf(address(_setToken));
352381

353382
(
354383
address callTarget,
355384
uint256 callValue,
356385
bytes memory callByteData
357386
) = adapter.getClaimCallData(
358-
address(_setToken),
387+
_setToken,
359388
_rewardPool
360389
);
361390

362391
_setToken.invoke(callTarget, callValue, callByteData);
363392

364-
emit RewardClaimed(_setToken, _rewardPool, adapter, rewards);
393+
uint256 finalRewardsBalance = rewardsToken.balanceOf(address(_setToken));
394+
395+
emit RewardClaimed(_setToken, _rewardPool, adapter, finalRewardsBalance.sub(initRewardsBalance));
365396
}
366397

367398
/**
368399
* Gets the adapter and validate it is associated to the list of claim integration of a rewardPool.
369400
*
370-
* @param _rewardPoolClaimSettings List of claim integrations associated to a rewardPool
371-
* @param _integrationName ID of claim module integration (mapping on integration registry)
401+
* @param _setToken Address of SetToken
402+
* @param _rewardsPool Sddress of rewards pool
403+
* @param _integrationName ID of claim module integration (mapping on integration registry)
372404
*/
373405
function _getAndValidateIntegrationAdapter(
374-
address[] memory _rewardPoolClaimSettings,
406+
ISetToken _setToken,
407+
address _rewardsPool,
375408
string calldata _integrationName
376409
)
377410
internal
378411
view
379412
returns (IClaimAdapter)
380413
{
381414
address adapter = getAndValidateAdapter(_integrationName);
382-
require(_rewardPoolClaimSettings.contains(adapter), "Adapter integration not present");
415+
require(claimSettingsStatus[_setToken][_rewardsPool][adapter], "Adapter integration not present");
383416
return IClaimAdapter(adapter);
384417
}
385418

@@ -395,11 +428,13 @@ contract ClaimModule is ModuleBase {
395428
address adapter = getAndValidateAdapter(_integrationName);
396429
address[] storage _rewardPoolClaimSettings = claimSettings[_setToken][_rewardPool];
397430

398-
require(!_rewardPoolClaimSettings.contains(adapter), "Integration names must be unique");
431+
require(!claimSettingsStatus[_setToken][_rewardPool][adapter], "Integration names must be unique");
399432
_rewardPoolClaimSettings.push(adapter);
433+
claimSettingsStatus[_setToken][_rewardPool][adapter] = true;
400434

401-
if (_rewardPoolClaimSettings.length == 1) {
435+
if (!rewardPoolStatus[_setToken][_rewardPool]) {
402436
rewardPoolList[_setToken].push(_rewardPool);
437+
rewardPoolStatus[_setToken][_rewardPool] = true;
403438
}
404439
}
405440

@@ -428,7 +463,7 @@ contract ClaimModule is ModuleBase {
428463
}
429464

430465
/**
431-
* Validates and store the adapter address used to claim rewards for the passed rewardPool. If no adapters
466+
* Validates and stores the adapter address used to claim rewards for the passed rewardPool. If no adapters
432467
* left after removal then remove rewardPool from rewardPoolList and delete entry in claimSettings.
433468
*
434469
* @param _setToken Address of SetToken
@@ -437,14 +472,14 @@ contract ClaimModule is ModuleBase {
437472
*/
438473
function _removeClaim(ISetToken _setToken, address _rewardPool, string calldata _integrationName) internal {
439474
address adapter = getAndValidateAdapter(_integrationName);
440-
address[] memory _rewardPoolClaimSettings = claimSettings[_setToken][_rewardPool];
441475

442-
require(_rewardPoolClaimSettings.contains(adapter), "Integration must be added");
443-
claimSettings[_setToken][_rewardPool] = _rewardPoolClaimSettings.remove(adapter);
476+
require(claimSettingsStatus[_setToken][_rewardPool][adapter], "Integration must be added");
477+
claimSettings[_setToken][_rewardPool].removeStorage(adapter);
478+
claimSettingsStatus[_setToken][_rewardPool][adapter] = false;
444479

445480
if (claimSettings[_setToken][_rewardPool].length == 0) {
446-
rewardPoolList[_setToken] = rewardPoolList[_setToken].remove(_rewardPool);
447-
delete claimSettings[_setToken][_rewardPool];
481+
rewardPoolList[_setToken].removeStorage(_rewardPool);
482+
rewardPoolStatus[_setToken][_rewardPool] = false;
448483
}
449484
}
450485

‎test/integration/claim/compClaimAdapter.spec.ts

+2-2
Original file line numberDiff line numberDiff line change
@@ -67,15 +67,15 @@ describe("CompClaimAdapter", function() {
6767
});
6868
});
6969

70-
describe("#getRewards", async function() {
70+
describe("#getRewardsAmount", async function() {
7171
const rewards: BigNumber = ether(1);
7272

7373
before(async function() {
7474
await comptroller.mock.compAccrued.returns(rewards);
7575
});
7676

7777
function subject(): Promise<BigNumber> {
78-
return compClaimAdapter.connect(owner.wallet).getRewards(ADDRESS_ZERO, ADDRESS_ZERO);
78+
return compClaimAdapter.connect(owner.wallet).getRewardsAmount(ADDRESS_ZERO, ADDRESS_ZERO);
7979
}
8080

8181
it("should return rewards", async function() {

0 commit comments

Comments
 (0)