diff --git a/pkg/core/native/validators_count.go b/pkg/core/native/validators_count.go index be551bfdb..d3b0feccf 100644 --- a/pkg/core/native/validators_count.go +++ b/pkg/core/native/validators_count.go @@ -1,6 +1,7 @@ package native import ( + "errors" "math/big" "github.com/nspcc-dev/neo-go/pkg/encoding/bigint" @@ -41,6 +42,7 @@ func (vc *ValidatorsCount) Bytes() []byte { // EncodeBinary implements io.Serializable interface. func (vc *ValidatorsCount) EncodeBinary(w *io.BinWriter) { + w.WriteVarUint(uint64(MaxValidatorsVoted)) for i := range vc { w.WriteVarBytes(bigint.ToBytes(&vc[i])) } @@ -48,7 +50,12 @@ func (vc *ValidatorsCount) EncodeBinary(w *io.BinWriter) { // DecodeBinary implements io.Serializable interface. func (vc *ValidatorsCount) DecodeBinary(r *io.BinReader) { - for i := range vc { + count := r.ReadVarUint() + if count < 0 || count > MaxValidatorsVoted { + r.Err = errors.New("invalid validators count") + return + } + for i := 0; i < int(count); i++ { buf := r.ReadVarBytes() if r.Err != nil { return diff --git a/pkg/core/state/native_state.go b/pkg/core/state/native_state.go index 9c8b3d915..d551c6711 100644 --- a/pkg/core/state/native_state.go +++ b/pkg/core/state/native_state.go @@ -4,8 +4,8 @@ import ( "math/big" "github.com/nspcc-dev/neo-go/pkg/crypto/keys" - "github.com/nspcc-dev/neo-go/pkg/encoding/bigint" "github.com/nspcc-dev/neo-go/pkg/io" + "github.com/nspcc-dev/neo-go/pkg/vm/stackitem" ) // NEP5BalanceState represents balance state of a NEP5-token. @@ -44,18 +44,27 @@ func (s *NEP5BalanceState) Bytes() []byte { return w.Bytes() } +func (s *NEP5BalanceState) toStackItem() stackitem.Item { + return stackitem.NewStruct([]stackitem.Item{stackitem.NewBigInteger(&s.Balance)}) +} + +func (s *NEP5BalanceState) fromStackItem(item stackitem.Item) { + s.Balance = *item.(*stackitem.Struct).Value().([]stackitem.Item)[0].Value().(*big.Int) +} + // EncodeBinary implements io.Serializable interface. func (s *NEP5BalanceState) EncodeBinary(w *io.BinWriter) { - w.WriteVarBytes(bigint.ToBytes(&s.Balance)) + si := s.toStackItem() + stackitem.EncodeBinaryStackItem(si, w) } // DecodeBinary implements io.Serializable interface. func (s *NEP5BalanceState) DecodeBinary(r *io.BinReader) { - buf := r.ReadVarBytes() + si := stackitem.DecodeBinaryStackItem(r) if r.Err != nil { return } - s.Balance = *bigint.FromBytes(buf) + s.fromStackItem(si) } // NEOBalanceStateFromBytes converts serialized NEOBalanceState to structure. @@ -85,14 +94,37 @@ func (s *NEOBalanceState) Bytes() []byte { // EncodeBinary implements io.Serializable interface. func (s *NEOBalanceState) EncodeBinary(w *io.BinWriter) { - s.NEP5BalanceState.EncodeBinary(w) - w.WriteU32LE(s.BalanceHeight) - w.WriteArray(s.Votes) + si := s.toStackItem() + stackitem.EncodeBinaryStackItem(si, w) } // DecodeBinary implements io.Serializable interface. func (s *NEOBalanceState) DecodeBinary(r *io.BinReader) { - s.NEP5BalanceState.DecodeBinary(r) - s.BalanceHeight = r.ReadU32LE() - r.ReadArray(&s.Votes) + si := stackitem.DecodeBinaryStackItem(r) + if r.Err != nil { + return + } + s.fromStackItem(si) +} + +func (s *NEOBalanceState) toStackItem() stackitem.Item { + result := s.NEP5BalanceState.toStackItem().(*stackitem.Struct) + result.Append(stackitem.NewBigInteger(big.NewInt(int64(s.BalanceHeight)))) + votes := make([]stackitem.Item, len(s.Votes)) + for i, v := range s.Votes { + votes[i] = stackitem.NewByteArray(v.Bytes()) + } + result.Append(stackitem.NewArray(votes)) + return result +} + +func (s *NEOBalanceState) fromStackItem(item stackitem.Item) { + structItem := item.Value().([]stackitem.Item) + s.Balance = *structItem[0].Value().(*big.Int) + s.BalanceHeight = uint32(structItem[1].Value().(*big.Int).Int64()) + votes := structItem[2].Value().([]stackitem.Item) + s.Votes = make([]*keys.PublicKey, len(votes)) + for i, v := range votes { + s.Votes[i].DecodeBytes(v.Value().([]byte)) + } }