Add extra deposit and withdraw methods

Scarb.toml

@@ -23,12 +23,20 @@ snforge_std = { git = "https://github.com/foundry-rs/starknet-foundry.git", tag
 [tool.fmt]
 sort-module-level-items = true
 
+[tool.snforge]
+fuzzer_runs = 10
+
 [[tool.snforge.fork]]
 name = "SEPOLIA"
-url = "http://178.32.172.148:6062/v0_7"
+url = "http://51.195.57.196:6062/v0_7"
 block_id.tag = "latest"
 
 [[tool.snforge.fork]]
 name = "MAINNET"
-url = "http://127.0.0.1:5050"
+url = "http://51.195.57.196:6060/v0_7"
 block_id.tag = "latest"
+
+[[tool.snforge.fork]]
+name = "MAINNET_FIXED_BLOCK"
+url = "http://51.195.57.196:6060/v0_7"
+block_id.number = "863242"

src/constants.cairo

@@ -1 +1,2 @@
 pub const ZK_SCALE_DECIMALS: u256 = 1000000000000000000000000000;
+pub const STRK_ADDRESS: felt252 = 0x04718f5a0fc34cc1af16a1cdee98ffb20c31f5cd61d6ab07201858f4287c938d;

src/deposit.cairo

@@ -9,7 +9,7 @@ mod Deposit {
 
     use openzeppelin_token::erc20::interface::{ERC20ABIDispatcher, ERC20ABIDispatcherTrait};
     use spotnet::{
-        constants::ZK_SCALE_DECIMALS,
+        constants::{ZK_SCALE_DECIMALS, STRK_ADDRESS},
         interfaces::{
             IMarketDispatcher, IMarketDispatcherTrait, IAirdropDispatcher, IAirdropDispatcherTrait,
             IDeposit
@@ -21,7 +21,7 @@ mod Deposit {
     };
 
     use starknet::{
-        ContractAddress, get_contract_address, get_tx_info, event::EventEmitter,
+        ContractAddress, get_contract_address, get_caller_address, get_tx_info, event::EventEmitter,
         storage::{StoragePointerWriteAccess, StoragePointerReadAccess}
     };
 
@@ -30,6 +30,7 @@ mod Deposit {
         owner: ContractAddress,
         ekubo_core: ICoreDispatcher,
         zk_market: IMarketDispatcher,
+        treasury: ContractAddress,
         is_position_open: bool
     }
 
@@ -38,11 +39,13 @@ mod Deposit {
         ref self: ContractState,
         owner: ContractAddress,
         ekubo_core: ICoreDispatcher,
-        zk_market: IMarketDispatcher
+        zk_market: IMarketDispatcher,
+        treasury: ContractAddress
     ) {
         self.owner.write(owner);
         self.ekubo_core.write(ekubo_core);
         self.zk_market.write(zk_market);
+        self.treasury.write(treasury);
     }
 
     fn get_borrow_amount(
@@ -133,8 +136,8 @@ mod Deposit {
             ekubo_limits: EkuboSlippageLimits,
             pool_price: TokenPrice
         ) {
-            let user_acount = get_tx_info().unbox().account_contract_address;
-            assert(user_acount == self.owner.read(), 'Caller is not the owner');
+            let user_account = get_tx_info().unbox().account_contract_address;
+            assert(user_account == self.owner.read(), 'Caller is not the owner');
             assert(!self.is_position_open.read(), 'Open position already exists');
             let DepositData { token, amount, multiplier, borrow_const } = deposit_data;
             assert(borrow_const > 0 && borrow_const < 100, 'Cannot calculate borrow amount');
@@ -146,10 +149,10 @@ mod Deposit {
 
             let curr_contract_address = get_contract_address();
             assert(
-                token_dispatcher.allowance(user_acount, curr_contract_address) >= amount,
+                token_dispatcher.allowance(user_account, curr_contract_address) >= amount,
                 'Approved amount insufficient'
             );
-            assert(token_dispatcher.balanceOf(user_acount) >= amount, 'Insufficient balance');
+            assert(token_dispatcher.balanceOf(user_account) >= amount, 'Insufficient balance');
 
             let zk_market = self.zk_market.read();
 
@@ -355,6 +358,11 @@ mod Deposit {
 
         /// Claims STRK airdrop on ZKlend
         ///
+        /// Can be called by anyone, e.g. a keeper
+        /// 
+        /// If the treasury address is zero, the funds are not sent to treasury to avoid burning them.
+        /// This does mean that a sophisticated user could deploy their own contract with a zero treasury address to avoid sending on fees.
+        /// 
         /// # Panics
         /// `is_position_open` storage variable is set to false('Open position not exists')
         /// `proof` span is empty
@@ -374,7 +382,69 @@ mod Deposit {
                 IAirdropDispatcher { contract_address: airdrop_addr }.claim(claim_data, proof),
                 'Claim failed'
             );
-            // TODO: Add transfer to the Treasury
+
+            let strk = ERC20ABIDispatcher { contract_address: STRK_ADDRESS.try_into().unwrap() };
+            let zk_market = self.zk_market.read();
+            let part_for_treasury = claim_data.amount - claim_data.amount / 5; // u128 integer division, rounds down
+
+            let treasury_addr = self.treasury.read();
+            let remainder = if(treasury_addr.into() != 0) { // Zeroable not publicly accessible in this Cairo version AFAIK
+               strk.transfer(treasury_addr, part_for_treasury.into());
+               claim_data.amount - part_for_treasury
+            } else {
+                claim_data.amount
+            };
+
+
+            strk.approve(zk_market.contract_address, remainder.into());
+            zk_market.deposit(STRK_ADDRESS.try_into().unwrap(), remainder.into());
+            zk_market.enable_collateral(STRK_ADDRESS.try_into().unwrap());
+        }
+
+        /// Makes a deposit into open zkLend position to control stability
+        /// 
+        /// Anyone can deposit theoretically
+        ///
+        /// # Panics
+        /// `is_position_open` variable is set to false
+        /// `amount` is equal to zero
+        ///
+        /// # Parameters
+        /// `token`: ContractAddress - token address to withdraw from zkLend
+        /// `amount`: TokenAmount - amount to withdraw
+        fn extra_deposit(ref self: ContractState, token: ContractAddress, amount: TokenAmount) {
+            assert(self.is_position_open.read(), 'Open position not exists');
+            assert(amount != 0, 'Deposit amount is zero');
+            let (zk_market, token_dispatcher) = (
+                self.zk_market.read(), ERC20ABIDispatcher { contract_address: token }
+            );
+            token_dispatcher.transferFrom(get_caller_address(), get_contract_address(), amount);
+            token_dispatcher.approve(zk_market.contract_address, amount);
+            zk_market.deposit(token, amount.try_into().unwrap());
+        }
+
+        /// Withdraws tokens from zkLend if looped tokens are repaid
+        ///
+        /// # Panics
+        /// address of account that started the transaction is not equal to `owner` storage variable
+        ///
+        /// # Parameters
+        /// `token`: TokenAddress - token address to withdraw from zkLend
+        /// `amount`: TokenAmount - amount to withdraw. Pass `0` to withdraw all
+        fn withdraw(ref self: ContractState, token: ContractAddress, amount: TokenAmount) {
+            assert(get_caller_address() == self.owner.read(), 'Caller is not the owner');
+            let zk_market = self.zk_market.read();
+            let token_dispatcher = ERC20ABIDispatcher { contract_address: token };
+            if amount == 0 {
+                zk_market.withdraw_all(token);
+                token_dispatcher
+                    .transfer(
+                        self.owner.read(), token_dispatcher.balanceOf(get_contract_address())
+                    );
+            } else {
+                zk_market.withdraw(token, amount.try_into().unwrap());
+                token_dispatcher.transfer(self.owner.read(), amount);
+            };
         }
     }
 

src/interfaces.cairo

@@ -1,5 +1,7 @@
 use ekubo::types::keys::PoolKey;
-use spotnet::types::{MarketReserveData, DepositData, Claim, EkuboSlippageLimits, TokenPrice};
+use spotnet::types::{
+    MarketReserveData, DepositData, Claim, EkuboSlippageLimits, TokenPrice, TokenAmount
+};
 use starknet::ContractAddress;
 
 #[starknet::interface]
@@ -28,6 +30,10 @@ pub trait IDeposit<TContractState> {
         proof: Span<felt252>,
         airdrop_addr: ContractAddress
     );
+
+    fn extra_deposit(ref self: TContractState, token: ContractAddress, amount: TokenAmount);
+
+    fn withdraw(ref self: TContractState, token: ContractAddress, amount: TokenAmount);
 }
 
 #[starknet::interface]

tests/lib.cairo

@@ -2,3 +2,6 @@ pub mod interfaces;
 
 #[cfg(test)]
 mod test_loop;
+
+#[cfg(test)]
+mod test_defispring;

tests/test_defispring.cairo

@@ -0,0 +1,80 @@
+use snforge_std::{declare, DeclareResultTrait, replace_bytecode, store, cheat_caller_address, CheatSpan};
+use starknet::ContractAddress;
+use spotnet::interfaces::{
+    IDepositDispatcher, IDepositDispatcherTrait
+};
+use spotnet::types::Claim;
+use spotnet::constants::STRK_ADDRESS;
+use openzeppelin_token::erc20::interface::{ERC20ABIDispatcher, ERC20ABIDispatcherTrait};
+
+const ADDRESS_ELIGIBLE_FOR_ZKLEND_REWARDS: felt252 = 0x020281104e6cb5884dabcdf3be376cf4ff7b680741a7bb20e5e07c26cd4870af;
+const HYPOTHETICAL_OWNER_ADDR: felt252 = 0x56789;
+
+#[test]
+#[fork("MAINNET_FIXED_BLOCK")]
+fn test_claim_as_keeper() {
+    let strk = ERC20ABIDispatcher { contract_address: STRK_ADDRESS.try_into().unwrap() };
+    let defispring_claim_contract: ContractAddress = 0x2d55d6f311413945595788818d4e89e151360a2c2c6b5270d5d0ed16475505f.try_into().unwrap();
+
+    let address_eligible_for_zklend_rewards: ContractAddress = ADDRESS_ELIGIBLE_FOR_ZKLEND_REWARDS.try_into().unwrap();
+    let contract = declare("Deposit").unwrap().contract_class();
+    replace_bytecode(address_eligible_for_zklend_rewards, *contract.class_hash).unwrap();
+
+    let strk_balance_at_start = strk.balance_of(address_eligible_for_zklend_rewards.try_into().unwrap());
+
+    // write the treasury address so we can check funds were sent
+    let hypothetical_treasury_address = 0x98765;
+    let storage_entry_for_treasury_address = array![hypothetical_treasury_address].span();
+    store(address_eligible_for_zklend_rewards, selector!("treasury"), storage_entry_for_treasury_address);
+    let ZKLEND_MARKET = 0x04c0a5193d58f74fbace4b74dcf65481e734ed1714121bdc571da345540efa05;
+    let storage_entry_for_zk_market = array![ZKLEND_MARKET].span();
+    store(address_eligible_for_zklend_rewards, selector!("zk_market"), storage_entry_for_zk_market);
+
+    let deposit_contract = IDepositDispatcher {contract_address: address_eligible_for_zklend_rewards};
+    let proof = array![
+            0x43a677604d8a532f023b8a1480e39f0a4f95460a88eb978bf86cf2e6af4a505,
+            0x69faedf42e0dccc605c8f5b773c58154bd51f1d807ce51d6a254b58379df414,
+            0x75afcd7c6775bd043279c5adf5cfc8519175ddb640d9bab3a80d6216fc434f2,
+            0x718d5326a3a934d067b4930ff2ffbc6dba50eb189ddecc50d559c74e30ce375,
+            0x60842dbbced8d585d720c3efe1f99fb32da09f6334f3ef679dbfbc9b47fcf2b,
+            0x7d22c5040360327dc761eb46959c15f888b68af777829d758150612ec13949c,
+            0x406de13ffdac0138c360921e9e51bb7fdabe9770c750157e40e04909589c0e7,
+            0x4f138575a80804622f8b92152ffaf19e634c63934348d13b92dd1eb91bfa3c,
+            0x1ca16be8f87a5dc8cbdd781a8e2e37b047ccdc0552ea048f8cbc28b6e0e9621,
+            0x6e5d1a64f19a0a541716f701413ae2bace2151b83da7b231ebb347c2be8272b,
+            0xf1e537b49ce8629386bfab3e390bb5dee770e0c8a176115b82607f9d9fa441,
+            0x517330339d2b79fff83a99bfa17d974267ff1a49fa517bda0ec6105130d15e1,
+            0x17230a4e2cb1bc3fb6a83c165e9f8f719c5198065e02d4aac7c55b75ac92fc0,
+            0x787a4ca028ae34239a07b4d023f4ef785c9aa934da05a0fd2ec1310dc1a6d83
+        ];
+    let claim: Claim = Claim {id: 11051, claimee: address_eligible_for_zklend_rewards, amount: 0x2a52c411698a729};
+    // eligible for 0x2a52c411698a729 = 190607217296713513 fri (fri is lowest denominator of strk token)
+    // treasury should get 152485773837370811 which is 80 %
+
+
+    deposit_contract.claim_reward(claim, proof.span(), defispring_claim_contract);
+
+
+    let fri_in_treasury = strk.balance_of(hypothetical_treasury_address.try_into().unwrap());
+    assert(fri_in_treasury == 152485773837370811, 'incorrect amount in treasury');
+    let strk_left_in_contract = strk.balance_of(address_eligible_for_zklend_rewards.try_into().unwrap());
+    assert!(strk_left_in_contract == strk_balance_at_start, "strk left in contract after airdrop claim");
+}
+
+#[test]
+#[fork("MAINNET_FIXED_BLOCK")]
+fn test_claim_and_withdraw() {
+    test_claim_as_keeper();
+    let address_eligible_for_zklend_rewards: ContractAddress = ADDRESS_ELIGIBLE_FOR_ZKLEND_REWARDS.try_into().unwrap();
+    let deposit_contract = IDepositDispatcher {contract_address: address_eligible_for_zklend_rewards};
+    let strk = ERC20ABIDispatcher { contract_address: STRK_ADDRESS.try_into().unwrap() };
+    let hypothetical_owner_address: ContractAddress = HYPOTHETICAL_OWNER_ADDR.try_into().unwrap();
+    let storage_entry_for_hypothetical_owner= array![HYPOTHETICAL_OWNER_ADDR].span();
+    store(address_eligible_for_zklend_rewards, selector!("owner"), storage_entry_for_hypothetical_owner);
+
+    cheat_caller_address(address_eligible_for_zklend_rewards, hypothetical_owner_address, CheatSpan::TargetCalls(1));
+    deposit_contract.withdraw(STRK_ADDRESS.try_into().unwrap(), 0); //passing 0 to withdraw all
+
+    assert(strk.balance_of(hypothetical_owner_address) != 0, 'no strk sent on to user');
+
+}