paymaster: Audit fixes

contracts/paymaster/interactor/src/interactor_main.rs

@@ -162,7 +162,13 @@ impl ContractInteract {
             .to(self.state.current_address())
             .gas(30_000_000u64)
             .typed(proxy::PaymasterContractProxy)
-            .forward_execution(relayer_addr, dest, endpoint_name, endpoint_args)
+            .forward_execution(
+                relayer_addr,
+                dest,
+                1_000_000u64,
+                endpoint_name,
+                endpoint_args,
+            )
             .payment(payments)
             .returns(ReturnsResultUnmanaged)
             .prepare_async()

contracts/paymaster/interactor/src/proxy.rs

@@ -53,6 +53,25 @@ where
     }
 }
 
+#[rustfmt::skip]
+impl<Env, From, To, Gas> PaymasterContractProxyMethods<Env, From, To, Gas>
+where
+    Env: TxEnv,
+    Env::Api: VMApi,
+    From: TxFrom<Env>,
+    To: TxTo<Env>,
+    Gas: TxGas<Env>,
+{
+    pub fn upgrade(
+        self,
+    ) -> TxTypedUpgrade<Env, From, To, NotPayable, Gas, ()> {
+        self.wrapped_tx
+            .payment(NotPayable)
+            .raw_upgrade()
+            .original_result()
+    }
+}
+
 #[rustfmt::skip]
 impl<Env, From, To, Gas> PaymasterContractProxyMethods<Env, From, To, Gas>
 where
@@ -65,19 +84,22 @@ where
     pub fn forward_execution<
         Arg0: ProxyArg<ManagedAddress<Env::Api>>,
         Arg1: ProxyArg<ManagedAddress<Env::Api>>,
-        Arg2: ProxyArg<ManagedBuffer<Env::Api>>,
-        Arg3: ProxyArg<MultiValueEncoded<Env::Api, ManagedBuffer<Env::Api>>>,
+        Arg2: ProxyArg<u64>,
+        Arg3: ProxyArg<ManagedBuffer<Env::Api>>,
+        Arg4: ProxyArg<MultiValueEncoded<Env::Api, ManagedBuffer<Env::Api>>>,
     >(
         self,
         relayer_addr: Arg0,
         dest: Arg1,
-        endpoint_name: Arg2,
-        endpoint_args: Arg3,
+        min_gas_limit: Arg2,
+        endpoint_name: Arg3,
+        endpoint_args: Arg4,
     ) -> TxTypedCall<Env, From, To, (), Gas, ()> {
         self.wrapped_tx
             .raw_call("forwardExecution")
             .argument(&relayer_addr)
             .argument(&dest)
+            .argument(&min_gas_limit)
             .argument(&endpoint_name)
             .argument(&endpoint_args)
             .original_result()

contracts/paymaster/src/forward_call.rs

@@ -9,10 +9,12 @@ pub trait ForwardCall {
     fn forward_call(
         &self,
         dest: ManagedAddress,
+        min_gas_limit: u64,
         endpoint_name: ManagedBuffer,
         payments: PaymentsVec<Self::Api>,
         endpoint_args: MultiValueEncoded<ManagedBuffer>,
     ) {
+        self.require_min_gas_limit(min_gas_limit);
         let original_caller = self.blockchain().get_caller();
 
         self.tx()
@@ -30,7 +32,6 @@ pub trait ForwardCall {
         original_caller: ManagedAddress,
         #[call_result] result: ManagedAsyncCallResult<MultiValueEncoded<ManagedBuffer>>,
     ) -> MultiValueEncoded<ManagedBuffer> {
-        // 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
@@ -61,4 +62,12 @@ pub trait ForwardCall {
             }
         }
     }
+
+    fn require_min_gas_limit(&self, min_gas_limit: u64) {
+        let gas_left = self.blockchain().get_gas_left();
+        require!(
+            gas_left >= min_gas_limit,
+            "Minimum required gas not provided"
+        );
+    }
 }

contracts/paymaster/src/paymaster.rs

@@ -4,38 +4,59 @@ use multiversx_sc::imports::*;
 
 pub mod forward_call;
 pub mod paymaster_proxy;
-const FEE_PAYMENT: usize = 0;
+const FEE_PAYMENT_INDEX: usize = 0;
 
 #[multiversx_sc::contract]
 pub trait PaymasterContract: forward_call::ForwardCall {
     #[init]
     fn init(&self) {}
 
+    #[upgrade]
+    fn upgrade(&self) {}
+
     #[endpoint(forwardExecution)]
     #[payable("*")]
     fn forward_execution(
         &self,
         relayer_addr: ManagedAddress,
         dest: ManagedAddress,
+        min_gas_limit: u64,
         endpoint_name: ManagedBuffer,
         endpoint_args: MultiValueEncoded<ManagedBuffer>,
     ) {
+        self.require_dest_same_shard(&dest);
         let payments = self.call_value().all_esdt_transfers();
         require!(!payments.is_empty(), "There is no fee for payment!");
 
-        let fee_payment = payments.get(FEE_PAYMENT);
+        let fee_payment = payments.get(FEE_PAYMENT_INDEX);
         self.tx()
             .to(&relayer_addr)
             .payment(EsdtTokenPayment::new(
                 fee_payment.token_identifier,
-                0,
+                fee_payment.token_nonce,
                 fee_payment.amount,
             ))
             .transfer();
 
         let mut payments_without_fee = payments.clone_value();
-        payments_without_fee.remove(FEE_PAYMENT);
+        payments_without_fee.remove(FEE_PAYMENT_INDEX);
+
+        self.forward_call(
+            dest,
+            min_gas_limit,
+            endpoint_name,
+            payments_without_fee,
+            endpoint_args,
+        );
+    }
 
-        self.forward_call(dest, endpoint_name, payments_without_fee, endpoint_args);
+    fn require_dest_same_shard(&self, dest: &ManagedAddress) {
+        let own_sc_address = self.blockchain().get_sc_address();
+        let own_shard = self.blockchain().get_shard_of_address(&own_sc_address);
+        let dest_shard = self.blockchain().get_shard_of_address(dest);
+        require!(
+            own_shard == dest_shard,
+            "Destination must be in the same shard"
+        );
     }
 }

contracts/paymaster/src/paymaster_proxy.rs

@@ -53,6 +53,25 @@ where
     }
 }
 
