From 34750796b7a90c6982ce9803ced63e9dec40f126 Mon Sep 17 00:00:00 2001
From: robertsasu <sasu.robert@gmail.com>
Date: Thu, 23 Nov 2023 15:50:57 +0200
Subject: [PATCH] fixed back transfers on callBack.

---
 .../json/scenariosContracts_test.go           |  7 ++++++
 vmhost/contexts/managedType.go                | 24 +++++++++++++++----
 vmhost/hostCore/execution.go                  |  1 +
 vmhost/interface.go                           |  1 +
 4 files changed, 28 insertions(+), 5 deletions(-)

diff --git a/integrationTests/json/scenariosContracts_test.go b/integrationTests/json/scenariosContracts_test.go
index d1b83a964..826e630f7 100644
--- a/integrationTests/json/scenariosContracts_test.go
+++ b/integrationTests/json/scenariosContracts_test.go
@@ -119,3 +119,10 @@ func TestCAttestation(t *testing.T) {
 		Run().
 		CheckNoError()
 }
+
+func TestRustPaymaster(t *testing.T) {
+	ScenariosTest(t).
+		Folder("paymaster/scenarios").
+		Run().
+		CheckNoError()
+}
diff --git a/vmhost/contexts/managedType.go b/vmhost/contexts/managedType.go
index 37baf4ad3..edeb84c8f 100644
--- a/vmhost/contexts/managedType.go
+++ b/vmhost/contexts/managedType.go
@@ -5,6 +5,7 @@ import (
 	"crypto/elliptic"
 	"encoding/binary"
 	"errors"
+	"github.com/multiversx/mx-chain-core-go/data/vm"
 	"io"
 	basicMath "math"
 	"math/big"
@@ -139,7 +140,7 @@ func (context *managedTypesContext) InitState() {
 // PushState appends the values map to the state stack
 func (context *managedTypesContext) PushState() {
 	newBigIntState, newBigFloatState, newEcState, newmBufferState, newmMapState := context.clone()
-	newTransfers := context.cloneBackTransfers()
+	newTransfers := cloneBackTransfers(context.managedTypesValues.backTransfers)
 	context.managedTypesStack = append(context.managedTypesStack, managedTypesState{
 		bigIntValues:   newBigIntState,
 		bigFloatValues: newBigFloatState,
@@ -150,6 +151,21 @@ func (context *managedTypesContext) PushState() {
 	})
 }
 
+// PopBackTransferIfAsyncCallBack copies the back transfer from the top of the stack in case of callbacks
+func (context *managedTypesContext) PopBackTransferIfAsyncCallBack(vmInput *vmcommon.ContractCallInput) {
+	if vmInput.CallType != vm.AsynchronousCallBack {
+		return
+	}
+
+	managedTypesStackLen := len(context.managedTypesStack)
+	if managedTypesStackLen == 0 {
+		return
+	}
+
+	prevState := context.managedTypesStack[managedTypesStackLen-1]
+	context.managedTypesValues.backTransfers = cloneBackTransfers(prevState.backTransfers)
+}
+
 // PopSetActiveState removes the latest entry from the state stack and sets it as the current values map
 func (context *managedTypesContext) PopSetActiveState() {
 	managedTypesStackLen := len(context.managedTypesStack)
@@ -797,7 +813,7 @@ func (context *managedTypesContext) AddValueOnlyBackTransfer(value *big.Int) {
 
 // GetBackTransfers returns all ESDT transfers and accumulated value as well, will clean accumulated values
 func (context *managedTypesContext) GetBackTransfers() ([]*vmcommon.ESDTTransfer, *big.Int) {
-	clonedTransfers := context.cloneBackTransfers()
+	clonedTransfers := cloneBackTransfers(context.managedTypesValues.backTransfers)
 	context.managedTypesValues.backTransfers = backTransfers{
 		ESDTTransfers: make([]*vmcommon.ESDTTransfer, 0),
 		CallValue:     big.NewInt(0),
@@ -806,9 +822,7 @@ func (context *managedTypesContext) GetBackTransfers() ([]*vmcommon.ESDTTransfer
 	return clonedTransfers.ESDTTransfers, clonedTransfers.CallValue
 }
 
-func (context *managedTypesContext) cloneBackTransfers() backTransfers {
-	currentBackTransfers := context.managedTypesValues.backTransfers
-
+func cloneBackTransfers(currentBackTransfers backTransfers) backTransfers {
 	newBackTransfers := backTransfers{
 		ESDTTransfers: make([]*vmcommon.ESDTTransfer, len(currentBackTransfers.ESDTTransfers)),
 		CallValue:     big.NewInt(0).Set(currentBackTransfers.CallValue),
diff --git a/vmhost/hostCore/execution.go b/vmhost/hostCore/execution.go
index 3e21ce743..1eb04379e 100644
--- a/vmhost/hostCore/execution.go
+++ b/vmhost/hostCore/execution.go
@@ -448,6 +448,7 @@ func (host *vmHost) executeOnDestContextNoBuiltinFunction(input *vmcommon.Contra
 	managedTypes, _, metering, output, runtime, async, storage := host.GetContexts()
 	managedTypes.PushState()
 	managedTypes.InitState()
+	managedTypes.PopBackTransferIfAsyncCallBack(input)
 
 	output.PushState()
 	output.CensorVMOutput()
diff --git a/vmhost/interface.go b/vmhost/interface.go
index e7edee399..03655d904 100644
--- a/vmhost/interface.go
+++ b/vmhost/interface.go
@@ -218,6 +218,7 @@ type ManagedTypesContext interface {
 	GetBackTransfers() ([]*vmcommon.ESDTTransfer, *big.Int)
 	AddValueOnlyBackTransfer(value *big.Int)
 	AddBackTransfers(transfers []*vmcommon.ESDTTransfer)
+	PopBackTransferIfAsyncCallBack(vmInput *vmcommon.ContractCallInput)
 }
 
 // OutputContext defines the functionality needed for interacting with the output context