diff --git a/src/lib.rs b/src/lib.rs index d88e6592..736173ae 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -25,7 +25,7 @@ pub use parse_reply::{ parse_reply_instantiate_data, MsgExecuteContractResponse, MsgInstantiateContractResponse, ParseReplyError, }; -pub use payment::{may_pay, must_pay, nonpayable, one_coin, PaymentError}; +pub use payment::{may_pay, must_pay, must_pay_many, nonpayable, one_coin, PaymentError}; pub use threshold::{Threshold, ThresholdError, ThresholdResponse}; pub use crate::balance::NativeBalance; diff --git a/src/payment.rs b/src/payment.rs index f6ccdd20..80993cc3 100644 --- a/src/payment.rs +++ b/src/payment.rs @@ -1,4 +1,5 @@ use cosmwasm_std::{Coin, MessageInfo, Uint128}; +use std::collections::HashMap; use thiserror::Error; /// returns an error if any coins were sent @@ -52,6 +53,39 @@ pub fn may_pay(info: &MessageInfo, denom: &str) -> Result } } +pub fn must_pay_many( + info: &MessageInfo, + required_denoms: &[&str], +) -> Result, PaymentError> { + if info.funds.is_empty() { + return Err(PaymentError::NoFunds {}); + } + if info.funds.len() != required_denoms.len() { + return Err(PaymentError::IncorrectNumberOfDenoms(required_denoms.len())); + } + + let mut funds_map = HashMap::new(); + for coin in &info.funds { + if funds_map.insert(coin.denom.as_str(), coin).is_some() { + return Err(PaymentError::DuplicateDenom(coin.denom.clone())); + } + } + + let mut matched_coins: Vec = Vec::new(); + for &denom in required_denoms { + match funds_map.get(denom) { + Some(coin) => { + if coin.amount == Uint128::zero() { + return Err(PaymentError::ZeroAmountDenom(denom.to_string())); + } + matched_coins.push((**coin).clone()); + } + None => return Err(PaymentError::MissingDenom(denom.to_string())), + } + } + Ok(matched_coins) +} + #[derive(Error, Debug, PartialEq, Eq)] pub enum PaymentError { #[error("Must send reserve token '{0}'")] @@ -68,6 +102,15 @@ pub enum PaymentError { #[error("This message does no accept funds")] NonPayable {}, + + #[error("Must send '{0}' denoms")] + IncorrectNumberOfDenoms(usize), + + #[error("Zero amound sent for '{0}'")] + ZeroAmountDenom(String), + + #[error("Received duplicate denom '{0}'")] + DuplicateDenom(String), } #[cfg(test)] @@ -108,6 +151,36 @@ mod test { let err = may_pay(&mixed_payment, atom).unwrap_err(); assert_eq!(err, PaymentError::ExtraDenom("wei".to_string())); } + #[test] + fn test_must_pay_many_works() { + let atom: &str = "uatom"; + let eth: &str = "wei"; + + let no_payment = mock_info(SENDER, &[]); + let atom_payment = mock_info(SENDER, &coins(100, atom)); + let zero_amount_payment = mock_info(SENDER, &[coin(0, atom), coin(120, eth)]); + let duplicate_atom_payment = mock_info(SENDER, &[coin(50, atom), coin(120, atom)]); + let eth_payment = mock_info(SENDER, &coins(100, eth)); + let mixed_payment = mock_info(SENDER, &[coin(50, atom), coin(120, eth)]); + + let err = must_pay_many(&no_payment, &[atom, eth]).unwrap_err(); + assert_eq!(err, PaymentError::NoFunds {}); + + let err = must_pay_many(&atom_payment, &[atom, eth]).unwrap_err(); + assert_eq!(err, PaymentError::IncorrectNumberOfDenoms(2)); + + let err = must_pay_many(&duplicate_atom_payment, &[atom, eth]).unwrap_err(); + assert_eq!(err, PaymentError::DuplicateDenom(atom.to_string())); + + let err = must_pay_many(&zero_amount_payment, &[atom, eth]).unwrap_err(); + assert_eq!(err, PaymentError::ZeroAmountDenom(atom.to_string())); + + let err = must_pay_many(ð_payment, &[atom]).unwrap_err(); + assert_eq!(err, PaymentError::MissingDenom(atom.to_string())); + + let res = must_pay_many(&mixed_payment, &[atom, eth]).unwrap(); + assert_eq!(res, vec![coin(50, atom), coin(120, eth)]); + } #[test] fn must_pay_works() {