Skip to content

Commit a98b928

Browse files
PeterMPhillipsPeterMPhillips
committed
airdrop
1 parent af300b3 commit a98b928

File tree

5 files changed

+207
-2
lines changed

5 files changed

+207
-2
lines changed

‎script/TokenDeployer.s.sol‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
11
// SPDX-License-Identifier: UNLICENSED
22
pragma solidity ^0.8.13;
33

4-
import { Distribution } from "../src/libraries/Distribution.sol";
54
import { EnsoToken } from "../src/EnsoToken.sol";
5+
import { Distribution } from "../src/libraries/Distribution.sol";
66
import { Script, console } from "forge-std/Script.sol";
77
import { ERC1967Proxy } from "openzeppelin-contracts/proxy/ERC1967/ERC1967Proxy.sol";
88

‎src/Airdrop.sol‎

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
// SPDX-License-Identifier: GPL-3.0-only
2+
pragma solidity ^0.8.20;
3+
4+
import { Ownable } from "openzeppelin-contracts/access/Ownable.sol";
5+
import { IERC20 } from "openzeppelin-contracts/token/ERC20/IERC20.sol";
6+
import { MerkleProof } from "openzeppelin-contracts/utils/cryptography/MerkleProof.sol";
7+
8+
contract Airdrop is Ownable {
9+
event Claim(address to, uint256 amount);
10+
11+
error InvalidMerkleProof();
12+
error AlreadyClaimed(bytes32 leaf);
13+
error AirdropExpired();
14+
error AirdropNotExpired();
15+
16+
IERC20 public immutable token;
17+
bytes32 public immutable root;
18+
uint256 public immutable expiration;
19+
mapping(bytes32 => bool) public claimed;
20+
21+
constructor(address _token, bytes32 _root, uint256 _expiration, address _owner) Ownable(_owner) {
22+
token = IERC20(_token);
23+
root = _root;
24+
expiration = _expiration;
25+
}
26+
27+
function getLeafHash(address to, uint256 amount) public pure returns (bytes32) {
28+
return keccak256(abi.encode(to, amount));
29+
}
30+
31+
function claim(bytes32[] memory proof, address to, uint256 amount) external {
32+
if (block.timestamp > expiration) revert AirdropExpired();
33+
34+
// NOTE: (to, amount) cannot have duplicates
35+
bytes32 leaf = getLeafHash(to, amount);
36+
37+
if (claimed[leaf]) revert AlreadyClaimed(leaf);
38+
if (!MerkleProof.verify(proof, root, leaf)) revert InvalidMerkleProof();
39+
claimed[leaf] = true;
40+
41+
token.transfer(to, amount);
42+
43+
emit Claim(to, amount);
44+
}
45+
46+
function sweep(address to) external onlyOwner {
47+
if (block.timestamp <= expiration) revert AirdropNotExpired();
48+
token.transfer(to, token.balanceOf(address(this)));
49+
}
50+
}

‎test/AirdropTest.t.sol‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
// SPDX-License-Identifier: MIT
2+
pragma solidity ^0.8.20;
3+
4+
import { TokenDeployer } from "../script/TokenDeployer.s.sol";
5+
import { Airdrop } from "../src/Airdrop.sol";
6+
import { EnsoToken } from "../src/EnsoToken.sol";
7+
import { MerkleHelper } from "./helpers/MerkleHelper.sol";
8+
import { Test } from "forge-std/Test.sol";
9+
import { ERC1967Proxy } from "openzeppelin-contracts/proxy/ERC1967/ERC1967Proxy.sol";
10+
11+
contract AirdropTest is Test {
12+
TokenDeployer public deployer;
13+
EnsoToken public token;
14+
ERC1967Proxy public proxy;
15+
16+
struct Reward {
17+
address to;
18+
uint256 amount;
19+
}
20+
21+
Reward[] private rewards;
22+
bytes32[] private hashes;
23+
bytes32 private root;
24+
uint256 private total;
25+
26+
mapping(bytes32 => Reward) private hashToReward;
27+
28+
uint256 constant N = 100;
29+
30+
function setUp() public {
31+
deployer = new TokenDeployer();
32+
(proxy,) = deployer.deploy();
33+
token = EnsoToken(address(proxy));
34+
35+
// Initialize users and airdrop amounts
36+
total = 0;
37+
for (uint256 i = 0; i < N; i++) {
38+
uint256 amount = (i + 1) * 100;
39+
rewards.push(Reward({ to: address(uint160(i + 1)), amount: amount }));
40+
hashes.push(keccak256(abi.encode(rewards[i].to, rewards[i].amount)));
41+
hashToReward[hashes[i]] = rewards[i];
42+
total += amount;
43+
}
44+
45+
hashes = MerkleHelper.sort(hashes);
46+
47+
root = MerkleHelper.calcRoot(hashes);
48+
}
49+
50+
function test_valid_proof() public {
51+
Airdrop airdrop = deployAirdrop(block.number + 60);
52+
for (uint256 i = 0; i < N; i++) {
53+
bytes32 h = hashes[i];
54+
Reward memory reward = hashToReward[h];
55+
bytes32[] memory proof = MerkleHelper.getProof(hashes, i);
56+
57+
airdrop.claim(proof, reward.to, reward.amount);
58+
assertEq(token.balanceOf(reward.to), reward.amount);
59+
}
60+
}
61+
62+
function test_expiration() public {
63+
Airdrop airdrop = deployAirdrop(block.number - 1); // revert on claim
64+
bytes32 h = hashes[0];
65+
Reward memory reward = hashToReward[h];
66+
bytes32[] memory proof = MerkleHelper.getProof(hashes, 0);
67+
68+
vm.expectRevert(abi.encodeWithSelector(Airdrop.AirdropExpired.selector));
69+
airdrop.claim(proof, reward.to, reward.amount);
70+
71+
vm.prank(deployer.OWNER());
72+
airdrop.sweep(address(this));
73+
vm.assertEq(token.balanceOf(address(this)), total);
74+
}
75+
76+
function deployAirdrop(uint256 expiration) internal returns (Airdrop airdrop) {
77+
airdrop = new Airdrop(address(token), root, expiration, deployer.OWNER());
78+
vm.prank(deployer.OWNER());
79+
token.unpause();
80+
vm.prank(deployer.OWNER());
81+
token.transfer(address(airdrop), total);
82+
}
83+
}

