diff --git a/src/dispense_tracker.rs b/src/dispense_tracker.rs index e405fef..1307573 100644 --- a/src/dispense_tracker.rs +++ b/src/dispense_tracker.rs @@ -94,7 +94,7 @@ impl DispenseTracker { } } - pub fn is_tracked(&self, address: &Address) -> bool { + pub fn has_tracked(&self, address: &Address) -> bool { self.tracked.get(address).is_some() || self.in_progress.contains(address) } } diff --git a/src/routes.rs b/src/routes.rs index 85c950c..6bcf8df 100644 --- a/src/routes.rs +++ b/src/routes.rs @@ -9,8 +9,8 @@ use axum::{ use fuel_core_client::client::FuelClient; use fuel_tx::UtxoId; -use fuel_types::{Address, AssetId}; -use fuels_accounts::{Account, Signer, ViewOnlyAccount}; +use fuel_types::{Address, AssetId, Bytes32}; +use fuels_accounts::{wallet::WalletUnlocked, Account, Signer, ViewOnlyAccount}; use fuels_core::types::transaction::{Transaction, TxPolicies}; use fuels_core::types::transaction_builders::BuildableTransaction; use fuels_core::types::{ @@ -120,20 +120,83 @@ impl IntoResponse for DispenseInfoResponse { } } -async fn check_and_mark_dispense_limit( +fn check_and_mark_dispense_limit( dispense_tracker: &SharedDispenseTracker, address: Address, interval: u64, -) -> bool { - let mut dispense_tracker = dispense_tracker.lock().unwrap(); - dispense_tracker.evict_expired_entries(interval); +) -> Result<(), DispenseError> { + let mut tracker = dispense_tracker.lock().unwrap(); + tracker.evict_expired_entries(interval); - if !dispense_tracker.is_tracked(&address) { - dispense_tracker.mark_in_progress(address); - return false; + if tracker.has_tracked(&address) { + return Err(error( + "Account has already received assets today".to_string(), + StatusCode::TOO_MANY_REQUESTS, + )); } - true + tracker.mark_in_progress(address); + Ok(()) +} + +async fn get_coin_output(wallet: &WalletUnlocked) -> Result { + let resources = wallet + .get_spendable_resources(AssetId::BASE, THE_BIGGEST_AMOUNT) + .await + .map_err(|e| { + error( + format!("Failed to get resources: {e}"), + StatusCode::INTERNAL_SERVER_ERROR, + ) + })?; + + let coin_output = resources + .into_iter() + .filter_map(|coin| match coin { + CoinType::Coin(coin) => Some(CoinOutput { + utxo_id: coin.utxo_id, + owner: coin.owner.into(), + amount: coin.amount, + }), + _ => None, + }) + .last() + .ok_or_else(|| { + error( + "The wallet is empty".to_string(), + StatusCode::INTERNAL_SERVER_ERROR, + ) + })?; + + Ok(coin_output) +} + +async fn submit_tx_with_timeout( + client: &FuelClient, + tx_id: &Bytes32, + timeout: u64, +) -> Result<(), DispenseError> { + tokio::time::timeout( + Duration::from_secs(timeout), + client.await_transaction_commit(tx_id), + ) + .await + .map(|r| { + r.map_err(|e| { + error( + format!("Failed to submit transaction with error: {e}"), + StatusCode::INTERNAL_SERVER_ERROR, + ) + }) + }) + .map_err(|e| { + error( + format!("Got a timeout during transaction submission: {e}"), + StatusCode::INTERNAL_SERVER_ERROR, + ) + })??; + + Ok(()) } #[tracing::instrument(skip(wallet, config))] @@ -158,15 +221,6 @@ pub async fn dispense_tokens( )); }?; - if check_and_mark_dispense_limit(&dispense_tracker, address, config.dispense_limit_interval) - .await - { - return Err(error( - "Account has already received assets today".to_string(), - StatusCode::TOO_MANY_REQUESTS, - )); - }; - // verify captcha if let Some(s) = config.captcha_secret.clone() { recaptcha::verify(s.expose_secret(), input.captcha.as_str(), None) @@ -180,8 +234,15 @@ pub async fn dispense_tokens( })?; } - let provider = wallet.provider().expect("client provider"); + check_and_mark_dispense_limit(&dispense_tracker, address, config.dispense_limit_interval)?; + let cleanup = || { + dispense_tracker + .lock() + .unwrap() + .remove_in_progress(&address); + }; + let provider = wallet.provider().expect("client provider"); let mut tx_id; loop { @@ -189,30 +250,10 @@ pub async fn dispense_tokens( let coin_output = if let Some(previous_coin_output) = &guard.last_output { *previous_coin_output } else { - wallet - .get_spendable_resources(AssetId::BASE, THE_BIGGEST_AMOUNT) - .await - .map_err(|e| { - dispense_tracker - .lock() - .unwrap() - .remove_in_progress(&address); - error( - format!("Failed to get resources: {e}"), - StatusCode::INTERNAL_SERVER_ERROR, - ) - })? - .into_iter() - .filter_map(|coin| match coin { - CoinType::Coin(coin) => Some(CoinOutput { - utxo_id: coin.utxo_id, - owner: coin.owner.into(), - amount: coin.amount, - }), - _ => None, - }) - .last() - .expect("The wallet is empty") + get_coin_output(&wallet).await.map_err(|e| { + cleanup(); + e + })? }; let coin_type = CoinType::Coin(Coin { @@ -286,25 +327,12 @@ pub async fn dispense_tokens( }; } - tokio::time::timeout( - Duration::from_secs(config.timeout), - client.await_transaction_commit(&tx_id), - ) - .await - .map(|r| { - r.map_err(|e| { - error( - format!("Failed to submit transaction with error: {e}"), - StatusCode::INTERNAL_SERVER_ERROR, - ) - }) - }) - .map_err(|e| { - error( - format!("Got a timeout during transaction submission: {e}"), - StatusCode::INTERNAL_SERVER_ERROR, - ) - })??; + submit_tx_with_timeout(&client, &tx_id, config.timeout) + .await + .map_err(|e| { + cleanup(); + e + })?; info!( "dispensed {} tokens to {:#x}",