Merge branch 'main' into bug-bounty

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<TContractState> {
-    // 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<TContractState> {
 
     // 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<u64, u128>,
     }
 
@@ -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<u64> = 128;
@@ -144,17 +158,18 @@ pub mod Airdrop {
 
     #[abi(embed_v0)]
     impl AirdropImpl of IAirdrop<ContractState> {
-        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<felt252>) -> 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,14 @@ 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(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());
+            token.transfer(config.refund_to, balance);
+        }
     }
 }

src/airdrop_test.cairo

@@ -1,16 +1,17 @@
+use core::num::traits::zero::{Zero};
 use core::array::{ArrayTrait, SpanTrait};
 use core::hash::{LegacyHash};
 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};
 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
@@ -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<felt252> = 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!(
@@ -829,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::Claimed>(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();
+    }
+}

src/e2e_test.cairo

@@ -45,7 +45,7 @@ impl DefaultDeploymentParameters of Default<DeploymentParameters> {
                 proposal_creation_threshold: 20
             },
             timelock_config: TimelockConfig { delay: 30, window: 90, },
-            airdrop_root: Option::None
+            airdrop_config: Option::None
         }
     }
 }

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<felt252>,
+    pub airdrop_config: Option<AirdropConfig>,
 }
 
 #[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<felt252> = 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(),

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);