Skip to content

Commit

Permalink
update: claim rewards by pool
Browse files Browse the repository at this point in the history
  • Loading branch information
@waldoclayton
waldoclayton committed Nov 3, 2023
1 parent 646b10c commit 1b9b85e
Show file tree
Hide file tree
Showing 4 changed files with 106 additions and 79 deletions.
9 changes: 5 additions & 4 deletions contracts/plugins/RewardCollectorV2.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,10 +25,11 @@ contract RewardCollectorV2 is RewardCollector {
}

/// @notice Allows a user to claim rewards from the distributor
/// @param _nonce The nonce for the claim
/// @param _totalReward The total reward amount of the sender
/// @param _pool The pool from which the reward is claimed
/// @param _nonce The nonce of the sender for the claim
/// @param _totalReward The total reward amount of the sender in the pool
/// @param _signature The signature for the claim
function claim(uint32 _nonce, uint224 _totalReward, bytes memory _signature) external {
distributor.claim(msg.sender, _nonce, _totalReward, _signature, address(this));
function claim(address _pool, uint32 _nonce, uint256 _totalReward, bytes memory _signature) external {
distributor.claim(_pool, msg.sender, _nonce, _totalReward, _signature, address(this));
}
}
69 changes: 40 additions & 29 deletions contracts/plugins/RewardDistributor.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,29 +11,31 @@ import "@openzeppelin/contracts/utils/cryptography/ECDSA.sol";
contract RewardDistributor is Ownable2Step {
using ECDSA for bytes32;

/// @notice Struct to store claimed info for an account
struct ClaimedInfo {
/// @notice The claimed count of the account
uint32 nonce;
/// @notice The claimed reward of the account
uint224 claimedReward;
}

/// @notice The address of the signer
address public immutable signer;
/// @notice The address of the token
IERC20 public immutable token;
/// @notice The collectors
mapping(address => bool) public collectors;
/// @notice The claimed info for each account
mapping(address => ClaimedInfo) public claimedInfos;
/// @notice The nonces for each account
mapping(address => uint32) public nonces;
/// @notice Mapping of pool accounts to their claimed rewards
/// pool => account => claimedReward
mapping(address => mapping(address => uint256)) public claimedRewards;

/// @dev Event emitted when a claim is made
/// @param receiver The address that received the reward
/// @param pool The pool from which to claim the reward
/// @param account The account that claimed the reward for
/// @param nonce The nonce of the claim
/// @param nonce The nonce of the sender for the claim
/// @param receiver The address that received the reward
/// @param amount The amount of the reward claimed
event Claimed(address indexed receiver, address indexed account, uint32 indexed nonce, uint224 amount);
event Claimed(
address indexed pool,
address indexed account,
uint32 indexed nonce,
address receiver,
uint256 amount
);

/// @notice Error thrown when the caller is not authorized
/// @param caller The caller address
Expand Down Expand Up @@ -70,49 +72,58 @@ contract RewardDistributor is Ownable2Step {

/// @notice Allows a user to claim their reward by providing a valid signature
/// @dev The account and the receiver are the sender
/// @param _nonce The nonce for the claim
/// @param _totalReward The total reward amount of the sender
/// @param _pool The pool from which to claim the reward
/// @param _nonce The nonce of the sender for the claim
/// @param _totalReward The total reward amount of the sender in the pool
/// @param _signature The signature of the signer to verify
/// @param _receiver The receiver of the claim
function claim(uint32 _nonce, uint224 _totalReward, bytes memory _signature, address _receiver) external {
function claim(
address _pool,
uint32 _nonce,
uint256 _totalReward,
bytes memory _signature,
address _receiver
) external {
if (_receiver == address(0)) _receiver = msg.sender;
_claim(msg.sender, _nonce, _totalReward, _signature, _receiver);
_claim(_pool, msg.sender, _nonce, _totalReward, _signature, _receiver);
}

/// @notice Claims a reward for a specific account
/// @param _pool The pool from which to claim the reward
/// @param _account The account to claim for
/// @param _nonce The nonce for the claim
/// @param _totalReward The total reward amount of the account
/// @param _nonce The nonce of the sender for the claim
/// @param _totalReward The total reward amount of the account in pool
/// @param _signature The signature for the claim
/// @param _receiver The receiver of the claim
function claim(
address _pool,
address _account,
uint32 _nonce,
uint224 _totalReward,
uint256 _totalReward,
bytes memory _signature,
address _receiver
) external onlyCollector {
if (_receiver == address(0)) _receiver = msg.sender;
_claim(_account, _nonce, _totalReward, _signature, _receiver);
_claim(_pool, _account, _nonce, _totalReward, _signature, _receiver);
}

function _claim(
address _pool,
address _account,
uint32 _nonce,
uint224 _totalReward,
uint256 _totalReward,
bytes memory _signature,
address _receiver
) private {
ClaimedInfo storage claimedInfo = claimedInfos[_account];
if (_nonce != claimedInfo.nonce + 1) revert InvalidNonce(_nonce);
address _signer = keccak256(abi.encode(_account, _nonce, _totalReward)).toEthSignedMessageHash().recover(
if (_nonce != nonces[_account] + 1) revert InvalidNonce(_nonce);
address _signer = keccak256(abi.encode(_pool, _account, _nonce, _totalReward)).toEthSignedMessageHash().recover(
_signature
);
if (_signer != signer) revert InvalidSignature();
uint224 claimableReward = _totalReward - claimedInfo.claimedReward;
emit Claimed(_receiver, _account, _nonce, claimableReward);
claimedInfo.nonce = _nonce;
claimedInfo.claimedReward += claimableReward;
uint256 claimableReward = _totalReward - claimedRewards[_pool][_account];
emit Claimed(_pool, _account, _nonce, _receiver, claimableReward);
nonces[_account] = _nonce;
claimedRewards[_pool][_account] += claimableReward;
Address.functionCall(
address(token),
abi.encodeWithSignature("mint(address,uint256)", _receiver, claimableReward)
Expand Down
21 changes: 14 additions & 7 deletions test/foundry/RewardCollectorV2.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,13 @@ contract RewardCollectorV2Test is Test {
RewardCollectorV2 private rewardCollectorV2;

event Transfer(address indexed from, address indexed to, uint256 value);
event Claimed(address indexed receiver, address indexed account, uint32 indexed nonce, uint224 amount);
event Claimed(
address indexed pool,
address indexed account,
uint32 indexed nonce,
address receiver,
uint256 amount
);

function setUp() public {
token = new EQU();
Expand All @@ -28,25 +34,26 @@ contract RewardCollectorV2Test is Test {
}

function testMulticall() public {
address account = address(1);
address pool = address(1);
address account = address(2);
uint32 nonce = 1;
uint224 totalReward = 100;
bytes32 hash = keccak256(abi.encode(account, nonce, totalReward)).toEthSignedMessageHash();
uint256 totalReward = 100;
bytes32 hash = keccak256(abi.encode(pool, account, nonce, totalReward)).toEthSignedMessageHash();
(uint8 v, bytes32 r, bytes32 s) = vm.sign(SIGNER_PRIVATE_KEY, hash);
bytes memory signature = abi.encodePacked(r, s, v);
bytes[] memory data = new bytes[](2);
data[0] = abi.encodeWithSignature("claim(uint32,uint224,bytes)", nonce, totalReward, signature);
data[0] = abi.encodeWithSignature("claim(address,uint32,uint256,bytes)", pool, nonce, totalReward, signature);
data[1] = abi.encodeWithSignature("sweepToken(address,uint256,address)", address(token), 0, account);
vm.prank(account);
vm.expectEmit(true, true, true, true);
emit Claimed(address(rewardCollectorV2), account, nonce, totalReward);
emit Claimed(pool, account, nonce, address(rewardCollectorV2), totalReward);
vm.expectEmit(true, true, false, true);
emit Transfer(address(0), address(rewardCollectorV2), totalReward);
vm.expectEmit(true, true, false, true);
emit Transfer(address(rewardCollectorV2), account, totalReward);
bytes[] memory results = rewardCollectorV2.multicall(data);
assertEq(results[0], bytes(""));
assertEq(abi.decode(results[1], (uint224)), totalReward);
assertEq(abi.decode(results[1], (uint256)), totalReward);
assertEq(token.balanceOf(account), totalReward);
}
}
86 changes: 47 additions & 39 deletions test/foundry/RewardDistributor.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,17 +10,24 @@ import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
contract RewardDistributorTest is Test {
uint256 private constant SIGNER_PRIVATE_KEY = 0x12345;
uint256 private constant OTHER_PRIVATE_KEY = 0x54321;
address private constant ACCOUNT = address(1);
address private constant RECEIVER = address(2);
address private constant OTHER_ACCOUNT = address(3);
address private constant POOL = address(1);
address private constant ACCOUNT = address(2);
address private constant RECEIVER = address(3);
address private constant OTHER_ACCOUNT = address(4);
uint32 private constant NONCE = 1;
uint224 private constant TOTAL_REWARD = 1000;
uint256 private constant TOTAL_REWARD = 1000;

RewardDistributor private rewardDistributor;
IERC20 private token;

event Transfer(address indexed from, address indexed to, uint256 value);
event Claimed(address indexed receiver, address indexed account, uint32 indexed nonce, uint224 amount);
event Claimed(
address indexed pool,
address indexed account,
uint32 indexed nonce,
address receiver,
uint256 amount
);

function setUp() public {
address signer = vm.addr(SIGNER_PRIVATE_KEY);
Expand All @@ -32,98 +39,99 @@ contract RewardDistributorTest is Test {

function testClaim_RevertIfTheCallerIsNotTheCollector() public {
vm.prank(OTHER_ACCOUNT);
bytes memory signature = sign(SIGNER_PRIVATE_KEY, ACCOUNT, NONCE, TOTAL_REWARD);
bytes memory signature = sign(SIGNER_PRIVATE_KEY, POOL, ACCOUNT, NONCE, TOTAL_REWARD);
vm.expectRevert(abi.encodeWithSignature("CallerUnauthorized(address)", OTHER_ACCOUNT));
rewardDistributor.claim(ACCOUNT, NONCE, TOTAL_REWARD, signature, RECEIVER);
rewardDistributor.claim(POOL, ACCOUNT, NONCE, TOTAL_REWARD, signature, RECEIVER);
}

function testClaim_RevertIfTheNonceIsInvalid() public {
bytes memory signature = sign(SIGNER_PRIVATE_KEY, ACCOUNT, NONCE, TOTAL_REWARD);
bytes memory signature = sign(SIGNER_PRIVATE_KEY, POOL, ACCOUNT, NONCE, TOTAL_REWARD);
vm.expectRevert(abi.encodeWithSignature("InvalidNonce(uint32)", NONCE + 1));
rewardDistributor.claim(ACCOUNT, NONCE + 1, TOTAL_REWARD, signature, RECEIVER);
}

function testClaim_RevertIfAccountIsNotTheSignedAccount() public {
bytes memory signature = sign(SIGNER_PRIVATE_KEY, ACCOUNT, NONCE, TOTAL_REWARD);
bytes memory signature = sign(SIGNER_PRIVATE_KEY, POOL, ACCOUNT, NONCE, TOTAL_REWARD);
vm.expectRevert(RewardDistributor.InvalidSignature.selector);
rewardDistributor.claim(OTHER_ACCOUNT, NONCE, TOTAL_REWARD, signature, RECEIVER);
}

function testClaim_RevertIfTheTotalRewardIsInvalid() public {
bytes memory signature = sign(SIGNER_PRIVATE_KEY, ACCOUNT, NONCE, TOTAL_REWARD);
bytes memory signature = sign(SIGNER_PRIVATE_KEY, POOL, ACCOUNT, NONCE, TOTAL_REWARD);
vm.expectRevert(RewardDistributor.InvalidSignature.selector);
rewardDistributor.claim(ACCOUNT, NONCE, TOTAL_REWARD + 1, signature, RECEIVER);
}

function testClaim_RevertIfPrivateKeyIsNotSignerPrivateKey() public {
bytes memory signature = sign(OTHER_PRIVATE_KEY, ACCOUNT, NONCE, TOTAL_REWARD);
bytes memory signature = sign(OTHER_PRIVATE_KEY, POOL, ACCOUNT, NONCE, TOTAL_REWARD);
vm.expectRevert(RewardDistributor.InvalidSignature.selector);
rewardDistributor.claim(ACCOUNT, NONCE, TOTAL_REWARD, signature, RECEIVER);
}

function testClaim_ShouldUpdateWithTheRightValueAndEmitTheRightEvent() public {
deal(address(token), address(RECEIVER), 1000);
bytes memory signature = sign(SIGNER_PRIVATE_KEY, ACCOUNT, NONCE, TOTAL_REWARD);
bytes memory signature = sign(SIGNER_PRIVATE_KEY, POOL, ACCOUNT, NONCE, TOTAL_REWARD);
vm.expectEmit(true, true, true, true);
emit Claimed(RECEIVER, ACCOUNT, NONCE, TOTAL_REWARD);
emit Claimed(POOL, ACCOUNT, NONCE, RECEIVER, TOTAL_REWARD);
vm.expectEmit(true, true, false, true);
emit Transfer(address(0), RECEIVER, TOTAL_REWARD);
rewardDistributor.claim(ACCOUNT, NONCE, TOTAL_REWARD, signature, RECEIVER);
rewardDistributor.claim(POOL, ACCOUNT, NONCE, TOTAL_REWARD, signature, RECEIVER);
assertEq(token.balanceOf(RECEIVER), TOTAL_REWARD + 1000);
(uint32 nonce, uint224 claimedReward) = rewardDistributor.claimedInfos(ACCOUNT);
assertEq(nonce, NONCE);
assertEq(claimedReward, TOTAL_REWARD);
assertEq(rewardDistributor.nonces(ACCOUNT), NONCE);
assertEq(rewardDistributor.claimedRewards(POOL, ACCOUNT), TOTAL_REWARD);

uint224 newAmount = 2000;
signature = sign(SIGNER_PRIVATE_KEY, ACCOUNT, NONCE + 1, TOTAL_REWARD + newAmount);
uint256 newAmount = 2000;
signature = sign(SIGNER_PRIVATE_KEY, POOL, ACCOUNT, NONCE + 1, TOTAL_REWARD + newAmount);
vm.expectEmit(true, true, true, true);
emit Claimed(RECEIVER, ACCOUNT, NONCE + 1, newAmount);
emit Claimed(POOL, ACCOUNT, NONCE + 1, RECEIVER, newAmount);
vm.expectEmit(true, true, false, true);
emit Transfer(address(0), RECEIVER, newAmount);
rewardDistributor.claim(ACCOUNT, NONCE + 1, TOTAL_REWARD + newAmount, signature, RECEIVER);
rewardDistributor.claim(POOL, ACCOUNT, NONCE + 1, TOTAL_REWARD + newAmount, signature, RECEIVER);
assertEq(token.balanceOf(RECEIVER), TOTAL_REWARD + newAmount + 1000);
(nonce, claimedReward) = rewardDistributor.claimedInfos(ACCOUNT);
assertEq(nonce, NONCE + 1);
assertEq(claimedReward, TOTAL_REWARD + newAmount);
assertEq(rewardDistributor.nonces(ACCOUNT), NONCE + 1);
assertEq(rewardDistributor.claimedRewards(POOL, ACCOUNT), TOTAL_REWARD + newAmount);
}

function testClaim_RevertIfTheSameNonceIsUsedAgain() public {
bytes memory signature = sign(SIGNER_PRIVATE_KEY, ACCOUNT, NONCE, TOTAL_REWARD);
rewardDistributor.claim(ACCOUNT, NONCE, TOTAL_REWARD, signature, RECEIVER);
bytes memory signature = sign(SIGNER_PRIVATE_KEY, POOL, ACCOUNT, NONCE, TOTAL_REWARD);
rewardDistributor.claim(POOL, ACCOUNT, NONCE, TOTAL_REWARD, signature, RECEIVER);
vm.expectRevert(abi.encodeWithSignature("InvalidNonce(uint32)", NONCE));
rewardDistributor.claim(ACCOUNT, NONCE, TOTAL_REWARD, signature, RECEIVER);
rewardDistributor.claim(POOL, ACCOUNT, NONCE, TOTAL_REWARD, signature, RECEIVER);
}

function testClaim_RevertIfTheTotalRewardIsLtClaimedReward() public {
bytes memory signature = sign(SIGNER_PRIVATE_KEY, ACCOUNT, NONCE, TOTAL_REWARD);
rewardDistributor.claim(ACCOUNT, NONCE, TOTAL_REWARD, signature, RECEIVER);
signature = sign(SIGNER_PRIVATE_KEY, ACCOUNT, NONCE + 1, TOTAL_REWARD - 1);
bytes memory signature = sign(SIGNER_PRIVATE_KEY, POOL, ACCOUNT, NONCE, TOTAL_REWARD);
rewardDistributor.claim(POOL, ACCOUNT, NONCE, TOTAL_REWARD, signature, RECEIVER);
signature = sign(SIGNER_PRIVATE_KEY, POOL, ACCOUNT, NONCE + 1, TOTAL_REWARD - 1);
vm.expectRevert(stdError.arithmeticError);
rewardDistributor.claim(ACCOUNT, NONCE + 1, TOTAL_REWARD - 1, signature, RECEIVER);
rewardDistributor.claim(POOL, ACCOUNT, NONCE + 1, TOTAL_REWARD - 1, signature, RECEIVER);
}

function testClaim_BySender() public {
vm.prank(ACCOUNT);
bytes memory signature = sign(SIGNER_PRIVATE_KEY, ACCOUNT, NONCE, TOTAL_REWARD);
bytes memory signature = sign(SIGNER_PRIVATE_KEY, POOL, ACCOUNT, NONCE, TOTAL_REWARD);
vm.expectEmit(true, true, true, true);
emit Claimed(ACCOUNT, ACCOUNT, NONCE, TOTAL_REWARD);
emit Claimed(POOL, ACCOUNT, NONCE, ACCOUNT, TOTAL_REWARD);
vm.expectEmit(true, true, false, true);
emit Transfer(address(0), ACCOUNT, TOTAL_REWARD);
rewardDistributor.claim(NONCE, TOTAL_REWARD, signature, address(0));
rewardDistributor.claim(POOL, NONCE, TOTAL_REWARD, signature, address(0));
assertEq(token.balanceOf(ACCOUNT), TOTAL_REWARD);
(uint32 nonce, uint224 claimedReward) = rewardDistributor.claimedInfos(ACCOUNT);
assertEq(nonce, NONCE);
assertEq(claimedReward, TOTAL_REWARD);
assertEq(rewardDistributor.nonces(ACCOUNT), NONCE);
assertEq(rewardDistributor.claimedRewards(POOL, ACCOUNT), TOTAL_REWARD);
}

function sign(
uint256 _privateKey,
address _pool,
address _account,
uint32 _nonce,
uint224 _totalReward
uint256 _totalReward
) private pure returns (bytes memory signature) {
bytes32 hash = keccak256(
abi.encodePacked("\x19Ethereum Signed Message:\n32", keccak256(abi.encode(_account, _nonce, _totalReward)))
abi.encodePacked(
"\x19Ethereum Signed Message:\n32",
keccak256(abi.encode(_pool, _account, _nonce, _totalReward))
)
);
(uint8 v, bytes32 r, bytes32 s) = vm.sign(_privateKey, hash);
signature = abi.encodePacked(r, s, v);
Expand Down

0 comments on commit 1b9b85e

Please sign in to comment.