@@ -8,10 +8,11 @@ mod Deposit {
88 types :: {i129 :: i129, keys :: PoolKey }
99 };
1010
11- use openzeppelin_security :: ReentrancyGuardComponent ;
12-
13- use openzeppelin_token :: erc20 :: interface :: {ERC20ABIDispatcher , ERC20ABIDispatcherTrait };
14- use openzeppelin_upgrades :: {UpgradeableComponent , interface :: IUpgradeable };
11+ use openzeppelin :: {
12+ security :: ReentrancyGuardComponent ,
13+ token :: erc20 :: interface :: {ERC20ABIDispatcher , ERC20ABIDispatcherTrait },
14+ upgrades :: {UpgradeableComponent , interface :: IUpgradeable }, access :: ownable :: OwnableComponent
15+ };
1516 use spotnet :: {
1617 constants :: {ZK_SCALE_DECIMALS , STRK_ADDRESS },
1718 interfaces :: {
@@ -33,18 +34,25 @@ mod Deposit {
3334 path : ReentrancyGuardComponent , storage : reentrancy_guard , event : ReentrancyGuardEvent
3435 );
3536 component! (path : UpgradeableComponent , storage : upgradeable , event : UpgradeableEvent );
37+ component! (path : OwnableComponent , storage : ownable , event : OwnableEvent );
38+
39+ impl OwnableInternalImpl = OwnableComponent :: InternalImpl <ContractState >;
40+ #[abi(embed_v0)]
41+ impl OwnableTwoStepMixinImpl =
42+ OwnableComponent :: OwnableTwoStepMixinImpl <ContractState >;
3643
3744 impl ReentrancyInternalImpl = ReentrancyGuardComponent :: InternalImpl <ContractState >;
3845 impl UpgradeableInternalImpl = UpgradeableComponent :: InternalImpl <ContractState >;
3946
4047 #[storage]
4148 struct Storage {
42- owner : ContractAddress ,
4349 ekubo_core : ICoreDispatcher ,
4450 zk_market : IMarketDispatcher ,
4551 treasury : ContractAddress ,
4652 is_position_open : bool ,
4753 #[substorage(v0)]
54+ ownable : OwnableComponent :: Storage ,
55+ #[substorage(v0)]
4856 reentrancy_guard : ReentrancyGuardComponent :: Storage ,
4957 #[substorage(v0)]
5058 upgradeable : UpgradeableComponent :: Storage
@@ -59,7 +67,7 @@ mod Deposit {
5967 treasury : ContractAddress
6068 ) {
6169 assert (owner . is_non_zero (), ' Owner address is zero'
62- self . owner . write (owner );
70+ self . ownable . initializer (owner );
6371 self . ekubo_core. write (ekubo_core );
6472 self . zk_market. write (zk_market );
6573 self . treasury. write (treasury );
@@ -93,7 +101,8 @@ mod Deposit {
93101 supply_decimals : DecimalScale ,
94102 debt_decimals : DecimalScale
95103 ) -> u256 {
96- let deposited = ((total_deposited * ZK_SCALE_DECIMALS * supply_token_price . into ()) / supply_decimals . into ());
104+ let deposited = ((total_deposited * ZK_SCALE_DECIMALS * supply_token_price . into ())
105+ / supply_decimals . into ());
97106 let free_amount = (((deposited * collateral_factor . into () / ZK_SCALE_DECIMALS )
98107 * borrow_factor . into ()
99108 / ZK_SCALE_DECIMALS ))
@@ -121,11 +130,27 @@ mod Deposit {
121130 repaid_amount : TokenAmount
122131 }
123132
133+ #[derive(starknet:: Event , Drop )]
134+ struct Withdraw {
135+ token : ContractAddress ,
136+ amount : TokenAmount
137+ }
138+
139+ #[derive(starknet:: Event , Drop )]
140+ struct ExtraDeposit {
141+ token : ContractAddress ,
142+ amount : TokenAmount
143+ }
144+
124145 #[event]
125146 #[derive(Drop , starknet:: Event )]
126147 enum Event {
127148 LiquidityLooped : LiquidityLooped ,
128149 PositionClosed : PositionClosed ,
150+ Withdraw : Withdraw ,
151+ ExtraDeposit : ExtraDeposit ,
152+ #[flat]
153+ OwnableEvent : OwnableComponent :: Event ,
129154 #[flat]
130155 ReentrancyGuardEvent : ReentrancyGuardComponent :: Event ,
131156 #[flat]
@@ -165,7 +190,7 @@ mod Deposit {
165190 pool_price : TokenPrice
166191 ) {
167192 let user_account = get_tx_info (). unbox (). account_contract_address;
168- assert (user_account == self . owner . read (), ' Caller is not the owner'
193+ assert (user_account == self . ownable . owner (), ' Caller is not the owner'
169194 assert (! self . is_position_open. read (), ' Open position already exists'
170195 let DepositData { token , amount , multiplier , borrow_portion_percent } = deposit_data ;
171196 assert (
@@ -193,7 +218,7 @@ mod Deposit {
193218 (false , pool_key . token0, ekubo_limits . lower)
194219 };
195220
196- token_dispatcher . transferFrom (self . owner . read (), curr_contract_address , amount );
221+ token_dispatcher . transferFrom (self . ownable . owner (), curr_contract_address , amount );
197222 let (deposit_reserve_data , debt_reserve_data ) = (
198223 zk_market . get_reserve_data (token ), zk_market . get_reserve_data (borrowing_token )
199224 );
@@ -203,6 +228,11 @@ mod Deposit {
203228 debt_reserve_data . borrow_factor. into ()
204229 );
205230
231+ assert (
232+ deposit_reserve_data . enabled && debt_reserve_data . enabled,
233+ ' Reserves must be enabled'
234+ );
235+
206236 zk_market . enable_collateral (token );
207237
208238 token_dispatcher . approve (zk_market . contract_address, amount );
@@ -293,7 +323,7 @@ mod Deposit {
293323 debt_price : TokenPrice
294324 ) {
295325 assert (
296- get_tx_info (). unbox (). account_contract_address == self . owner . read (),
326+ get_tx_info (). unbox (). account_contract_address == self . ownable . owner (),
297327 ' Caller is not the owner'
298328 );
299329 assert (self . is_position_open. read (), ' Open position not exists'
@@ -309,6 +339,11 @@ mod Deposit {
309339 debt_reserve_data . borrow_factor. into ()
310340 );
311341
342+ assert (
343+ deposit_reserve_data . enabled && debt_reserve_data . enabled,
344+ ' Reserves must be enabled'
345+ );
346+
312347 let z_token_disp = ERC20ABIDispatcher {
313348 contract_address : deposit_reserve_data . z_token_address
314349 };
@@ -386,7 +421,7 @@ mod Deposit {
386421 zk_market . disable_collateral (supply_token );
387422 self . is_position_open. write (false );
388423 let withdrawn_amount = token_disp . balanceOf (contract_address );
389- token_disp . transfer (self . owner . read (), withdrawn_amount );
424+ token_disp . transfer (self . ownable . owner (), withdrawn_amount );
390425 self
391426 . emit (
392427 PositionClosed {
@@ -465,10 +500,11 @@ mod Deposit {
465500 token_dispatcher . approve (zk_market . contract_address, amount );
466501 zk_market . enable_collateral (token );
467502 zk_market . deposit (token , amount . try_into (). unwrap ());
503+ self . emit (ExtraDeposit { token , amount });
468504 self . reentrancy_guard. end ();
469505 }
470506
471- /// Withdraws tokens from zkLend if looped tokens are repaid
507+ /// Withdraws tokens from zkLend
472508 ///
473509 /// # Panics
474510 /// address of account that started the transaction is not equal to `owner` storage variable
@@ -477,19 +513,25 @@ mod Deposit {
477513 /// `token`: TokenAddress - token address to withdraw from zkLend
478514 /// `amount`: TokenAmount - amount to withdraw. Pass `0` to withdraw all
479515 fn withdraw (ref self : ContractState , token : ContractAddress , amount : TokenAmount ) {
480- assert ( get_caller_address () == self . owner . read (), ' Caller is not the owner '
516+ self . ownable . assert_only_owner ( );
481517 let zk_market = self . zk_market. read ();
518+
482519 let token_dispatcher = ERC20ABIDispatcher { contract_address : token };
520+ let mut withdrawn_amount = amount ;
483521 if amount == 0 {
522+ let current_contract = get_contract_address ();
523+ let initial_balance = ERC20ABIDispatcher { contract_address : token }
524+ . balanceOf (current_contract );
484525 zk_market . withdraw_all (token );
485- token_dispatcher
486- . transfer (
487- self . owner . read (), token_dispatcher . balanceOf ( get_contract_address ())
488- );
526+ let new_balance = ERC20ABIDispatcher { contract_address : token }
527+ . balanceOf ( current_contract );
528+ withdrawn_amount = new_balance - initial_balance ;
529+ token_dispatcher . transfer ( self . ownable . owner (), new_balance );
489530 } else {
490531 zk_market . withdraw (token , amount . try_into (). unwrap ());
491- token_dispatcher . transfer (self . owner . read (), amount );
532+ token_dispatcher . transfer (self . ownable . owner (), amount );
492533 };
534+ self . emit (Withdraw { token , amount : withdrawn_amount });
493535 }
494536 }
495537
@@ -520,7 +562,7 @@ mod Deposit {
520562 impl UpgradeableImpl of IUpgradeable <ContractState > {
521563 fn upgrade (ref self : ContractState , new_class_hash : ClassHash ) {
522564 // This function can only be called by the owner
523- assert ( get_caller_address () == self . owner . read (), ' Caller is not the owner '
565+ self . ownable . assert_only_owner ( );
524566
525567 self . upgradeable. upgrade (new_class_hash );
526568 }
0 commit comments