‎test/TokenTest.t.sol‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ contract TokenTest is Test {
3333

3434
function test_OwnerBalance() public view {
3535
uint256 balance = token.balanceOf(deployer.OWNER());
36-
assertEq(balance, 95_600_000 * 10**18);
36+
assertEq(balance, 95_600_000 * 10 ** 18);
3737
}
3838

3939
function test_PausedFail() public {

‎test/helpers/MerkleHelper.sol‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
// SPDX-License-Identifier: MIT
2+
pragma solidity ^0.8.20;
3+
4+
import { Hashes } from "openzeppelin-contracts/utils/cryptography/Hashes.sol";
5+
6+
library MerkleHelper {
7+
// Bubble sort
8+
function sort(bytes32[] memory arr) internal pure returns (bytes32[] memory) {
9+
uint256 n = arr.length;
10+
for (uint256 i = 0; i < n; i++) {
11+
for (uint256 j = 0; j < n - 1 - i; j++) {
12+
if (arr[j] > arr[j + 1]) {
13+
(arr[j], arr[j + 1]) = (arr[j + 1], arr[j]);
14+
}
15+
}
16+
}
17+
18+
return arr;
19+
}
20+
21+
function calcRoot(bytes32[] memory hashes) internal pure returns (bytes32) {
22+
uint256 n = hashes.length;
23+
24+
while (n > 1) {
25+
for (uint256 i = 0; i < n; i += 2) {
26+
bytes32 left = hashes[i];
27+
bytes32 right = hashes[i + 1 < n ? i + 1 : i];
28+
(left, right) = left <= right ? (left, right) : (right, left);
29+
hashes[i >> 1] = Hashes.efficientKeccak256(left, right);
30+
}
31+
n = (n + (n & 1)) >> 1;
32+
}
33+
34+
return hashes[0];
35+
}
36+
37+
function getProof(bytes32[] memory hashes, uint256 index) internal pure returns (bytes32[] memory) {
38+
bytes32[] memory proof = new bytes32[](0);
39+
uint256 len = 0;
40+
41+
uint256 n = hashes.length;
42+
uint256 k = index;
43+
44+
while (n > 1) {
45+
// Get proof for this level
46+
uint256 j = k & 1 == 1 ? k - 1 : (k + 1 < n ? k + 1 : k);
47+
bytes32 h = hashes[j];
48+
49+
// proof.push(h)
50+
assembly {
51+
len := add(len, 1)
52+
let pos := add(proof, shl(5, len))
53+
mstore(pos, h)
54+
mstore(proof, len)
55+
mstore(0x40, add(pos, 0x20))
56+
}
57+
58+
k >>= 1;
59+
60+
// Calculate next level of hashes
61+
for (uint256 i = 0; i < n; i += 2) {
62+
bytes32 left = hashes[i];
63+
bytes32 right = hashes[i + 1 < n ? i + 1 : i];
64+
(left, right) = left <= right ? (left, right) : (right, left);
65+
hashes[i >> 1] = Hashes.efficientKeccak256(left, right);
66+
}
67+
n = (n + (n & 1)) >> 1;
68+
}
69+
70+
return proof;
71+
}
72+
}

0 commit comments

Comments
 (0)