core: refactor (*NEO).GetCandidates and add (*NEO).GetAllCandidates

pkg/core/native/native_neo.go

@@ -1,6 +1,7 @@
 package native
 
 import (
+	"context"
 	"crypto/elliptic"
 	"encoding/binary"
 	"errors"
@@ -13,6 +14,7 @@ import (
 	"github.com/nspcc-dev/neo-go/pkg/core/dao"
 	"github.com/nspcc-dev/neo-go/pkg/core/interop"
 	"github.com/nspcc-dev/neo-go/pkg/core/interop/runtime"
+	istorage "github.com/nspcc-dev/neo-go/pkg/core/interop/storage"
 	"github.com/nspcc-dev/neo-go/pkg/core/native/nativenames"
 	"github.com/nspcc-dev/neo-go/pkg/core/state"
 	"github.com/nspcc-dev/neo-go/pkg/core/storage"
@@ -91,6 +93,10 @@ const (
 	committeeRewardRatio = 10
 	// neoHolderRewardRatio is a percent of generated GAS that is distributed to voters.
 	voterRewardRatio = 80
+
+	// maxGetCandidatesRespLen is the maximum number of candidates to return from the
+	// getCandidates method.
+	maxGetCandidatesRespLen = 256
 )
 
 var (
@@ -194,6 +200,15 @@ func newNEO(cfg config.ProtocolConfiguration) *NEO {
 	md = newMethodAndPrice(n.getCandidatesCall, 1<<22, callflag.ReadStates)
 	n.AddMethod(md, desc)
 
+	desc = newDescriptor("getAllCandidates", smartcontract.InteropInterfaceType)
+	md = newMethodAndPrice(n.getAllCandidatesCall, 1<<22, callflag.ReadStates)
+	n.AddMethod(md, desc)
+
+	desc = newDescriptor("getCandidateVote", smartcontract.IntegerType,
+		manifest.NewParameter("pubkey", smartcontract.PublicKeyType))
+	md = newMethodAndPrice(n.getCandidateVoteCall, 1<<15, callflag.ReadStates)
+	n.AddMethod(md, desc)
+
 	desc = newDescriptor("getAccountState", smartcontract.ArrayType,
 		manifest.NewParameter("account", smartcontract.Hash160Type))
 	md = newMethodAndPrice(n.getAccountState, 1<<15, callflag.ReadStates)
@@ -851,7 +866,7 @@ func (n *NEO) ModifyAccountVotes(acc *state.NEOBalance, d *dao.Simple, value *bi
 	return nil
 }
 
-func (n *NEO) getCandidates(d *dao.Simple, sortByKey bool) ([]keyWithVotes, error) {
+func (n *NEO) getCandidates(d *dao.Simple, sortByKey bool, max int) ([]keyWithVotes, error) {
 	arr := make([]keyWithVotes, 0)
 	buf := io.NewBufBinWriter()
 	d.Seek(n.ID, storage.SeekRange{Prefix: []byte{prefixCandidate}}, func(k, v []byte) bool {
@@ -861,7 +876,7 @@ func (n *NEO) getCandidates(d *dao.Simple, sortByKey bool) ([]keyWithVotes, erro
 			arr = append(arr, keyWithVotes{Key: string(k), Votes: &c.Votes})
 		}
 		buf.Reset()
-		return true
+		return !sortByKey || max > 0 && len(arr) < max
 	})
 
 	if !sortByKey {
@@ -893,7 +908,7 @@ func (n *NEO) getCandidates(d *dao.Simple, sortByKey bool) ([]keyWithVotes, erro
 // GetCandidates returns current registered validators list with keys
 // and votes.
 func (n *NEO) GetCandidates(d *dao.Simple) ([]state.Validator, error) {
-	kvs, err := n.getCandidates(d, true)
+	kvs, err := n.getCandidates(d, true, maxGetCandidatesRespLen)
 	if err != nil {
 		return nil, err
 	}
@@ -909,7 +924,7 @@ func (n *NEO) GetCandidates(d *dao.Simple) ([]state.Validator, error) {
 }
 
 func (n *NEO) getCandidatesCall(ic *interop.Context, _ []stackitem.Item) stackitem.Item {
-	validators, err := n.getCandidates(ic.DAO, true)
+	validators, err := n.getCandidates(ic.DAO, true, maxGetCandidatesRespLen)
 	if err != nil {
 		panic(err)
 	}
@@ -923,6 +938,51 @@ func (n *NEO) getCandidatesCall(ic *interop.Context, _ []stackitem.Item) stackit
 	return stackitem.NewArray(arr)
 }
 
+func (n *NEO) getAllCandidatesCall(ic *interop.Context, _ []stackitem.Item) stackitem.Item {
+	ctx, cancel := context.WithCancel(context.Background())
+	prefix := []byte{prefixCandidate}
+	buf := io.NewBufBinWriter()
+	keep := func(kv storage.KeyValue) bool {
+		c := new(candidate).FromBytes(kv.Value)
+		emit.CheckSig(buf.BinWriter, kv.Key)
+		if c.Registered && !n.Policy.IsBlocked(ic.DAO, hash.Hash160(buf.Bytes())) {
+			buf.Reset()
+			return true
+		}
+		buf.Reset()
+		return false
+	}
+	seekres := ic.DAO.SeekAsync(ctx, n.ID, storage.SeekRange{Prefix: prefix})
+	filteredRes := make(chan storage.KeyValue)
+	go func() {
+		for kv := range seekres {
+			if keep(kv) {
+				filteredRes <- kv
+			}
+		}
+		close(filteredRes)
+	}()
+
+	opts := istorage.FindRemovePrefix | istorage.FindDeserialize | istorage.FindPick1
+	item := istorage.NewIterator(filteredRes, prefix, int64(opts))
+	ic.RegisterCancelFunc(cancel)
+	return stackitem.NewInterop(item)
+}
+
+func (n *NEO) getCandidateVoteCall(ic *interop.Context, args []stackitem.Item) stackitem.Item {
+	pub := toPublicKey(args[0])
+	key := makeValidatorKey(pub)
+	si := ic.DAO.GetStorageItem(n.ID, key)
+	if si == nil {
+		return stackitem.NewBigInteger(big.NewInt(-1))
+	}
+	c := new(candidate).FromBytes(si)
+	if !c.Registered {
+		return stackitem.NewBigInteger(big.NewInt(-1))
+	}
+	return stackitem.NewBigInteger(&c.Votes)
+}
+
 func (n *NEO) getAccountState(ic *interop.Context, args []stackitem.Item) stackitem.Item {
 	key := makeAccountKey(toUint160(args[0]))
 	si := ic.DAO.GetStorageItem(n.ID, key)
@@ -1021,7 +1081,7 @@ func (n *NEO) computeCommitteeMembers(blockHeight uint32, d *dao.Simple) (keys.P
 	count := n.cfg.GetCommitteeSize(blockHeight + 1)
 	// Can be sorted and/or returned to outside users, thus needs to be copied.
 	sbVals := keys.PublicKeys(n.standbyKeys[:count]).Copy()
-	cs, err := n.getCandidates(d, false)
+	cs, err := n.getCandidates(d, false, -1)
 	if err != nil {
 		return nil, nil, err
 	}

pkg/core/native/native_test/neo_test.go

@@ -1,6 +1,7 @@
 package native_test
 
 import (
+	"bytes"
 	"encoding/json"
 	"math"
 	"math/big"
@@ -9,15 +10,20 @@ import (
 
 	"github.com/nspcc-dev/neo-go/internal/contracts"
 	"github.com/nspcc-dev/neo-go/internal/random"
+	"github.com/nspcc-dev/neo-go/pkg/core/interop/interopnames"
 	"github.com/nspcc-dev/neo-go/pkg/core/native"
 	"github.com/nspcc-dev/neo-go/pkg/core/native/nativenames"
 	"github.com/nspcc-dev/neo-go/pkg/core/state"
 	"github.com/nspcc-dev/neo-go/pkg/core/transaction"
 	"github.com/nspcc-dev/neo-go/pkg/crypto/keys"
+	"github.com/nspcc-dev/neo-go/pkg/io"
 	"github.com/nspcc-dev/neo-go/pkg/neotest"
 	"github.com/nspcc-dev/neo-go/pkg/neotest/chain"
+	"github.com/nspcc-dev/neo-go/pkg/smartcontract/callflag"
 	"github.com/nspcc-dev/neo-go/pkg/smartcontract/trigger"
 	"github.com/nspcc-dev/neo-go/pkg/util"
+	"github.com/nspcc-dev/neo-go/pkg/vm/emit"
+	"github.com/nspcc-dev/neo-go/pkg/vm/opcode"
 	"github.com/nspcc-dev/neo-go/pkg/vm/stackitem"
 	"github.com/stretchr/testify/require"
 )
@@ -482,3 +488,88 @@ func TestNEO_CalculateBonus(t *testing.T) {
 			claimTx.SystemFee-claimTx.NetworkFee + +firstPart + secondPart))
 	})
 }
+
+func TestNEO_GetCandidates(t *testing.T) {
+	neoCommitteeInvoker := newNeoCommitteeClient(t, 100_0000_0000)
+	neoValidatorsInvoker := neoCommitteeInvoker.WithSigners(neoCommitteeInvoker.Validator)
+	policyInvoker := neoCommitteeInvoker.CommitteeInvoker(neoCommitteeInvoker.NativeHash(t, nativenames.Policy))
+	e := neoCommitteeInvoker.Executor
+
+	cfg := e.Chain.GetConfig()
+	candidatesCount := cfg.GetCommitteeSize(0) - 1
+
+	// Register a set of candidates and vote for them.
+	voters := make([]neotest.Signer, candidatesCount)
+	candidates := make([]neotest.Signer, candidatesCount)
+	for i := 0; i < candidatesCount; i++ {
+		voters[i] = e.NewAccount(t, 10_0000_0000)
+		candidates[i] = e.NewAccount(t, 2000_0000_0000) // enough for one registration
+	}
+	txes := make([]*transaction.Transaction, 0, candidatesCount*3)
+	for i := 0; i < candidatesCount; i++ {
+		transferTx := neoValidatorsInvoker.PrepareInvoke(t, "transfer", e.Validator.ScriptHash(), voters[i].(neotest.SingleSigner).Account().PrivateKey().GetScriptHash(), int64(candidatesCount+1-i)*1000000, nil)
+		txes = append(txes, transferTx)
+		registerTx := neoValidatorsInvoker.WithSigners(candidates[i]).PrepareInvoke(t, "registerCandidate", candidates[i].(neotest.SingleSigner).Account().PrivateKey().PublicKey().Bytes())
+		txes = append(txes, registerTx)
+		voteTx := neoValidatorsInvoker.WithSigners(voters[i]).PrepareInvoke(t, "vote", voters[i].(neotest.SingleSigner).Account().PrivateKey().GetScriptHash(), candidates[i].(neotest.SingleSigner).Account().PrivateKey().PublicKey().Bytes())
+		txes = append(txes, voteTx)
+	}
+
+	neoValidatorsInvoker.AddNewBlock(t, txes...)
+	for _, tx := range txes {
+		e.CheckHalt(t, tx.Hash(), stackitem.Make(true)) // luckily, both `transfer`, `registerCandidate` and `vote` return boolean values
+	}
+	expected := make([]stackitem.Item, candidatesCount)
+	for i := range expected {
+		pub := candidates[i].(neotest.SingleSigner).Account().PrivateKey().PublicKey().Bytes()
+		v := stackitem.NewBigInteger(big.NewInt(int64(candidatesCount-i+1) * 1000000))
+		expected[i] = stackitem.NewStruct([]stackitem.Item{
+			stackitem.NewByteArray(pub),
+			v,
+		})
+		neoCommitteeInvoker.Invoke(t, v, "getCandidateVote", pub)
+	}
+	sort.Slice(expected, func(i, j int) bool {
+		return bytes.Compare(expected[i].Value().([]stackitem.Item)[0].Value().([]byte), expected[j].Value().([]stackitem.Item)[0].Value().([]byte)) < 0
+	})
+	neoCommitteeInvoker.Invoke(t, stackitem.NewArray(expected), "getCandidates")
+
+	// Check that GetAllCandidates works the same way as GetCandidates.
+	checkGetAllCandidates := func(t *testing.T, expected []stackitem.Item) {
+		for i := 0; i < len(expected)+1; i++ {
+			w := io.NewBufBinWriter()
+			emit.AppCall(w.BinWriter, neoCommitteeInvoker.Hash, "getAllCandidates", callflag.All)
+			for j := 0; j < i+1; j++ {
+				emit.Opcodes(w.BinWriter, opcode.DUP)
+				emit.Syscall(w.BinWriter, interopnames.SystemIteratorNext)
+				emit.Opcodes(w.BinWriter, opcode.DROP) // drop the value returned from Next.
+			}
+			emit.Syscall(w.BinWriter, interopnames.SystemIteratorValue)
+			require.NoError(t, w.Err)
+			h := neoCommitteeInvoker.InvokeScript(t, w.Bytes(), neoCommitteeInvoker.Signers)
+			if i < len(expected) {
+				e.CheckHalt(t, h, expected[i])
+			} else {
+				e.CheckFault(t, h, "iterator index out of range") // ensure there are no extra elements.
+			}
+			w.Reset()
+		}
+	}
+	checkGetAllCandidates(t, expected)
+
+	// Block candidate and check it won't be returned from getCandidates and getAllCandidates.
+	unlucky := candidates[len(candidates)-1].(neotest.SingleSigner).Account().PrivateKey().PublicKey()
+	policyInvoker.Invoke(t, true, "blockAccount", unlucky.GetScriptHash())
+	for i := range expected {
+		if bytes.Equal(expected[i].Value().([]stackitem.Item)[0].Value().([]byte), unlucky.Bytes()) {
+			if i != len(expected)-1 {
+				expected = append(expected[:i], expected[i+1:]...)
+			} else {
+				expected = expected[:i]
+			}
+			break
+		}
+	}
+	neoCommitteeInvoker.Invoke(t, expected, "getCandidates")
+	checkGetAllCandidates(t, expected)
+}