core, vm: implement snapshot isolation

Initial implementation without optimisations, but it should work.

pkg/core/dao/dao.go

@@ -89,6 +89,11 @@ func (dao *Simple) GetWrapped() *Simple {
 	return d
 }
 
+// GetUnwrapped returns the underlying DAO. It does not perform changes persist.
+func (dao *Simple) GetUnwrapped() *Simple {
+	return dao.nativeCachePS
+}
+
 // GetPrivate returns a new DAO instance with another layer of private
 // MemCachedStore around the current DAO Store.
 func (dao *Simple) GetPrivate() *Simple {

pkg/core/interop/context.go

@@ -318,6 +318,35 @@ func (ic *Context) SpawnVM() *vm.VM {
 	v := vm.NewWithTrigger(ic.Trigger)
 	v.GasLimit = -1
 	v.SyscallHandler = ic.SyscallHandler
+	wrapper := func() {
+		if ic.DAO == nil {
+			return
+		}
+		ic.DAO = ic.DAO.GetPrivate()
+	}
+	committer := func() error {
+		if ic.DAO == nil {
+			return nil
+		}
+		_, err := ic.DAO.Persist()
+		if err != nil {
+			return fmt.Errorf("failed to persist changes %w", err)
+		}
+		ic.DAO = ic.DAO.GetUnwrapped()
+		return nil
+	}
+	reverter := func(ntfToRemove int) {
+		have := len(ic.Notifications)
+		if have < ntfToRemove {
+			panic(fmt.Errorf("inconsistent notifications count: should remove %d, have %d", ntfToRemove, len(ic.Notifications)))
+		}
+		ic.Notifications = ic.Notifications[:have-ntfToRemove]
+		if ic.DAO == nil {
+			return
+		}
+		ic.DAO = ic.DAO.GetUnwrapped() // Discard all changes made in this layer.
+	}
+	v.SetIsolationCallbacks(wrapper, committer, reverter)
 	ic.VM = v
 	return v
 }
@@ -388,4 +417,5 @@ func (ic *Context) AddNotification(hash util.Uint160, name string, item *stackit
 		Name:       name,
 		Item:       item,
 	})
+	ic.VM.EmitNotification()
 }

pkg/core/interop_system_neotest_test.go

