Merge pull request #1188 from nspcc-dev/forward-port-from-2.x

Forward port from 2.x

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,
 	}
 }

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
 }
 

pkg/network/server.go

@@ -187,7 +187,7 @@ func (s *Server) Shutdown() {
 	s.log.Info("shutting down server", zap.Int("peers", s.PeerCount()))
 	s.transport.Close()
 	s.discovery.Close()
-	for p := range s.peers {
+	for p := range s.Peers() {
 		p.Disconnect(errServerShutdown)
 	}
 	s.bQueue.discard()

pkg/rpc/request/param.go

@@ -163,6 +163,16 @@ func (p *Param) GetUint160FromAddress() (util.Uint160, error) {
 	return address.StringToUint160(s)
 }
 
+// GetUint160FromAddressOrHex returns Uint160 value of the parameter that was
+// supplied either as raw hex or as an address.
+func (p *Param) GetUint160FromAddressOrHex() (util.Uint160, error) {
+	u, err := p.GetUint160FromHex()
+	if err == nil {
+		return u, err
+	}
+	return p.GetUint160FromAddress()
+}
+
 // GetFuncParam returns current parameter as a function call parameter.
 func (p *Param) GetFuncParam() (FuncParam, error) {
 	if p == nil {

pkg/rpc/request/param_test.go

@@ -205,6 +205,25 @@ func TestParamGetUint160FromAddress(t *testing.T) {
 	require.NotNil(t, err)
 }
 
+func TestParam_GetUint160FromAddressOrHex(t *testing.T) {
+	in := "NPAsqZkx9WhNd4P72uhZxBhLinSuNkxfB8"
+	inHex, _ := address.StringToUint160(in)
+
+	t.Run("Address", func(t *testing.T) {
+		p := Param{StringT, in}
+		u, err := p.GetUint160FromAddressOrHex()
+		require.NoError(t, err)
+		require.Equal(t, inHex, u)
+	})
+
+	t.Run("Hex", func(t *testing.T) {
+		p := Param{StringT, inHex.StringLE()}
+		u, err := p.GetUint160FromAddressOrHex()
+		require.NoError(t, err)
+		require.Equal(t, inHex, u)
+	})
+}
+
 func TestParamGetFuncParam(t *testing.T) {
 	fp := FuncParam{
 		Type: smartcontract.StringType,

pkg/rpc/server/server.go

@@ -496,7 +496,7 @@ func (s *Server) getApplicationLog(reqParams request.Params) (interface{}, *resp
 }
 
 func (s *Server) getNEP5Balances(ps request.Params) (interface{}, *response.Error) {
-	u, err := ps.ValueWithType(0, request.StringT).GetUint160FromHex()
+	u, err := ps.Value(0).GetUint160FromAddressOrHex()
 	if err != nil {
 		return nil, response.ErrInvalidParams
 	}
@@ -525,7 +525,7 @@ func (s *Server) getNEP5Balances(ps request.Params) (interface{}, *response.Erro
 }
 
 func (s *Server) getNEP5Transfers(ps request.Params) (interface{}, *response.Error) {
-	u, err := ps.ValueWithType(0, request.StringT).GetUint160FromAddress()
+	u, err := ps.Value(0).GetUint160FromAddressOrHex()
 	if err != nil {
 		return nil, response.ErrInvalidParams
 	}

pkg/rpc/server/server_test.go

@@ -129,33 +129,13 @@ var rpcTestCases = map[string][]rpcTestCase{
 			name:   "positive",
 			params: `["` + testchain.PrivateKeyByID(0).GetScriptHash().StringLE() + `"]`,
 			result: func(e *executor) interface{} { return &result.NEP5Balances{} },
-			check: func(t *testing.T, e *executor, acc interface{}) {
-				res, ok := acc.(*result.NEP5Balances)
-				require.True(t, ok)
-				rubles, err := util.Uint160DecodeStringLE(testContractHash)
-				require.NoError(t, err)
-				expected := result.NEP5Balances{
-					Balances: []result.NEP5Balance{
-						{
-							Asset:       rubles,
-							Amount:      "8.77",
-							LastUpdated: 6,
+			check:  checkNep5Balances,
 		},
 		{
-							Asset:       e.chain.GoverningTokenHash(),
-							Amount:      "99998000",
-							LastUpdated: 4,
-						},
-						{
-							Asset:       e.chain.UtilityTokenHash(),
-							Amount:      "915.79002700",
-							LastUpdated: 6,
-						}},
-					Address: testchain.PrivateKeyByID(0).GetScriptHash().StringLE(),
-				}
-				require.Equal(t, testchain.PrivateKeyByID(0).Address(), res.Address)
-				require.ElementsMatch(t, expected.Balances, res.Balances)
-			},
+			name:   "positive_address",
+			params: `["` + address.Uint160ToString(testchain.PrivateKeyByID(0).GetScriptHash()) + `"]`,
+			result: func(e *executor) interface{} { return &result.NEP5Balances{} },
+			check:  checkNep5Balances,
 		},
 	},
 	"getnep5transfers": {
@@ -173,127 +153,13 @@ var rpcTestCases = map[string][]rpcTestCase{
 			name:   "positive",
 			params: `["` + testchain.PrivateKeyByID(0).Address() + `"]`,
 			result: func(e *executor) interface{} { return &result.NEP5Transfers{} },
-			check: func(t *testing.T, e *executor, acc interface{}) {
-				res, ok := acc.(*result.NEP5Transfers)
-				require.True(t, ok)
-				rublesHash, err := util.Uint160DecodeStringLE(testContractHash)
-				require.NoError(t, err)
-				blockSendRubles, err := e.chain.GetBlock(e.chain.GetHeaderHash(6))
-				require.NoError(t, err)
-				require.Equal(t, 1, len(blockSendRubles.Transactions))
-				txSendRublesHash := blockSendRubles.Transactions[0].Hash()
-				blockReceiveRubles, err := e.chain.GetBlock(e.chain.GetHeaderHash(5))
-				require.NoError(t, err)
-				require.Equal(t, 2, len(blockReceiveRubles.Transactions))
-				txReceiveRublesHash := blockReceiveRubles.Transactions[1].Hash()
-				blockReceiveGAS, err := e.chain.GetBlock(e.chain.GetHeaderHash(1))
-				require.NoError(t, err)
-				require.Equal(t, 2, len(blockReceiveGAS.Transactions))
-				txReceiveNEOHash := blockReceiveGAS.Transactions[0].Hash()
-				txReceiveGASHash := blockReceiveGAS.Transactions[1].Hash()
-				blockSendNEO, err := e.chain.GetBlock(e.chain.GetHeaderHash(4))
-				require.NoError(t, err)
-				require.Equal(t, 1, len(blockSendNEO.Transactions))
-				txSendNEOHash := blockSendNEO.Transactions[0].Hash()
-				expected := result.NEP5Transfers{
-					Sent: []result.NEP5Transfer{
-						{
-							Timestamp:   blockSendRubles.Timestamp,
-							Asset:       rublesHash,
-							Address:     testchain.PrivateKeyByID(1).Address(),
-							Amount:      "1.23",
-							Index:       6,
-							NotifyIndex: 0,
-							TxHash:      txSendRublesHash,
+			check:  checkNep5Transfers,
 		},
 		{
-							Timestamp:   blockSendNEO.Timestamp,
-							Asset:       e.chain.GoverningTokenHash(),
-							Address:     testchain.PrivateKeyByID(1).Address(),
-							Amount:      "1000",
-							Index:       4,
-							NotifyIndex: 0,
-							TxHash:      txSendNEOHash,
-						},
-					},
-					Received: []result.NEP5Transfer{
-						{
-							Timestamp:   blockReceiveRubles.Timestamp,
-							Asset:       rublesHash,
-							Address:     address.Uint160ToString(rublesHash),
-							Amount:      "10",
-							Index:       5,
-							NotifyIndex: 0,
-							TxHash:      txReceiveRublesHash,
-						},
-						{
-							Timestamp:   blockSendNEO.Timestamp,
-							Asset:       e.chain.UtilityTokenHash(),
-							Address:     "", // Minted GAS.
-							Amount:      "17.99982000",
-							Index:       4,
-							NotifyIndex: 0,
-							TxHash:      txSendNEOHash,
-						},
-						{
-							Timestamp:   blockReceiveGAS.Timestamp,
-							Asset:       e.chain.UtilityTokenHash(),
-							Address:     testchain.MultisigAddress(),
-							Amount:      "1000",
-							Index:       1,
-							NotifyIndex: 0,
-							TxHash:      txReceiveGASHash,
-						},
-						{
-							Timestamp:   blockReceiveGAS.Timestamp,
-							Asset:       e.chain.GoverningTokenHash(),
-							Address:     testchain.MultisigAddress(),
-							Amount:      "99999000",
-							Index:       1,
-							NotifyIndex: 0,
-							TxHash:      txReceiveNEOHash,
-						},
-					},
-					Address: testchain.PrivateKeyByID(0).Address(),
-				}
-
-				// take burned gas into account
-				u := testchain.PrivateKeyByID(0).GetScriptHash()
-				for i := 0; i <= int(e.chain.BlockHeight()); i++ {
-					var netFee int64
-					h := e.chain.GetHeaderHash(i)
-					b, err := e.chain.GetBlock(h)
-					require.NoError(t, err)
-					for j := range b.Transactions {
-						if u.Equals(b.Transactions[j].Sender) {
-							amount := b.Transactions[j].SystemFee + b.Transactions[j].NetworkFee
-							expected.Sent = append(expected.Sent, result.NEP5Transfer{
-								Timestamp: b.Timestamp,
-								Asset:     e.chain.UtilityTokenHash(),
-								Address:   "", // burn has empty receiver
-								Amount:    amountToString(big.NewInt(amount), 8),
-								Index:     b.Index,
-								TxHash:    b.Hash(),
-							})
-						}
-						netFee += b.Transactions[j].NetworkFee
-					}
-					if i > 0 {
-						expected.Received = append(expected.Received, result.NEP5Transfer{
-							Timestamp: b.Timestamp,
-							Asset:     e.chain.UtilityTokenHash(),
-							Address:   "", // minted from network fees.
-							Amount:    amountToString(big.NewInt(netFee), 8),
-							Index:     b.Index,
-							TxHash:    b.Hash(),
-						})
-					}
-
-				}
-				require.Equal(t, expected.Address, res.Address)
-				require.ElementsMatch(t, expected.Sent, res.Sent)
-				require.ElementsMatch(t, expected.Received, res.Received)
-			},
+			name:   "positive_hash",
+			params: `["` + testchain.PrivateKeyByID(0).GetScriptHash().StringLE() + `"]`,
+			result: func(e *executor) interface{} { return &result.NEP5Transfers{} },
+			check:  checkNep5Transfers,
 		},
 	},
 	"getstorage": {
@@ -1110,3 +976,153 @@ func doRPCCallOverHTTP(rpcCall string, url string, t *testing.T) []byte {
 	assert.NoErrorf(t, err, "could not read response from the request: %s", rpcCall)
 	return bytes.TrimSpace(body)
 }
+
+func checkNep5Balances(t *testing.T, e *executor, acc interface{}) {
+	res, ok := acc.(*result.NEP5Balances)
+	require.True(t, ok)
+	rubles, err := util.Uint160DecodeStringLE(testContractHash)
+	require.NoError(t, err)
+	expected := result.NEP5Balances{
+		Balances: []result.NEP5Balance{
+			{
+				Asset:       rubles,
+				Amount:      "8.77",
+				LastUpdated: 6,
+			},
+			{
+				Asset:       e.chain.GoverningTokenHash(),
+				Amount:      "99998000",
+				LastUpdated: 4,
+			},
+			{
+				Asset:       e.chain.UtilityTokenHash(),
+				Amount:      "915.79002700",
+				LastUpdated: 6,
+			}},
+		Address: testchain.PrivateKeyByID(0).GetScriptHash().StringLE(),
+	}
+	require.Equal(t, testchain.PrivateKeyByID(0).Address(), res.Address)
+	require.ElementsMatch(t, expected.Balances, res.Balances)
+}
+
+func checkNep5Transfers(t *testing.T, e *executor, acc interface{}) {
+	res, ok := acc.(*result.NEP5Transfers)
+	require.True(t, ok)
+	rublesHash, err := util.Uint160DecodeStringLE(testContractHash)
+	require.NoError(t, err)
+	blockSendRubles, err := e.chain.GetBlock(e.chain.GetHeaderHash(6))
+	require.NoError(t, err)
+	require.Equal(t, 1, len(blockSendRubles.Transactions))
+	txSendRublesHash := blockSendRubles.Transactions[0].Hash()
+	blockReceiveRubles, err := e.chain.GetBlock(e.chain.GetHeaderHash(5))
+	require.NoError(t, err)
+	require.Equal(t, 2, len(blockReceiveRubles.Transactions))
+	txReceiveRublesHash := blockReceiveRubles.Transactions[1].Hash()
+	blockReceiveGAS, err := e.chain.GetBlock(e.chain.GetHeaderHash(1))
+	require.NoError(t, err)
+	require.Equal(t, 2, len(blockReceiveGAS.Transactions))
+	txReceiveNEOHash := blockReceiveGAS.Transactions[0].Hash()
+	txReceiveGASHash := blockReceiveGAS.Transactions[1].Hash()
+	blockSendNEO, err := e.chain.GetBlock(e.chain.GetHeaderHash(4))
+	require.NoError(t, err)
+	require.Equal(t, 1, len(blockSendNEO.Transactions))
+	txSendNEOHash := blockSendNEO.Transactions[0].Hash()
+	expected := result.NEP5Transfers{
+		Sent: []result.NEP5Transfer{
+			{
+				Timestamp:   blockSendRubles.Timestamp,
+				Asset:       rublesHash,
+				Address:     testchain.PrivateKeyByID(1).Address(),
+				Amount:      "1.23",
+				Index:       6,
+				NotifyIndex: 0,
+				TxHash:      txSendRublesHash,
+			},
+			{
+				Timestamp:   blockSendNEO.Timestamp,
+				Asset:       e.chain.GoverningTokenHash(),
+				Address:     testchain.PrivateKeyByID(1).Address(),
+				Amount:      "1000",
+				Index:       4,
+				NotifyIndex: 0,
+				TxHash:      txSendNEOHash,
+			},
+		},
+		Received: []result.NEP5Transfer{
+			{
+				Timestamp:   blockReceiveRubles.Timestamp,
+				Asset:       rublesHash,
+				Address:     address.Uint160ToString(rublesHash),
+				Amount:      "10",
+				Index:       5,
+				NotifyIndex: 0,
+				TxHash:      txReceiveRublesHash,
+			},
+			{
+				Timestamp:   blockSendNEO.Timestamp,
+				Asset:       e.chain.UtilityTokenHash(),
+				Address:     "", // Minted GAS.
+				Amount:      "17.99982000",
+				Index:       4,
+				NotifyIndex: 0,
+				TxHash:      txSendNEOHash,
+			},
+			{
+				Timestamp:   blockReceiveGAS.Timestamp,
+				Asset:       e.chain.UtilityTokenHash(),
+				Address:     testchain.MultisigAddress(),
+				Amount:      "1000",
+				Index:       1,
+				NotifyIndex: 0,
+				TxHash:      txReceiveGASHash,
+			},
+			{
+				Timestamp:   blockReceiveGAS.Timestamp,
+				Asset:       e.chain.GoverningTokenHash(),
+				Address:     testchain.MultisigAddress(),
+				Amount:      "99999000",
+				Index:       1,
+				NotifyIndex: 0,
+				TxHash:      txReceiveNEOHash,
+			},
+		},
+		Address: testchain.PrivateKeyByID(0).Address(),
+	}
+
+	// take burned gas into account
+	u := testchain.PrivateKeyByID(0).GetScriptHash()
+	for i := 0; i <= int(e.chain.BlockHeight()); i++ {
+		var netFee int64
+		h := e.chain.GetHeaderHash(i)
+		b, err := e.chain.GetBlock(h)
+		require.NoError(t, err)
+		for j := range b.Transactions {
+			if u.Equals(b.Transactions[j].Sender) {
+				amount := b.Transactions[j].SystemFee + b.Transactions[j].NetworkFee
+				expected.Sent = append(expected.Sent, result.NEP5Transfer{
+					Timestamp: b.Timestamp,
+					Asset:     e.chain.UtilityTokenHash(),
+					Address:   "", // burn has empty receiver
+					Amount:    amountToString(big.NewInt(amount), 8),
+					Index:     b.Index,
+					TxHash:    b.Hash(),
+				})
+			}
+			netFee += b.Transactions[j].NetworkFee
+		}
+		if i > 0 {
+			expected.Received = append(expected.Received, result.NEP5Transfer{
+				Timestamp: b.Timestamp,
+				Asset:     e.chain.UtilityTokenHash(),
+				Address:   "", // minted from network fees.
+				Amount:    amountToString(big.NewInt(netFee), 8),
+				Index:     b.Index,
+				TxHash:    b.Hash(),
+			})
+		}
+
+	}
+	require.Equal(t, expected.Address, res.Address)
+	require.ElementsMatch(t, expected.Sent, res.Sent)
+	require.ElementsMatch(t, expected.Received, res.Received)
+}