Merge pull request #1149 from nspcc-dev/fix-nep5-balances-2.x

Fix nep5 balances migration (2.x)

pkg/core/dao/cacheddao.go

@@ -6,6 +6,7 @@ import (
 	"sort"
 
 	"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"
 )
@@ -21,6 +22,8 @@ type Cached struct {
 	balances  map[util.Uint160]*state.NEP5Balances
 	transfers map[util.Uint160]map[uint32]*state.NEP5TransferLog
 	storage   *itemCache
+
+	dropNEP5Cache bool
 }
 
 // NewCached returns new Cached wrapping around given backing store.
@@ -39,7 +42,7 @@ func NewCached(d DAO) *Cached {
 			}
 		}
 	}
-	return &Cached{dao, accs, ctrs, unspents, balances, transfers, st}
+	return &Cached{dao, accs, ctrs, unspents, balances, transfers, st, false}
 }
 
 // GetAccountStateOrNew retrieves Account from cache or underlying store
@@ -149,6 +152,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) {
@@ -162,6 +224,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 {
 			if err := lowerCache.FlushStorage(); err != nil {
@@ -220,6 +285,7 @@ func (cd *Cached) GetWrapped() DAO {
 		cd.balances,
 		cd.transfers,
 		cd.storage,
+		false,
 	}
 }
 

pkg/core/interop_neo.go

@@ -597,6 +597,7 @@ func (ic *interopContext) contractMigrate(v *vm.VM) error {
 					return err
 				}
 			}
+			ic.dao.MigrateNEP5Balances(hash, contract.ScriptHash())
 		}
 	}
 	v.Estack().PushVal(vm.NewInteropItem(contract))