diff --git a/contracts/asset_manager/src/contract.rs b/contracts/asset_manager/src/contract.rs index 959f5e8..a144f2a 100644 --- a/contracts/asset_manager/src/contract.rs +++ b/contracts/asset_manager/src/contract.rs @@ -96,10 +96,8 @@ impl AssetManager { } pub fn get_rate_limit(env: Env, token_address: Address) -> Result<(u64, u32, u64, u64), ContractError> { - let data: TokenData = read_token_data(&env, token_address); - if data.period<=0 { - return Err(ContractError::TokenDoesNotExists); - } + let data: TokenData = read_token_data(&env, token_address)?; + Ok(( data.period, data.percentage, @@ -108,13 +106,14 @@ impl AssetManager { )) } - pub fn reset_limit(env: Env, token: Address) { + pub fn reset_limit(env: Env, token: Address) -> Result { let admin = read_administrator(&env); admin.require_auth(); let balance = Self::get_token_balance(&env, token.clone()); - let mut data: TokenData = read_token_data(&env, token.clone()); + let mut data: TokenData = read_token_data(&env, token.clone())?; data.current_limit = (balance * data.percentage as u128 / POINTS) as u64; write_token_data(&env, token, data); + Ok(true) } pub fn get_withdraw_limit(env: Env, token: Address) -> Result { @@ -133,7 +132,7 @@ impl AssetManager { if balance - amount < limit { return Err(ContractError::ExceedsWithdrawLimit); }; - let mut data: TokenData = read_token_data(&env, token.clone()); + let mut data: TokenData = read_token_data(&env, token.clone())?; data.current_limit = limit as u64; data.last_update = env.ledger().timestamp(); write_token_data(&env, token, data); @@ -145,7 +144,7 @@ impl AssetManager { balance: u128, token: Address, ) -> Result { - let data: TokenData = read_token_data(&env, token); + let data: TokenData = read_token_data(&env, token)?; let period: u128 = data.period as u128; let percentage: u128 = data.percentage as u128; if period == 0 { diff --git a/contracts/asset_manager/src/states.rs b/contracts/asset_manager/src/states.rs index e6fbd80..0b387cf 100644 --- a/contracts/asset_manager/src/states.rs +++ b/contracts/asset_manager/src/states.rs @@ -1,6 +1,6 @@ use soroban_sdk::{Address, Env, Vec}; -use crate::storage_types::{DataKey, TokenData}; +use crate::{errors::ContractError, storage_types::{DataKey, TokenData}}; pub(crate) const DAY_IN_LEDGERS: u32 = 17280; pub(crate) const INSTANCE_BUMP_AMOUNT: u32 = 30 * DAY_IN_LEDGERS; @@ -41,10 +41,14 @@ pub fn write_token_data(env: &Env, token_address: Address, data: TokenData) { env.storage().persistent().set(&key, &data); } -pub fn read_token_data(env: &Env, token_address: Address) -> TokenData { - let default = TokenData{percentage: 0, period: 0, last_update: 0, current_limit: 0 }; +pub fn read_token_data(env: &Env, token_address: Address) -> Result { let key = DataKey::TokenData(token_address); - env.storage().persistent().get(&key).unwrap_or(default) + let token_data: TokenData = env + .storage() + .persistent() + .get(&key) + .ok_or(ContractError::TokenDoesNotExists)?; + Ok(token_data) } pub fn write_tokens(e: &Env, token: Address) { diff --git a/contracts/asset_manager/src/tests/asset_manager_test.rs b/contracts/asset_manager/src/tests/asset_manager_test.rs index b25f231..29b2d5c 100644 --- a/contracts/asset_manager/src/tests/asset_manager_test.rs +++ b/contracts/asset_manager/src/tests/asset_manager_test.rs @@ -242,7 +242,7 @@ fn test_handle_call_message_for_withdraw_to() { let stellar_asset_client: token::StellarAssetClient = token::StellarAssetClient::new(&ctx.env, &ctx.token); stellar_asset_client.mint(&ctx.registry, &((bnusd_amount * 2) as i128)); - + client.configure_rate_limit(&ctx.token, &300, &300); let data = WithdrawTo::new( ctx.token.to_string(), ctx.withdrawer.to_string(),