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 }