Paymaster audit fixes

contracts/paymaster/Cargo.toml

@@ -15,9 +15,14 @@ num-bigint = "0.4.2"
 
 [dependencies.multiversx-sc]
 version = "0.44.0"
+features = ["back-transfers"]
+
 
 [dev-dependencies.multiversx-sc-scenario]
 version = "0.44.0"
 
 [dev-dependencies.adder]
 path = "../adder"
+
+[dev-dependencies.multiversx-wegld-swap-sc]
+path = "../wegld-swap"

contracts/paymaster/src/forward_call.rs

@@ -23,28 +23,35 @@ pub trait ForwardCall {
             .with_callback(self.callbacks().transfer_callback(original_caller))
             .call_and_exit();
     }
-    
+
     #[callback]
     fn transfer_callback(
         &self,
         original_caller: ManagedAddress,
         #[call_result] result: ManagedAsyncCallResult<MultiValueEncoded<ManagedBuffer>>,
     ) -> MultiValueEncoded<ManagedBuffer> {
-        let initial_payments = self.call_value().all_esdt_transfers();
+        // TODO: use ManagedGetBackTransfers once rc1.6 is activated
+        let back_transfers = self.blockchain().get_back_transfers();
+
+        // Send the original input tokens back to the original caller
+        if !back_transfers.esdt_payments.is_empty() {
+            self.send()
+                .direct_multi(&original_caller, &back_transfers.esdt_payments);
+        }
+        if back_transfers.total_egld_amount != BigUint::zero() {
+            self.send()
+                .direct_egld(&original_caller, &back_transfers.total_egld_amount)
+        }
 
         match result {
-            ManagedAsyncCallResult::Ok(return_values) => return_values,
+            ManagedAsyncCallResult::Ok(return_values) => {
+                // Send the resulted tokens to the original caller
+                return_values
+            }
             ManagedAsyncCallResult::Err(err) => {
-                if !initial_payments.is_empty() {
-                    self.send()
-                        .direct_multi(&original_caller, &initial_payments);
-                }
-
                 let mut err_result = MultiValueEncoded::new();
                 err_result.push(ManagedBuffer::new_from_bytes(ERR_CALLBACK_MSG));
-                err_result.push(err.err_msg.clone());
-
-                sc_print!("{}", err.err_msg);
+                err_result.push(err.err_msg);
 
                 err_result
             }

contracts/paymaster/src/paymaster.rs

@@ -5,7 +5,6 @@ multiversx_sc::imports!();
 pub mod forward_call;
 const FEE_PAYMENT: usize = 0;
 
-/// An empty contract. To be used as a template when starting a new contract from scratch.
 #[multiversx_sc::contract]
 pub trait PaymasterContract: forward_call::ForwardCall {
     #[init]

contracts/paymaster/tests/paymaster_blackbox_test.rs

@@ -1,7 +1,7 @@
 use multiversx_sc::{
     codec::{multi_types::MultiValueVec, top_encode_to_vec_u8_or_panic},
     storage::mappers::SingleValue,
-    types::{Address, BigUint},
+    types::{Address, BigUint, MultiValueEncoded},
 };
 use multiversx_sc_scenario::{
     api::StaticApi,
@@ -13,35 +13,42 @@ use multiversx_sc_scenario::{
 };
 
 use adder::ProxyTrait as _;
+use multiversx_wegld_swap_sc::ProxyTrait as _;
 use paymaster::ProxyTrait as _;
 
 const PAYMASTER_ADDRESS_EXPR: &str = "sc:paymaster";
 const RELAYER_ADDRESS_EXPR: &str = "address:relayer";
 const CALLEE_SC_ADDER_ADDRESS_EXPR: &str = "sc:adder";
+const CALLEE_SC_WEGLD_ADDRESS_EXPR: &str = "sc:wegld";
 const PAYMASTER_PATH_EXPR: &str = "file:output/paymaster.wasm";
 const ADDER_PATH_EXPR: &str = "file:tests/test-contracts/adder.wasm";
+const WEGLD_PATH_EXPR: &str = "file:tests/test-contracts/multiversx-wegld-swap-sc.wasm.wasm";
 const CALLER_ADDRESS_EXPR: &str = "address:caller";
 const CALLEE_USER_ADDRESS_EXPR: &str = "address:callee_user";
 const OWNER_ADDRESS_EXPR: &str = "address:owner";
 const BALANCE: &str = "100,000,000";
 const PAYMASTER_TOKEN_ID_EXPR: &str = "str:PAYMSTR-123456";
+const WEGLD_TOKEN_ID_EXPR: &str = "str:WEGLD-123456";
+const WEGLD_TOKEN_ID: &[u8] = b"WEGLD-123456";
 const FEE_TOKEN_ID_EXPR: &str = "str:FEE-123456";
 const ADDITIONAL_TOKEN_ID_EXPR: &str = "str:ADDIT-123456";
 const FEE_AMOUNT: &str = "20,000";
 const INITIAL_ADD_VALUE: u64 = 5;
 const ADDITIONAL_ADD_VALUE: u64 = 5;
-
-
+const UNWRAP_ENDPOINT_NAME: &[u8] = b"unwrap";
 
 type PaymasterContract = ContractInfo<paymaster::Proxy<StaticApi>>;
 type AdderContract = ContractInfo<adder::Proxy<StaticApi>>;
+type WegldContract = ContractInfo<multiversx_wegld_swap_sc::Proxy<StaticApi>>;
+
 
 fn world() -> ScenarioWorld {
     let mut blockchain = ScenarioWorld::new();
     blockchain.set_current_dir_from_workspace("contracts/paymaster");
 
     blockchain.register_contract(PAYMASTER_PATH_EXPR, paymaster::ContractBuilder);
     blockchain.register_contract(ADDER_PATH_EXPR, adder::ContractBuilder);
+    blockchain.register_contract(WEGLD_PATH_EXPR, multiversx_wegld_swap_sc::ContractBuilder);
 
     blockchain
 }
@@ -52,6 +59,7 @@ struct PaymasterTestState {
     paymaster_contract: PaymasterContract,
     relayer_address: Address,
     callee_sc_adder_contract: AdderContract,
+    callee_sc_wegld_address: WegldContract,
 }
 
 impl PaymasterTestState {
@@ -66,6 +74,7 @@ impl PaymasterTestState {
                         .nonce(1)
                         .balance(BALANCE)
                         .esdt_balance(PAYMASTER_TOKEN_ID_EXPR, BALANCE)
+                        .esdt_balance(WEGLD_TOKEN_ID_EXPR, BALANCE)
                         .esdt_balance(FEE_TOKEN_ID_EXPR, BALANCE)
                         .esdt_balance(ADDITIONAL_TOKEN_ID_EXPR, BALANCE),
                 )
@@ -81,14 +90,15 @@ impl PaymasterTestState {
         let relayer_address = AddressValue::from(RELAYER_ADDRESS_EXPR).to_address();
         let paymaster_contract = PaymasterContract::new(PAYMASTER_ADDRESS_EXPR);
         let callee_sc_adder_contract = AdderContract::new(CALLEE_SC_ADDER_ADDRESS_EXPR);
-        // let callee_sc_adder_address = AddressValue::from(CALLEE_SC_ADDER_ADDRESS_EXPR).to_address();
+        let callee_sc_wegld_address = WegldContract::new(CALLEE_SC_WEGLD_ADDRESS_EXPR);
 
         Self {
             world,
             callee_user_address,
             paymaster_contract,
             relayer_address,
             callee_sc_adder_contract,
+            callee_sc_wegld_address,
         }
     }
 
@@ -130,6 +140,25 @@ impl PaymasterTestState {
         self
     }
 
+    fn deploy_wegld_contract(&mut self) -> &mut Self {
+        let wegld_code = self.world.code_expression(WEGLD_PATH_EXPR);
+
+        self.world
+            .set_state_step(SetStateStep::new().new_address(
+                OWNER_ADDRESS_EXPR,
+                3,
+                CALLEE_SC_WEGLD_ADDRESS_EXPR,
+            ))
+            .sc_deploy(
+                ScDeployStep::new()
+                    .from(OWNER_ADDRESS_EXPR)
+                    .code(wegld_code)
+                    .call(self.callee_sc_wegld_address.init(WEGLD_TOKEN_ID)),
+            );
+
+        self
+    }
+
     fn check_esdt_balance(
         &mut self,
         address_expr: &str,
@@ -144,13 +173,27 @@ impl PaymasterTestState {
 
         self
     }
+    fn check_egld_balance(
+        &mut self,
+        address_expr: &str,
+        balance_expr: &str,
+    ) -> &mut Self {
+        self.world
+            .check_state_step(CheckStateStep::new().put_account(
+                address_expr,
+                CheckAccount::new().balance(balance_expr),
+            ));
+
+        self
+    }
 }
 
 #[test]
 fn test_deploy_paymasters() {
     let mut state = PaymasterTestState::new();
     state.deploy_paymaster_contract();
     state.deploy_adder_contract();
+    state.deploy_wegld_contract();
 }
 
 #[test]
@@ -164,7 +207,7 @@ fn test_forward_call_no_fee_payment() {
             .call(state.paymaster_contract.forward_execution(
                 state.relayer_address.clone(),
                 state.callee_user_address.clone(),
-                b"add"  ,
+                b"add",
                 MultiValueVec::<Vec<u8>>::new(),
             ))
             .expect(TxExpect::user_error("str:There is no fee for payment!")),
@@ -187,7 +230,7 @@ fn test_forward_call_user() {
                     b"add",
                     MultiValueVec::<Vec<u8>>::new(),
                 ))
-                .esdt_transfer(FEE_TOKEN_ID_EXPR, 0, FEE_AMOUNT)
+                .esdt_transfer(FEE_TOKEN_ID_EXPR, 0, FEE_AMOUNT),
         )
         .check_state_step(CheckStateStep::new().put_account(
             RELAYER_ADDRESS_EXPR,
@@ -214,7 +257,7 @@ fn test_forward_call_sc_adder() {
                 state.callee_sc_adder_contract.to_address(),
                 b"add",
                 MultiValueVec::from([top_encode_to_vec_u8_or_panic(&ADDITIONAL_ADD_VALUE)]),
-            ))
+            )),
     );
 
     let expected_adder_sum = INITIAL_ADD_VALUE + ADDITIONAL_ADD_VALUE;
@@ -231,7 +274,6 @@ fn test_forward_call_sc_adder() {
     );
 }
 
-
 #[test]
 fn test_forward_call_sc_adder_multiple_payments() {
     let mut state = PaymasterTestState::new();
@@ -252,7 +294,7 @@ fn test_forward_call_sc_adder_multiple_payments() {
                 state.callee_sc_adder_contract.to_address(),
                 b"add",
                 MultiValueVec::from([top_encode_to_vec_u8_or_panic(&ADDITIONAL_ADD_VALUE)]),
-            ))
+            )),
     );
 
     let expected_adder_sum = INITIAL_ADD_VALUE + ADDITIONAL_ADD_VALUE;