@@ -2,11 +2,14 @@ package core_test
 
 import (
 	"encoding/json"
+	"fmt"
 	"math"
 	"math/big"
+	"strings"
 	"testing"
 
 	"github.com/nspcc-dev/neo-go/internal/contracts"
+	"github.com/nspcc-dev/neo-go/pkg/compiler"
 	"github.com/nspcc-dev/neo-go/pkg/config"
 	"github.com/nspcc-dev/neo-go/pkg/core/interop"
 	"github.com/nspcc-dev/neo-go/pkg/core/interop/interopnames"
@@ -19,6 +22,7 @@ import (
 	"github.com/nspcc-dev/neo-go/pkg/neotest"
 	"github.com/nspcc-dev/neo-go/pkg/neotest/chain"
 	"github.com/nspcc-dev/neo-go/pkg/smartcontract"
+	"github.com/nspcc-dev/neo-go/pkg/smartcontract/manifest"
 	"github.com/nspcc-dev/neo-go/pkg/util"
 	"github.com/nspcc-dev/neo-go/pkg/util/slice"
 	"github.com/nspcc-dev/neo-go/pkg/vm/emit"
@@ -311,3 +315,195 @@ func TestSystemContractCreateAccount_Hardfork(t *testing.T) {
 	require.True(t, tx2Standard.SystemFee < tx3Standard.SystemFee)
 	require.True(t, tx2Multisig.SystemFee < tx3Multisig.SystemFee)
 }
+
+func TestSnapshotIsolation_Exceptions(t *testing.T) {
+	bc, acc := chain.NewSingle(t)
+	e := neotest.NewExecutor(t, bc, acc, acc)
+
+	// Contract A puts value in the storage, emits notifications and panics.
+	srcA := `package contractA
+		import (
+			"github.com/nspcc-dev/neo-go/pkg/interop/contract"
+			"github.com/nspcc-dev/neo-go/pkg/interop/runtime"
+			"github.com/nspcc-dev/neo-go/pkg/interop/storage"
+		)
+		func DoAndPanic(key, value []byte, nNtf int) int { // avoid https://github.com/nspcc-dev/neo-go/issues/2509
+			c := storage.GetContext()
+			storage.Put(c, key, value)
+			for i := 0; i < nNtf; i++ {
+				runtime.Notify("NotificationFromA", i)
+			}
+			panic("panic from A")
+		}
+		func CheckA(key []byte, nNtf int) bool {
+			c := storage.GetContext()
+			value := storage.Get(c, key)
+			// If called from B, then no storage changes made by A should be visible by this moment (they have been discarded after exception handling).
+			if value != nil {
+				return false
+			}
+			notifications := runtime.GetNotifications(nil)
+			if len(notifications) != nNtf {
+				return false
+			}
+			// If called from B, then no notifications made by A should be visible by this moment (they have been discarded after exception handling).
+			for i := 0; i < len(notifications); i++ {
+				ntf := notifications[i]
+				name := string(ntf[1].([]byte))
+				if name == "NotificationFromA" {
+					return false
+				}
+			}
+			return true
+		}
+		func CheckB() bool {
+			return contract.Call(runtime.GetCallingScriptHash(), "checkStorageChanges", contract.All).(bool)
+		}`
+	ctrA := neotest.CompileSource(t, acc.ScriptHash(), strings.NewReader(srcA), &compiler.Options{
+		NoEventsCheck:      true,
+		NoPermissionsCheck: true,
+		Name:               "contractA",
+		Permissions:        []manifest.Permission{{Methods: manifest.WildStrings{Value: nil}}},
+	})
+	e.DeployContract(t, ctrA, nil)
+
+	var hashAStr string
+	for i := 0; i < util.Uint160Size; i++ {
+		hashAStr += fmt.Sprintf("%#x", ctrA.Hash[i])
+		if i != util.Uint160Size-1 {
+			hashAStr += ", "
+		}
+	}
+	// Contract B puts value in the storage, emits notifications and calls A either
+	// in try-catch block or without it. After that checks that proper notifications
+	// and storage changes are available from different contexts.
+	srcB := `package contractB
+		import (
+			"github.com/nspcc-dev/neo-go/pkg/interop"
+			"github.com/nspcc-dev/neo-go/pkg/interop/contract"
+			"github.com/nspcc-dev/neo-go/pkg/interop/runtime"
+			"github.com/nspcc-dev/neo-go/pkg/interop/storage"
+			"github.com/nspcc-dev/neo-go/pkg/interop/util"
+		)
+		var caughtKey = []byte("caught")
+		func DoAndCatch(shouldRecover bool, keyA, valueA, keyB, valueB []byte, nNtfA, nNtfB1, nNtfB2 int) {
+			if shouldRecover {
+				defer func() {
+					if r := recover(); r != nil {
+						keyA := []byte("keyA") // defer can not capture variables from outside
+						nNtfB1 := 2
+						nNtfB2 := 4
+						c := storage.GetContext()
+						storage.Put(c, caughtKey, []byte{})
+						for i := 0; i < nNtfB2; i++ {
+							runtime.Notify("NotificationFromB after panic", i)
+						}
+						// Check that storage changes and notifications made by A are reverted.
+						ok := contract.Call(interop.Hash160{` + hashAStr + `}, "checkA", contract.All, keyA, nNtfB1+nNtfB2).(bool)
+						if !ok {
+							util.Abort() // should never ABORT if snapshot isolation is correctly implemented.
+						}
+						// Check that storage changes made by B after catch are still available in current context.
+						ok = CheckStorageChanges()
+						if !ok {
+							util.Abort() // should never ABORT if snapshot isolation is correctly implemented.
+						}
+						// Check that storage changes made by B after catch are still available from the outside context.
+						ok = contract.Call(interop.Hash160{` + hashAStr + `}, "checkB", contract.All).(bool)
+						if !ok {
+							util.Abort() // should never ABORT if snapshot isolation is correctly implemented.
+						}
+					}
+				}()
+			}
+			c := storage.GetContext()
+			storage.Put(c, keyB, valueB)
+			for i := 0; i < nNtfB1; i++ {
+				runtime.Notify("NotificationFromB before panic", i)
+			}
+			contract.Call(interop.Hash160{` + hashAStr + `}, "doAndPanic", contract.All, keyA, valueA, nNtfA)
+		}
+		func CheckStorageChanges() bool {
+			c := storage.GetContext()
+			itm := storage.Get(c, caughtKey)
+			return itm != nil
+		}`
+	ctrB := neotest.CompileSource(t, acc.ScriptHash(), strings.NewReader(srcB), &compiler.Options{
+		Name:               "contractB",
+		NoEventsCheck:      true,
+		NoPermissionsCheck: true,
+		Permissions:        []manifest.Permission{{Methods: manifest.WildStrings{Value: nil}}},
+	})
+	e.DeployContract(t, ctrB, nil)
+
+	keyA := []byte("keyA")     // hard-coded in the contract code due to `defer` inability to capture variables from outside.
+	valueA := []byte("valueA") // hard-coded in the contract code
+	keyB := []byte("keyB")
+	valueB := []byte("valueB")
+	nNtfA := 3
+	nNtfBBeforePanic := 2 // hard-coded in the contract code
+	nNtfBAfterPanic := 4  // hard-coded in the contract code
+	ctrInvoker := e.NewInvoker(ctrB.Hash, e.Committee)
+
+	// Firstly, do not catch exception and check that all notifications are presented in the notifications list.
+	h := ctrInvoker.InvokeFail(t, `unhandled exception: "panic from A"`, "doAndCatch", false, keyA, valueA, keyB, valueB, nNtfA, nNtfBBeforePanic, nNtfBAfterPanic)
+	aer := e.GetTxExecResult(t, h)
+	require.Equal(t, nNtfBBeforePanic+nNtfA, len(aer.Events))
+
+	// Then catch exception thrown by A and check that only notifications/storage changes from B are saved.
+	h = ctrInvoker.Invoke(t, stackitem.Null{}, "doAndCatch", true, keyA, valueA, keyB, valueB, nNtfA, nNtfBBeforePanic, nNtfBAfterPanic)
+	aer = e.GetTxExecResult(t, h)
+	require.Equal(t, nNtfBBeforePanic+nNtfBAfterPanic, len(aer.Events))
+}
+
+// This test is written to avoid https://github.com/neo-project/neo/issues/2746.
+func TestSnapshotIsolation_CallToItself(t *testing.T) {
+	bc, acc := chain.NewSingle(t)
+	e := neotest.NewExecutor(t, bc, acc, acc)
+
+	// Contract A calls method of self and throws if storage changes made by Do are unavailable after call to it.
+	srcA := `package contractA
+		import (
+			"github.com/nspcc-dev/neo-go/pkg/interop/contract"
+			"github.com/nspcc-dev/neo-go/pkg/interop/runtime"
+			"github.com/nspcc-dev/neo-go/pkg/interop/storage"
+		)
+		var key = []byte("key")
+		func Test() {
+			contract.Call(runtime.GetExecutingScriptHash(), "callMyselfAndCheck", contract.All)
+		}
+		func CallMyselfAndCheck() {
+			contract.Call(runtime.GetExecutingScriptHash(), "do", contract.All)
+			c := storage.GetContext()
+			val := storage.Get(c, key)
+			if val == nil {
+				panic("changes from previous context were not persisted")
+			}
+		}
+		func Do() {
+			c := storage.GetContext()
+			storage.Put(c, key, []byte("value"))
+		}
+		func Check() {
+			c := storage.GetContext()
+			val := storage.Get(c, key)
+			if val == nil {
+				panic("value is nil")
+			}
+		}
+`
+	ctrA := neotest.CompileSource(t, acc.ScriptHash(), strings.NewReader(srcA), &compiler.Options{
+		NoEventsCheck:      true,
+		NoPermissionsCheck: true,
+		Name:               "contractA",
+		Permissions:        []manifest.Permission{{Methods: manifest.WildStrings{Value: nil}}},
+	})
+	e.DeployContract(t, ctrA, nil)
+
+	ctrInvoker := e.NewInvoker(ctrA.Hash, e.Committee)
+	ctrInvoker.Invoke(t, stackitem.Null{}, "test")
+
+	// A separate call is needed to check whether all VM contexts were properly
+	// unwrapped and persisted during the previous call.
+	ctrInvoker.Invoke(t, stackitem.Null{}, "check")
+}

pkg/vm/context.go

@@ -54,22 +54,32 @@ type Context struct {
 	NEF *nef.File
 	// invTree is an invocation tree (or branch of it) for this context.
 	invTree *InvocationTree
+	// notificationsCount stores number of notifications emitted during current context
+	// handling.
+	notificationsCount *int
+	// isWrapped tells whether the context's DAO was wrapped into another layer of
+	// MemCachedStore on creation and whether it should be unwrapped on context unloading.
+	isWrapped bool
 }
 
 var errNoInstParam = errors.New("failed to read instruction parameter")
 
 // NewContext returns a new Context object.
 func NewContext(b []byte) *Context {
-	return NewContextWithParams(b, -1, 0)
+	return NewContextWithParams(b, -1, 0, nil)
 }
 
 // NewContextWithParams creates new Context objects using script, parameter count,
 // return value count and initial position in script.
-func NewContextWithParams(b []byte, rvcount int, pos int) *Context {
+func NewContextWithParams(b []byte, rvcount int, pos int, notificationsCount *int) *Context {
+	if notificationsCount == nil {
+		notificationsCount = new(int)
+	}
 	return &Context{
-		prog:     b,
-		retCount: rvcount,
-		nextip:   pos,
+		prog:               b,
+		retCount:           rvcount,
+		nextip:             pos,
+		notificationsCount: notificationsCount,
 	}
 }
 

pkg/vm/vm.go

@@ -67,6 +67,13 @@ type VM struct {
 	// callback to get interop price
 	getPrice func(opcode.Opcode, []byte) int64
 
+	// wraps DAO with private MemCachedStore
+	wrapDao func()
+	// commits DAO changes and unwraps DAO.
+	commitChanges func() error
+	// unwraps DAO and removes last notificationsCount notifications from the context
+	revertChanges func(notificationsCount int)
+
 	istack Stack  // invocation stack.
 	estack *Stack // execution stack.
 
@@ -116,6 +123,25 @@ func NewWithTrigger(t trigger.Type) *VM {
 	return vm
 }
 
+func (v *VM) EmitNotification() {
+	currCtx := v.Context()
+	if currCtx == nil {
+		return
+	}
+	*currCtx.notificationsCount++
+}
+
+// SetIsolationCallbacks registers given callbacks to perform DAO and interop context
+// isolation between contract calls.
+// wrapper performs DAO cloning;
+// committer persists changes made in the upper snapshot to the underlying DAO;
+// reverter rolls back the whole set of changes made in the current snapshot.
+func (v *VM) SetIsolationCallbacks(wrapper func(), committer func() error, reverter func(ntfToRemove int)) {
+	v.wrapDao = wrapper
+	v.commitChanges = committer
+	v.revertChanges = reverter
+}
+
 // SetPriceGetter registers the given PriceGetterFunc in v.
 // f accepts vm's Context, current instruction and instruction parameter.
 func (v *VM) SetPriceGetter(f func(opcode.Opcode, []byte) int64) {
@@ -320,7 +346,7 @@ func (v *VM) loadScriptWithCallingHash(b []byte, exe *nef.File, caller util.Uint
 	var sl slot
 
 	v.checkInvocationStackSize()
-	ctx := NewContextWithParams(b, rvcount, offset)
+	ctx := NewContextWithParams(b, rvcount, offset, nil)
 	if rvcount != -1 || v.estack.Len() != 0 {
 		v.estack = newStack("evaluation", &v.refs)
 	}
@@ -341,6 +367,10 @@ func (v *VM) loadScriptWithCallingHash(b []byte, exe *nef.File, caller util.Uint
 		curTree.Calls = append(curTree.Calls, newTree)
 		ctx.invTree = newTree
 	}
+	if v.wrapDao != nil {
+		v.wrapDao()
+		ctx.isWrapped = true
+	}
 	v.istack.PushItem(ctx)
 }
 
@@ -1591,6 +1621,23 @@ func (v *VM) unloadContext(ctx *Context) {
 	if ctx.static != nil && (currCtx == nil || ctx.static != currCtx.static) {
 		ctx.static.ClearRefs(&v.refs)
 	}
+	if currCtx == nil || ctx.isWrapped { // In case of CALL, CALLA, CALLL we don't need to commit/discard changes, unwrap DAO and change notificationsCount.
+		if v.uncaughtException == nil {
+			if v.commitChanges != nil {
+				if err := v.commitChanges(); err != nil {
+					// TODO: return an error instead?
+					panic(fmt.Errorf("failed to commit changes: %w", err))
+				}
+			}
+			if currCtx != nil {
+				*currCtx.notificationsCount += *ctx.notificationsCount
+			}
+		} else {
+			if v.revertChanges != nil {
+				v.revertChanges(*ctx.notificationsCount)
+			}
+		}
+	}
 }
 
 // getTryParams splits TRY(L) instruction parameter into offsets for catch and finally blocks.
@@ -1647,6 +1694,11 @@ func (v *VM) call(ctx *Context, offset int) {
 	newCtx.tryStack.elems = nil
 	initStack(&newCtx.tryStack, "exception", nil)
 	newCtx.NEF = ctx.NEF
+	// Use exactly the same counter and don't use v.wrapDao() for this context.
+	// Unloading of such unwrapped context will be properly handled inside
+	// unloadContext without unnecessary DAO unwrapping and notificationsCount changes.
+	newCtx.notificationsCount = ctx.notificationsCount
+	newCtx.isWrapped = false
 	v.istack.PushItem(newCtx)
 	newCtx.Jump(offset)
 }