+#[rustfmt::skip]
+impl<Env, From, To, Gas> PaymasterContractProxyMethods<Env, From, To, Gas>
+where
+    Env: TxEnv,
+    Env::Api: VMApi,
+    From: TxFrom<Env>,
+    To: TxTo<Env>,
+    Gas: TxGas<Env>,
+{
+    pub fn upgrade(
+        self,
+    ) -> TxTypedUpgrade<Env, From, To, NotPayable, Gas, ()> {
+        self.wrapped_tx
+            .payment(NotPayable)
+            .raw_upgrade()
+            .original_result()
+    }
+}
+
 #[rustfmt::skip]
 impl<Env, From, To, Gas> PaymasterContractProxyMethods<Env, From, To, Gas>
 where
@@ -65,19 +84,22 @@ where
     pub fn forward_execution<
         Arg0: ProxyArg<ManagedAddress<Env::Api>>,
         Arg1: ProxyArg<ManagedAddress<Env::Api>>,
-        Arg2: ProxyArg<ManagedBuffer<Env::Api>>,
-        Arg3: ProxyArg<MultiValueEncoded<Env::Api, ManagedBuffer<Env::Api>>>,
+        Arg2: ProxyArg<u64>,
+        Arg3: ProxyArg<ManagedBuffer<Env::Api>>,
+        Arg4: ProxyArg<MultiValueEncoded<Env::Api, ManagedBuffer<Env::Api>>>,
     >(
         self,
         relayer_addr: Arg0,
         dest: Arg1,
-        endpoint_name: Arg2,
-        endpoint_args: Arg3,
+        min_gas_limit: Arg2,
+        endpoint_name: Arg3,
+        endpoint_args: Arg4,
     ) -> TxTypedCall<Env, From, To, (), Gas, ()> {
         self.wrapped_tx
             .raw_call("forwardExecution")
             .argument(&relayer_addr)
             .argument(&dest)
+            .argument(&min_gas_limit)
             .argument(&endpoint_name)
             .argument(&endpoint_args)
             .original_result()

contracts/paymaster/tests/paymaster_blackbox_test.rs

@@ -160,6 +160,7 @@ fn test_forward_call_no_fee_payment() {
         .forward_execution(
             RELAYER_ADDRESS_EXPR,
             CALLEE_USER_ADDRESS_EXPR,
+            0u64,
             b"add",
             MultiValueVec::<Vec<u8>>::new(),
         )
@@ -181,6 +182,7 @@ fn test_forward_call_user() {
         .forward_execution(
             RELAYER_ADDRESS_EXPR,
             CALLEE_USER_ADDRESS_EXPR,
+            0u64,
             b"add",
             MultiValueVec::<Vec<u8>>::new(),
         )
@@ -215,6 +217,7 @@ fn test_forward_call_sc_adder() {
         .forward_execution(
             RELAYER_ADDRESS_EXPR,
             CALLEE_SC_ADDER_ADDRESS_EXPR,
+            0u64,
             b"add",
             MultiValueVec::from([top_encode_to_vec_u8_or_panic(&ADDITIONAL_ADD_VALUE)]),
         )
@@ -256,6 +259,7 @@ fn test_forward_call_sc_adder_with_relayer_address() {
         .forward_execution(
             RELAYER_ADDRESS_EXPR,
             CALLEE_SC_ADDER_ADDRESS_EXPR,
+            0u64,
             b"add",
             MultiValueVec::from([top_encode_to_vec_u8_or_panic(&ADDITIONAL_ADD_VALUE)]),
         )
@@ -310,6 +314,7 @@ fn test_forward_call_wegld() {
         .forward_execution(
             RELAYER_ADDRESS_EXPR,
             CALLEE_SC_WEGLD_ADDRESS_EXPR,
+            0u64,
             UNWRAP_ENDPOINT_NAME,
             MultiValueEncoded::new(),
         )
@@ -358,6 +363,7 @@ fn test_forward_call_fails_wegld_0_amount() {
         .forward_execution(
             RELAYER_ADDRESS_EXPR,
             CALLEE_SC_WEGLD_ADDRESS_EXPR,
+            0u64,
             UNWRAP_ENDPOINT_NAME,
             MultiValueEncoded::new(),
         )

contracts/paymaster/wasm/src/lib.rs

@@ -5,9 +5,10 @@
 ////////////////////////////////////////////////////
 
 // Init:                                 1
+// Upgrade:                              1
 // Endpoints:                            1
 // Async Callback:                       1
-// Total number of exported functions:   3
+// Total number of exported functions:   4
 
 #![no_std]
 
@@ -18,6 +19,7 @@ multiversx_sc_wasm_adapter::endpoints! {
     paymaster
     (
         init => init
+        upgrade => upgrade
         forwardExecution => forward_execution
     )
 }