From 420f42345484fafeccdc5be93c439e1b19180300 Mon Sep 17 00:00:00 2001 From: Moody Salem Date: Thu, 14 Mar 2024 19:37:30 -0400 Subject: [PATCH 1/2] add the ability to set a refundable timestamp in the airdrop --- src/airdrop.cairo | 55 ++++++++++++++++++++++++++++++------------ src/airdrop_test.cairo | 15 +++++++++--- src/e2e_test.cairo | 2 +- src/factory.cairo | 9 ++++--- src/factory_test.cairo | 19 ++++++++++++--- 5 files changed, 74 insertions(+), 26 deletions(-) diff --git a/src/airdrop.cairo b/src/airdrop.cairo index 65c6794..86e1498 100644 --- a/src/airdrop.cairo +++ b/src/airdrop.cairo @@ -12,14 +12,24 @@ pub struct Claim { pub amount: u128, } +#[derive(Copy, Drop, Serde, Hash, PartialEq, Debug, starknet::Store)] +pub struct Config { + // the root of the airdrop + pub root: felt252, + // the timestamp after which anyone can refund any remaining tokens + pub refundable_timestamp: u64, + // the address that receives the refund + pub refund_to: ContractAddress, +} + #[starknet::interface] pub trait IAirdrop { - // Return the root of the airdrop - fn get_root(self: @TContractState) -> felt252; - - // Return the token being dropped + // Returns the address of the token distributed by this airdrop fn get_token(self: @TContractState) -> ContractAddress; + // Returns the config of this deployed airdrop + fn get_config(self: @TContractState) -> Config; + // Claims the given allotment of tokens. // Because this method is idempotent, it does not revert in case of a second submission of the same claim. // This makes it simpler to batch many claims together in a single transaction. @@ -36,6 +46,9 @@ pub trait IAirdrop { // Return whether the claim with the given ID has been claimed fn is_claimed(self: @TContractState, claim_id: u64) -> bool; + + // Refunds the current token balance. Can be called by anyone after the refund timestamp. + fn refund(ref self: TContractState); } #[starknet::contract] @@ -46,7 +59,8 @@ pub mod Airdrop { use core::num::traits::zero::{Zero}; use governance::interfaces::erc20::{IERC20DispatcherTrait}; use governance::utils::exp2::{exp2}; - use super::{IAirdrop, ContractAddress, Claim, IERC20Dispatcher}; + use starknet::{get_block_timestamp, get_contract_address}; + use super::{Config, IAirdrop, ContractAddress, Claim, IERC20Dispatcher}; pub fn hash_function(a: felt252, b: felt252) -> felt252 { @@ -70,8 +84,8 @@ pub mod Airdrop { #[storage] struct Storage { - root: felt252, token: IERC20Dispatcher, + config: Config, claimed_bitmap: LegacyMap, } @@ -87,9 +101,9 @@ pub mod Airdrop { } #[constructor] - fn constructor(ref self: ContractState, token: IERC20Dispatcher, root: felt252) { - self.root.write(root); - self.token.write(token); + fn constructor(ref self: ContractState, token: ContractAddress, config: Config) { + self.token.write(IERC20Dispatcher { contract_address: token }); + self.config.write(config); } const BITMAP_SIZE: NonZero = 128; @@ -144,17 +158,18 @@ pub mod Airdrop { #[abi(embed_v0)] impl AirdropImpl of IAirdrop { - fn get_root(self: @ContractState) -> felt252 { - self.root.read() - } - fn get_token(self: @ContractState) -> ContractAddress { self.token.read().contract_address } + fn get_config(self: @ContractState) -> Config { + self.config.read() + } + fn claim(ref self: ContractState, claim: Claim, proof: Span) -> bool { let leaf = hash_claim(claim); - assert(self.root.read() == compute_pedersen_root(leaf, proof), 'INVALID_PROOF'); + let config = self.config.read(); + assert(config.root == compute_pedersen_root(leaf, proof), 'INVALID_PROOF'); // this is copied in from is_claimed because we only want to read the bitmap once let (word, index) = claim_id_to_bitmap_index(claim.id); @@ -187,8 +202,10 @@ pub mod Airdrop { let root_of_group = compute_root_of_group(claims); + let config = self.config.read(); + assert( - self.root.read() == compute_pedersen_root(root_of_group, remaining_proof), + config.root == compute_pedersen_root(root_of_group, remaining_proof), 'INVALID_PROOF' ); @@ -231,5 +248,13 @@ pub mod Airdrop { let bitmap = self.claimed_bitmap.read(word); (bitmap & exp2(index)).is_non_zero() } + + fn refund(ref self: ContractState) { + let config = self.config.read(); + assert(get_block_timestamp() >= config.refundable_timestamp, 'TOO_EARLY'); + let token = self.token.read(); + let balance = token.balanceOf(get_contract_address()); + token.transfer(config.refund_to, balance); + } } } diff --git a/src/airdrop_test.cairo b/src/airdrop_test.cairo index ea76fef..c279e0b 100644 --- a/src/airdrop_test.cairo +++ b/src/airdrop_test.cairo @@ -1,3 +1,4 @@ +use core::num::traits::zero::{Zero}; use core::array::{ArrayTrait, SpanTrait}; use core::hash::{LegacyHash}; use core::option::{OptionTrait}; @@ -5,7 +6,7 @@ use core::option::{OptionTrait}; use core::result::{Result, ResultTrait}; use core::traits::{TryInto, Into}; use governance::airdrop::{ - IAirdropDispatcher, IAirdropDispatcherTrait, Airdrop, + IAirdropDispatcher, IAirdropDispatcherTrait, Airdrop, Config, Airdrop::{compute_pedersen_root, hash_function, hash_claim, compute_root_of_group}, Claim }; use governance::interfaces::erc20::{IERC20Dispatcher, IERC20DispatcherTrait}; @@ -29,9 +30,13 @@ pub(crate) fn deploy_token(owner: ContractAddress, amount: u128) -> IERC20Dispat } -fn deploy(token: ContractAddress, root: felt252) -> IAirdropDispatcher { +fn deploy_with_refundee( + token: ContractAddress, root: felt252, refundable_timestamp: u64, refund_to: ContractAddress +) -> IAirdropDispatcher { let mut constructor_args: Array = ArrayTrait::new(); - Serde::serialize(@(token, root), ref constructor_args); + Serde::serialize( + @(token, Config { root, refundable_timestamp, refund_to }), ref constructor_args + ); let (address, _) = deploy_syscall( Airdrop::TEST_CLASS_HASH.try_into().unwrap(), 0, constructor_args.span(), true @@ -40,6 +45,10 @@ fn deploy(token: ContractAddress, root: felt252) -> IAirdropDispatcher { IAirdropDispatcher { contract_address: address } } +fn deploy(token: ContractAddress, root: felt252) -> IAirdropDispatcher { + deploy_with_refundee(token, root, Zero::zero(), Zero::zero()) +} + #[test] fn test_selector() { assert_eq!( diff --git a/src/e2e_test.cairo b/src/e2e_test.cairo index 63d7a2c..38c736c 100644 --- a/src/e2e_test.cairo +++ b/src/e2e_test.cairo @@ -45,7 +45,7 @@ impl DefaultDeploymentParameters of Default { proposal_creation_threshold: 20 }, timelock_config: TimelockConfig { delay: 30, window: 90, }, - airdrop_root: Option::None + airdrop_config: Option::None } } } diff --git a/src/factory.cairo b/src/factory.cairo index 070e3bb..10b6bf6 100644 --- a/src/factory.cairo +++ b/src/factory.cairo @@ -1,5 +1,6 @@ use governance::airdrop::{IAirdropDispatcher}; use governance::governor::{Config as GovernorConfig}; +use governance::airdrop::{Config as AirdropConfig}; use governance::governor::{IGovernorDispatcher}; use governance::staker::{IStakerDispatcher}; use governance::timelock::{ITimelockDispatcher, TimelockConfig}; @@ -9,7 +10,7 @@ use starknet::{ContractAddress}; pub struct DeploymentParameters { pub governor_config: GovernorConfig, pub timelock_config: TimelockConfig, - pub airdrop_root: Option, + pub airdrop_config: Option, } #[derive(Copy, Drop, Serde)] @@ -88,10 +89,10 @@ pub mod Factory { ) .unwrap(); - let airdrop = match params.airdrop_root { - Option::Some(root) => { + let airdrop = match params.airdrop_config { + Option::Some(config) => { let mut airdrop_constructor_args: Array = ArrayTrait::new(); - Serde::serialize(@(token, root), ref airdrop_constructor_args); + Serde::serialize(@(token, config), ref airdrop_constructor_args); let (airdrop_address, _) = deploy_syscall( class_hash: self.airdrop_class_hash.read(), diff --git a/src/factory_test.cairo b/src/factory_test.cairo index 12ce08e..e1117e4 100644 --- a/src/factory_test.cairo +++ b/src/factory_test.cairo @@ -3,7 +3,7 @@ use core::option::{OptionTrait}; use core::result::{Result}; use core::traits::{TryInto}; -use governance::airdrop::{Airdrop}; +use governance::airdrop::{Airdrop, Config as AirdropConfig}; use governance::airdrop::{IAirdropDispatcherTrait}; use governance::airdrop_test::{deploy_token}; use governance::factory::{ @@ -54,7 +54,13 @@ fn test_deploy() { .deploy( token, DeploymentParameters { - airdrop_root: Option::Some('root'), + airdrop_config: Option::Some( + AirdropConfig { + root: 'root', + refundable_timestamp: 312, + refund_to: contract_address_const::<'refund'>(), + } + ), governor_config: GovernorConfig { voting_start_delay: 0, voting_period: 180, @@ -67,7 +73,14 @@ fn test_deploy() { ); let drop = result.airdrop.unwrap(); - assert_eq!(drop.get_root(), 'root'); + assert_eq!( + drop.get_config(), + AirdropConfig { + root: 'root', + refundable_timestamp: 312, + refund_to: contract_address_const::<'refund'>(), + } + ); assert_eq!(drop.get_token(), token); assert_eq!(result.staker.get_token(), token); From 367ddb0df7e4ab00774455867ec05faff5a691de Mon Sep 17 00:00:00 2001 From: Moody Salem Date: Thu, 14 Mar 2024 20:10:58 -0400 Subject: [PATCH 2/2] add unit tests for the refundable feature --- src/airdrop.cairo | 1 + src/airdrop_test.cairo | 86 +++++++++++++++++++++++++++++++++++++++++- 2 files changed, 86 insertions(+), 1 deletion(-) diff --git a/src/airdrop.cairo b/src/airdrop.cairo index 86e1498..c76693a 100644 --- a/src/airdrop.cairo +++ b/src/airdrop.cairo @@ -251,6 +251,7 @@ pub mod Airdrop { fn refund(ref self: ContractState) { let config = self.config.read(); + assert(config.refundable_timestamp.is_non_zero(), 'NOT_REFUNDABLE'); assert(get_block_timestamp() >= config.refundable_timestamp, 'TOO_EARLY'); let token = self.token.read(); let balance = token.balanceOf(get_contract_address()); diff --git a/src/airdrop_test.cairo b/src/airdrop_test.cairo index c279e0b..e361252 100644 --- a/src/airdrop_test.cairo +++ b/src/airdrop_test.cairo @@ -11,7 +11,7 @@ use governance::airdrop::{ }; use governance::interfaces::erc20::{IERC20Dispatcher, IERC20DispatcherTrait}; use governance::test::test_token::{TestToken}; -use starknet::testing::{pop_log}; +use starknet::testing::{pop_log, set_block_timestamp}; use starknet::{ get_contract_address, syscalls::{deploy_syscall}, ClassHash, contract_address_const, ContractAddress @@ -838,3 +838,87 @@ fn test_claim_128_double_claim() { assert_eq!(airdrop.claim_128(claims.span().slice(0, 128), array![s2, rr].span()), 0); assert_eq!(pop_log::(airdrop.contract_address).is_none(), true); } + + +mod refundable { + use super::{ + deploy_token, deploy_with_refundee, get_contract_address, contract_address_const, + set_block_timestamp, IAirdropDispatcherTrait, IERC20DispatcherTrait, TestToken, + ContractAddress + }; + + fn refund_to() -> ContractAddress { + contract_address_const::<'refund_to'>() + } + + #[test] + fn test_refundable_transfers_token_at_refund_timestamp() { + let token = deploy_token(get_contract_address(), 1234); + let airdrop = deploy_with_refundee( + token: token.contract_address, + root: 'root', + refundable_timestamp: 1, + refund_to: refund_to() + ); + + set_block_timestamp(1); + token.transfer(airdrop.contract_address, 567); + assert_eq!(token.balanceOf(refund_to()), 0); + airdrop.refund(); + assert_eq!(token.balanceOf(refund_to()), 567); + airdrop.refund(); + assert_eq!(token.balanceOf(refund_to()), 567); + + set_block_timestamp(2); + airdrop.refund(); + assert_eq!(token.balanceOf(refund_to()), 567); + token.transfer(airdrop.contract_address, 123); + airdrop.refund(); + assert_eq!(token.balanceOf(refund_to()), 690); + } + + #[test] + #[should_panic(expected: ('TOO_EARLY', 'ENTRYPOINT_FAILED'))] + fn test_refundable_reverts_before_timestamp() { + let token = deploy_token(get_contract_address(), 1234); + let airdrop = deploy_with_refundee( + token: token.contract_address, + root: 'root', + refundable_timestamp: 1, + refund_to: refund_to() + ); + + token.transfer(airdrop.contract_address, 567); + airdrop.refund(); + } + + #[test] + #[should_panic(expected: ('NOT_REFUNDABLE', 'ENTRYPOINT_FAILED'))] + fn test_not_refundable_reverts() { + let token = deploy_token(get_contract_address(), 1234); + let airdrop = deploy_with_refundee( + token: token.contract_address, + root: 'root', + refundable_timestamp: 0, + refund_to: refund_to() + ); + + airdrop.refund(); + } + + #[test] + #[should_panic(expected: ('NOT_REFUNDABLE', 'ENTRYPOINT_FAILED'))] + fn test_not_refundable_reverts_after_time() { + let token = deploy_token(get_contract_address(), 1234); + let airdrop = deploy_with_refundee( + token: token.contract_address, + root: 'root', + refundable_timestamp: 0, + refund_to: refund_to() + ); + + set_block_timestamp(1); + + airdrop.refund(); + } +}