From 2d572c3a7df0fe24e3ce82658da93ecbdd90e351 Mon Sep 17 00:00:00 2001
From: Dusan Maksimovic <dusan.maksimovic@ethernal.tech>
Date: Tue, 16 Jul 2024 12:49:29 +0200
Subject: [PATCH] Further simplify funds validation during lock_tokens()

---
 contracts/atom_wars/src/contract.rs | 26 +++++---------------------
 contracts/atom_wars/src/error.rs    |  6 +++++-
 2 files changed, 10 insertions(+), 22 deletions(-)

diff --git a/contracts/atom_wars/src/contract.rs b/contracts/atom_wars/src/contract.rs
index 82c1519..0198c70 100644
--- a/contracts/atom_wars/src/contract.rs
+++ b/contracts/atom_wars/src/contract.rs
@@ -143,24 +143,9 @@ fn lock_tokens(
     lock_duration: u64,
 ) -> Result<Response, ContractError> {
     let constants = CONSTANTS.load(deps.storage)?;
-    validate_lock_duration(constants.lock_epoch_length, lock_duration)?;
 
-    // Validate that sent funds are the required denom
-    if info.funds.len() != 1 {
-        return Err(ContractError::Std(StdError::generic_err(
-            "Must send exactly one coin",
-        )));
-    }
-
-    let sent_funds = info.funds[0].clone();
-
-    // validate that the sent funds are the required denom using must_pay
-    must_pay(&info, &constants.denom).map_err(|_| {
-        ContractError::Std(StdError::generic_err(format!(
-            "Must send {} tokens",
-            constants.denom
-        )))
-    })?;
+    validate_lock_duration(constants.lock_epoch_length, lock_duration)?;
+    must_pay(&info, &constants.denom)?;
 
     // validate that the user does not have too many locks
     if get_lock_count(deps.as_ref(), info.sender.clone()) >= MAX_LOCK_ENTRIES {
@@ -170,9 +155,8 @@ fn lock_tokens(
         ))));
     }
 
-    // Create entry in LocksMap
     let lock_entry = LockEntry {
-        funds: sent_funds,
+        funds: info.funds[0].clone(),
         lock_start: env.block.time,
         lock_end: env.block.time.plus_nanos(lock_duration),
     };
@@ -600,7 +584,7 @@ fn add_to_whitelist(
     // Validate that the sender is a whitelist admin
     let whitelist_admins = WHITELIST_ADMINS.load(deps.storage)?;
     if !whitelist_admins.contains(&info.sender) {
-        return Err(ContractError::Unauthorized {});
+        return Err(ContractError::Unauthorized);
     }
 
     // Add covenant_params to whitelist
@@ -628,7 +612,7 @@ fn remove_from_whitelist(
     // Validate that the sender is a whitelist admin
     let whitelist_admins = WHITELIST_ADMINS.load(deps.storage)?;
     if !whitelist_admins.contains(&info.sender) {
-        return Err(ContractError::Unauthorized {});
+        return Err(ContractError::Unauthorized);
     }
 
     // Remove covenant_params from whitelist
diff --git a/contracts/atom_wars/src/error.rs b/contracts/atom_wars/src/error.rs
index 5c87b31..aa81a8e 100644
--- a/contracts/atom_wars/src/error.rs
+++ b/contracts/atom_wars/src/error.rs
@@ -1,4 +1,5 @@
 use cosmwasm_std::{OverflowError, StdError};
+use cw_utils::PaymentError;
 use thiserror::Error;
 
 #[derive(Error, Debug)]
@@ -10,5 +11,8 @@ pub enum ContractError {
     OverflowError(#[from] OverflowError),
 
     #[error("Unauthorized")]
-    Unauthorized {},
+    Unauthorized,
+
+    #[error("{0}")]
+    PaymentError(#[from] PaymentError),
 }