Merge pull request #27 from Dharitri-org/raj

v0.2.1
Dharitri-org · Mar 22, 2024 · 3329152 · 3329152
2 parents 2299552 + 601694e
commit 3329152
Show file tree

Hide file tree

Showing 7 changed files with 238 additions and 55 deletions.
diff --git a/builtInFunctions/migrateDataTrie.go b/builtInFunctions/migrateDataTrie.go
@@ -1,6 +1,7 @@
 package builtInFunctions
 
 import (
+	"bytes"
 	"fmt"
 	"sync"
 
@@ -35,7 +36,7 @@ func NewMigrateDataTrieFunc(
 		accounts:    accounts,
 	}
 
-	mdt.baseActiveHandler.activeHandler = enableEpochsHandler.IsAutoBalanceDataTriesEnabled
+	mdt.baseActiveHandler.activeHandler = enableEpochsHandler.IsMigrateDataTrieEnabled
 
 	return mdt, nil
 }
@@ -64,9 +65,18 @@ func (mdt *migrateDataTrie) ProcessBuiltinFunction(
 		NewVersion:   core.AutoBalanceEnabled,
 		TrieMigrator: dtm,
 	}
-	err = acntDst.AccountDataHandler().MigrateDataTrieLeaves(argsMigrateDataTrie)
-	if err != nil {
-		return nil, err
+
+	shouldMigrateAcntDst := bytes.Equal(acntDst.AddressBytes(), vmcommon.SystemAccountAddress) || !vmcommon.IsSystemAccountAddress(acntDst.AddressBytes())
+	if shouldMigrateAcntDst {
+		err = acntDst.AccountDataHandler().MigrateDataTrieLeaves(argsMigrateDataTrie)
+		if err != nil {
+			return nil, err
+		}
+	} else {
+		err = mdt.migrateSystemAccount(argsMigrateDataTrie)
+		if err != nil {
+			return nil, err
+		}
 	}
 
 	vmOutput := &vmcommon.VMOutput{
@@ -77,6 +87,34 @@ func (mdt *migrateDataTrie) ProcessBuiltinFunction(
 	return vmOutput, nil
 }
 
+func (mdt *migrateDataTrie) migrateSystemAccount(argsMigrateDataTrie vmcommon.ArgsMigrateDataTrieLeaves) error {
+	account, err := mdt.getExistingAccount(vmcommon.SystemAccountAddress)
+	if err != nil {
+		return err
+	}
+
+	err = account.AccountDataHandler().MigrateDataTrieLeaves(argsMigrateDataTrie)
+	if err != nil {
+		return err
+	}
+
+	return mdt.accounts.SaveAccount(account)
+}
+
+func (mdt *migrateDataTrie) getExistingAccount(address []byte) (vmcommon.UserAccountHandler, error) {
+	account, err := mdt.accounts.GetExistingAccount(address)
+	if err != nil {
+		return nil, err
+	}
+
+	userAccount, ok := account.(vmcommon.UserAccountHandler)
+	if !ok {
+		return nil, ErrWrongTypeAssertion
+	}
+
+	return userAccount, nil
+}
+
 // SetNewGasConfig is called whenever gas cost is changed
 func (mdt *migrateDataTrie) SetNewGasConfig(gasCost *vmcommon.GasCost) {
 	if gasCost == nil {

diff --git a/builtInFunctions/migrateDataTrie_test.go b/builtInFunctions/migrateDataTrie_test.go
@@ -1,6 +1,7 @@
 package builtInFunctions
 
 import (
+	"bytes"
 	"errors"
 	"math/big"
 	"strings"
@@ -24,27 +25,25 @@ func TestNewMigrateDataTrieFunc(t *testing.T) {
 		assert.True(t, check.IfNil(mdtf))
 		assert.Equal(t, ErrNilEnableEpochsHandler, err)
 	})
-
 	t.Run("nil accountsDB", func(t *testing.T) {
 		t.Parallel()
 
 		mdtf, err := NewMigrateDataTrieFunc(vmcommon.BuiltInCost{}, &mock.EnableEpochsHandlerStub{}, nil)
 		assert.True(t, check.IfNil(mdtf))
 		assert.Equal(t, ErrNilAccountsAdapter, err)
 	})
-
 	t.Run("should work", func(t *testing.T) {
 		t.Parallel()
 
 		enableEpochs := &mock.EnableEpochsHandlerStub{
-			IsAutoBalanceDataTriesEnabledField: true,
+			IsMigrateDataTrieEnabledField: true,
 		}
 		mdtf, err := NewMigrateDataTrieFunc(vmcommon.BuiltInCost{}, enableEpochs, &mock.AccountsStub{})
 		assert.False(t, check.IfNil(mdtf))
 		assert.Nil(t, err)
 		assert.True(t, mdtf.IsActive())
 
-		enableEpochs.IsAutoBalanceDataTriesEnabledField = false
+		enableEpochs.IsMigrateDataTrieEnabledField = false
 		assert.False(t, mdtf.IsActive())
 	})
 }
@@ -60,7 +59,6 @@ func TestMigrateDataTrie_ProcessBuiltinFunction(t *testing.T) {
 		assert.Nil(t, vmOutput)
 		assert.Equal(t, ErrNilVmInput, err)
 	})
-
 	t.Run("not enough gas provided for at least one migration", func(t *testing.T) {
 		t.Parallel()
 
@@ -79,7 +77,6 @@ func TestMigrateDataTrie_ProcessBuiltinFunction(t *testing.T) {
 		assert.Nil(t, vmOutput)
 		assert.True(t, strings.Contains(err.Error(), "not enough gas"))
 	})
-
 	t.Run("invalid call value", func(t *testing.T) {
 		t.Parallel()
 
@@ -94,13 +91,28 @@ func TestMigrateDataTrie_ProcessBuiltinFunction(t *testing.T) {
 		assert.Nil(t, vmOutput)
 		assert.Equal(t, ErrBuiltInFunctionCalledWithValue, err)
 	})
+	t.Run("invalid number of arguments", func(t *testing.T) {
+		t.Parallel()
 
+		input := &vmcommon.ContractCallInput{
+			VMInput: vmcommon.VMInput{
+				CallValue: big.NewInt(0),
+				Arguments: [][]byte{[]byte("arg1")},
+			},
+		}
+
+		mdtf, _ := NewMigrateDataTrieFunc(vmcommon.BuiltInCost{}, &mock.EnableEpochsHandlerStub{}, &mock.AccountsStub{})
+		vmOutput, err := mdtf.ProcessBuiltinFunction(mock.NewUserAccount([]byte("sender")), nil, input)
+		assert.Nil(t, vmOutput)
+		assert.True(t, errors.Is(err, ErrInvalidNumberOfArguments))
+	})
 	t.Run("nil dest account", func(t *testing.T) {
 		t.Parallel()
 
 		input := &vmcommon.ContractCallInput{
 			VMInput: vmcommon.VMInput{
-				CallValue: big.NewInt(0),
+				CallValue:  big.NewInt(0),
+				CallerAddr: []byte("caller"),
 			},
 		}
 
@@ -109,6 +121,113 @@ func TestMigrateDataTrie_ProcessBuiltinFunction(t *testing.T) {
 		assert.Nil(t, vmOutput)
 		assert.True(t, errors.Is(err, ErrNilSCDestAccount))
 	})
+	t.Run("address is system account address from shard 0", func(t *testing.T) {
+		t.Parallel()
+
+		migrateCalled := false
+		saveCalled := false
+		input := &vmcommon.ContractCallInput{
+			VMInput: vmcommon.VMInput{
+				CallValue:  big.NewInt(0),
+				CallerAddr: []byte("12345678912345678912345678912345"),
+			},
+		}
+		systemAcc := &mock.UserAccountStub{
+			AccountDataHandlerCalled: func() vmcommon.AccountDataHandler {
+				return &mock.DataTrieTrackerStub{
+					MigrateDataTrieLeavesCalled: func(args vmcommon.ArgsMigrateDataTrieLeaves) error {
+						migrateCalled = true
+						return nil
+					},
+				}
+			},
+		}
+		adb := &mock.AccountsStub{
+			GetExistingAccountCalled: func(address []byte) (vmcommon.AccountHandler, error) {
+				return systemAcc, nil
+			},
+			SaveAccountCalled: func(account vmcommon.AccountHandler) error {
+				assert.Equal(t, systemAcc, account)
+				saveCalled = true
+				return nil
+			},
+		}
+
+		mdtf, _ := NewMigrateDataTrieFunc(vmcommon.BuiltInCost{}, &mock.EnableEpochsHandlerStub{}, adb)
+
+		shard0SystemAccAddr := bytes.Repeat([]byte{255}, 30)
+		shard0SystemAccAddr = append(shard0SystemAccAddr, []byte{0, 0}...)
+		vmOutput, err := mdtf.ProcessBuiltinFunction(mock.NewUserAccount([]byte("sender")), mock.NewUserAccount(shard0SystemAccAddr), input)
+		assert.Nil(t, err)
+		assert.NotNil(t, vmOutput)
+		assert.True(t, migrateCalled)
+		assert.True(t, saveCalled)
+	})
+	t.Run("address is system account address from shard 1", func(t *testing.T) {
+		t.Parallel()
+
+		migrateCalled := false
+		input := &vmcommon.ContractCallInput{
+			VMInput: vmcommon.VMInput{
+				CallValue:  big.NewInt(0),
+				CallerAddr: []byte("12345678912345678912345678912345"),
+			},
+		}
+		systemAcc := &mock.UserAccountStub{
+			AccountDataHandlerCalled: func() vmcommon.AccountDataHandler {
+				return &mock.DataTrieTrackerStub{
+					MigrateDataTrieLeavesCalled: func(args vmcommon.ArgsMigrateDataTrieLeaves) error {
+						migrateCalled = true
+						return nil
+					},
+				}
+			},
+		}
+		adb := &mock.AccountsStub{
+			GetExistingAccountCalled: func(address []byte) (vmcommon.AccountHandler, error) {
+				assert.Fail(t, "should not be called")
+				return nil, nil
+			},
+			SaveAccountCalled: func(account vmcommon.AccountHandler) error {
+				assert.Fail(t, "should not be called")
+				return nil
+			},
+		}
+
+		mdtf, _ := NewMigrateDataTrieFunc(vmcommon.BuiltInCost{}, &mock.EnableEpochsHandlerStub{}, adb)
+
+		vmOutput, err := mdtf.ProcessBuiltinFunction(mock.NewUserAccount([]byte("sender")), systemAcc, input)
+		assert.Nil(t, err)
+		assert.NotNil(t, vmOutput)
+		assert.True(t, migrateCalled)
+	})
+	t.Run("should work", func(t *testing.T) {
+		t.Parallel()
+
+		migrateCalled := false
+		input := &vmcommon.ContractCallInput{
+			VMInput: vmcommon.VMInput{
+				CallValue:  big.NewInt(0),
+				CallerAddr: []byte("sender"),
+			},
+		}
+		destAcc := &mock.UserAccountStub{
+			AccountDataHandlerCalled: func() vmcommon.AccountDataHandler {
+				return &mock.DataTrieTrackerStub{
+					MigrateDataTrieLeavesCalled: func(args vmcommon.ArgsMigrateDataTrieLeaves) error {
+						migrateCalled = true
+						return nil
+					},
+				}
+			},
+		}
+
+		mdtf, _ := NewMigrateDataTrieFunc(vmcommon.BuiltInCost{}, &mock.EnableEpochsHandlerStub{}, &mock.AccountsStub{})
+		vmOutput, err := mdtf.ProcessBuiltinFunction(mock.NewUserAccount([]byte("sender")), destAcc, input)
+		assert.Nil(t, err)
+		assert.NotNil(t, vmOutput)
+		assert.True(t, migrateCalled)
+	})
 }
 
 func TestMigrateDataTrie_SetNewGasConfig(t *testing.T) {
@@ -126,10 +245,16 @@ func TestMigrateDataTrie_Concurrency(t *testing.T) {
 		VMInput: vmcommon.VMInput{
 			CallValue:   big.NewInt(0),
 			GasProvided: 10000,
+			CallerAddr:  []byte("sender"),
+		},
+	}
+	adb := &mock.AccountsStub{
+		LoadAccountCalled: func(address []byte) (vmcommon.AccountHandler, error) {
+			return mock.NewUserAccount(address), nil
 		},
 	}
 
-	mdtf, _ := NewMigrateDataTrieFunc(vmcommon.BuiltInCost{}, &mock.EnableEpochsHandlerStub{}, &mock.AccountsStub{})
+	mdtf, _ := NewMigrateDataTrieFunc(vmcommon.BuiltInCost{}, &mock.EnableEpochsHandlerStub{}, adb)
 	numOperations := 1000
 
 	wg := sync.WaitGroup{}