From ae9658a10838b4616dd3206bb7d1b66818cac1e9 Mon Sep 17 00:00:00 2001
From: Roman Khimov <roman@nspcc.ru>
Date: Mon, 6 Jul 2020 17:17:49 +0300
Subject: [PATCH] dao: migrate nep5 balances with the contract

Fixes #1144. It's quite simple approach, we just update balance info right
upon contract migration. It will slow down migration transactions, but it
takes about 1-2 seconds to Seek through balances at mainnet's 3.8M, so the
approach should still work good enough. The other idea was to make lazy
updates (maintaining contract migration map), but it's more complicated to
implement (and implies that a balance get might also do a write).

There also is a concern about memory usage, it can give a spike of some tens
of megabytes, but that also is considered to be acceptable.
---
 pkg/core/dao/cacheddao.go | 69 ++++++++++++++++++++++++++++++++++++++-
 pkg/core/interop_neo.go   |  2 ++
 2 files changed, 70 insertions(+), 1 deletion(-)

diff --git a/pkg/core/dao/cacheddao.go b/pkg/core/dao/cacheddao.go
index 9fc792308..3862ed767 100644
--- a/pkg/core/dao/cacheddao.go
+++ b/pkg/core/dao/cacheddao.go
@@ -1,9 +1,11 @@
 package dao
 
 import (
+	"bytes"
 	"errors"
 
 	"github.com/nspcc-dev/neo-go/pkg/core/state"
+	"github.com/nspcc-dev/neo-go/pkg/core/storage"
 	"github.com/nspcc-dev/neo-go/pkg/io"
 	"github.com/nspcc-dev/neo-go/pkg/util"
 )
@@ -17,6 +19,8 @@ type Cached struct {
 	contracts map[util.Uint160]*state.Contract
 	balances  map[util.Uint160]*state.NEP5Balances
 	transfers map[util.Uint160]map[uint32]*state.NEP5TransferLog
+
+	dropNEP5Cache bool
 }
 
 // NewCached returns new Cached wrapping around given backing store.
@@ -25,7 +29,7 @@ func NewCached(d DAO) *Cached {
 	ctrs := make(map[util.Uint160]*state.Contract)
 	balances := make(map[util.Uint160]*state.NEP5Balances)
 	transfers := make(map[util.Uint160]map[uint32]*state.NEP5TransferLog)
-	return &Cached{d.GetWrapped(), accs, ctrs, balances, transfers}
+	return &Cached{d.GetWrapped(), accs, ctrs, balances, transfers, false}
 }
 
 // GetAccountStateOrNew retrieves Account from cache or underlying store
@@ -121,6 +125,65 @@ func (cd *Cached) AppendNEP5Transfer(acc util.Uint160, index uint32, tr *state.N
 	return lg.Size() >= nep5TransferBatchSize, cd.PutNEP5TransferLog(acc, index, lg)
 }
 
+// MigrateNEP5Balances migrates NEP5 balances from old contract to the new one.
+func (cd *Cached) MigrateNEP5Balances(from, to util.Uint160) error {
+	var (
+		simpleDAO *Simple
+		cachedDAO = cd
+		ok        bool
+		w         = io.NewBufBinWriter()
+	)
+	for simpleDAO == nil {
+		simpleDAO, ok = cachedDAO.DAO.(*Simple)
+		if !ok {
+			cachedDAO, ok = cachedDAO.DAO.(*Cached)
+			if !ok {
+				panic("uknown DAO")
+			}
+		}
+	}
+	for acc, bs := range cd.balances {
+		err := simpleDAO.putNEP5Balances(acc, bs, w)
+		if err != nil {
+			return err
+		}
+		w.Reset()
+	}
+	cd.dropNEP5Cache = true
+	var store = simpleDAO.Store
+	// Create another layer of cache because we can't change original storage
+	// while seeking.
+	var upStore = storage.NewMemCachedStore(store)
+	store.Seek([]byte{byte(storage.STNEP5Balances)}, func(k, v []byte) {
+		if !bytes.Contains(v, from[:]) {
+			return
+		}
+		bs := state.NewNEP5Balances()
+		reader := io.NewBinReaderFromBuf(v)
+		bs.DecodeBinary(reader)
+		if reader.Err != nil {
+			panic("bad nep5 balances")
+		}
+		tr, ok := bs.Trackers[from]
+		if !ok {
+			return
+		}
+		delete(bs.Trackers, from)
+		bs.Trackers[to] = tr
+		w.Reset()
+		bs.EncodeBinary(w.BinWriter)
+		if w.Err != nil {
+			panic("error on nep5 balance encoding")
+		}
+		err := upStore.Put(k, w.Bytes())
+		if err != nil {
+			panic("can't put value in the DB")
+		}
+	})
+	_, err := upStore.Persist()
+	return err
+}
+
 // Persist flushes all the changes made into the (supposedly) persistent
 // underlying store.
 func (cd *Cached) Persist() (int, error) {
@@ -130,6 +193,9 @@ func (cd *Cached) Persist() (int, error) {
 	// usage scenario it should be good enough if cd doesn't modify object
 	// caches (accounts/contracts/etc) in any way.
 	if ok {
+		if cd.dropNEP5Cache {
+			lowerCache.balances = make(map[util.Uint160]*state.NEP5Balances)
+		}
 		var simpleCache *Simple
 		for simpleCache == nil {
 			simpleCache, ok = lowerCache.DAO.(*Simple)
@@ -176,5 +242,6 @@ func (cd *Cached) GetWrapped() DAO {
 		cd.contracts,
 		cd.balances,
 		cd.transfers,
+		false,
 	}
 }
diff --git a/pkg/core/interop_neo.go b/pkg/core/interop_neo.go
index cf41306c4..43e46dd6e 100644
--- a/pkg/core/interop_neo.go
+++ b/pkg/core/interop_neo.go
@@ -154,6 +154,7 @@ func contractUpdate(ic *interop.Context, v *vm.VM) error {
 		if err := ic.DAO.DeleteContractState(oldHash); err != nil {
 			return fmt.Errorf("failed to update script: %v", err)
 		}
+		ic.DAO.MigrateNEP5Balances(oldHash, newHash)
 	}
 	// if manifest was provided, update the old contract manifest and check associated
 	// storage items if needed
@@ -181,6 +182,7 @@ func contractUpdate(ic *interop.Context, v *vm.VM) error {
 			return fmt.Errorf("failed to update manifest: %v", err)
 		}
 	}
+
 	return nil
 }