@@ -273,3 +315,70 @@ fn test_forward_call_sc_adder_multiple_payments() {
         FEE_AMOUNT,
     );
 }
+
+#[test]
+fn test_forward_call_wegld() {
+    let mut state = PaymasterTestState::new();
+    state.deploy_paymaster_contract();
+    state.deploy_adder_contract();
+
+    state.check_esdt_balance(CALLER_ADDRESS_EXPR, FEE_TOKEN_ID_EXPR, BALANCE);
+    state.check_esdt_balance(CALLER_ADDRESS_EXPR, WEGLD_TOKEN_ID_EXPR, BALANCE);
+
+    // Call fails because unwrap amount is 0
+    state.world.sc_call(
+        ScCallStep::new()
+            .from(CALLER_ADDRESS_EXPR)
+            .esdt_transfer(FEE_TOKEN_ID_EXPR, 0, FEE_AMOUNT)
+            .esdt_transfer(WEGLD_TOKEN_ID_EXPR, 0, BALANCE)
+            .call(state.paymaster_contract.forward_execution(
+                state.relayer_address.clone(),
+                state.callee_sc_wegld_address.to_address(),
+                UNWRAP_ENDPOINT_NAME,
+                MultiValueEncoded::new(),
+            ))
+    );
+
+    // Fee is kept by the relayer
+    let new_fee_amount: &str =  "99980000";
+    state.check_esdt_balance(RELAYER_ADDRESS_EXPR, FEE_TOKEN_ID_EXPR, FEE_AMOUNT);
+    state.check_esdt_balance(CALLER_ADDRESS_EXPR, FEE_TOKEN_ID_EXPR, new_fee_amount);
+
+    // Caller has the original balance
+    state.check_egld_balance(CALLER_ADDRESS_EXPR, BALANCE);
+}
+
+
+#[test]
+fn test_forward_call_fails_wegld_0_amount() {
+    let mut state = PaymasterTestState::new();
+    state.deploy_paymaster_contract();
+    state.deploy_adder_contract();
+
+    state.check_esdt_balance(CALLER_ADDRESS_EXPR, FEE_TOKEN_ID_EXPR, BALANCE);
+    state.check_esdt_balance(CALLER_ADDRESS_EXPR, WEGLD_TOKEN_ID_EXPR, BALANCE);
+
+    let failling_amount = 0u64;
+
+    // Call fails because unwrap amount is 0
+    state.world.sc_call(
+        ScCallStep::new()
+            .from(CALLER_ADDRESS_EXPR)
+            .esdt_transfer(FEE_TOKEN_ID_EXPR, 0, FEE_AMOUNT)
+            .esdt_transfer(WEGLD_TOKEN_ID_EXPR, 0, failling_amount)
+            .call(state.paymaster_contract.forward_execution(
+                state.relayer_address.clone(),
+                state.callee_sc_wegld_address.to_address(),
+                UNWRAP_ENDPOINT_NAME,
+                MultiValueEncoded::new(),
+            ))
+    );
+
+    // Fee is kept by the relayer
+    let new_fee_amount: &str =  "99980000";
+    state.check_esdt_balance(RELAYER_ADDRESS_EXPR, FEE_TOKEN_ID_EXPR, FEE_AMOUNT);
+    state.check_esdt_balance(CALLER_ADDRESS_EXPR, FEE_TOKEN_ID_EXPR, new_fee_amount);
+
+    // Caller has the original balance
+    state.check_esdt_balance(CALLER_ADDRESS_EXPR, WEGLD_TOKEN_ID_EXPR, BALANCE);
+}