rpcclient/management: add new methods

pkg/rpcclient/management/management.go

@@ -7,10 +7,12 @@ various methods to perform state-changing calls.
 package management
 
 import (
+	"encoding/binary"
 	"encoding/json"
 	"fmt"
 	"math/big"
 
+	"github.com/google/uuid"
 	"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"
@@ -20,11 +22,15 @@ import (
 	"github.com/nspcc-dev/neo-go/pkg/smartcontract/manifest"
 	"github.com/nspcc-dev/neo-go/pkg/smartcontract/nef"
 	"github.com/nspcc-dev/neo-go/pkg/util"
+	"github.com/nspcc-dev/neo-go/pkg/vm/stackitem"
 )
 
 // Invoker is used by ContractReader to call various methods.
 type Invoker interface {
 	Call(contract util.Uint160, operation string, params ...interface{}) (*result.Invoke, error)
+	CallAndExpandIterator(contract util.Uint160, method string, maxItems int, params ...interface{}) (*result.Invoke, error)
+	TerminateSession(sessionID uuid.UUID) error
+	TraverseIterator(sessionID uuid.UUID, iterator *result.Iterator, num int) ([]stackitem.Item, error)
 }
 
 // Actor is used by Contract to create and send transactions.
@@ -55,6 +61,19 @@ type Contract struct {
 	actor Actor
 }
 
+// IDHash is an ID/Hash pair returned by the iterator from the GetContractHashes method.
+type IDHash struct {
+	ID   int32
+	Hash util.Uint160
+}
+
+// HashesIterator is used for iterating over GetContractHashes results.
+type HashesIterator struct {
+	client   Invoker
+	session  uuid.UUID
+	iterator result.Iterator
+}
+
 // Hash stores the hash of the native ContractManagement contract.
 var Hash = state.CreateNativeContractHash(nativenames.Management)
 
@@ -83,7 +102,16 @@ func New(actor Actor) *Contract {
 // getcontractstate RPC API that has more options and works faster than going
 // via contract invocation.
 func (c *ContractReader) GetContract(hash util.Uint160) (*state.Contract, error) {
-	itm, err := unwrap.Item(c.invoker.Call(Hash, "getContract", hash))
+	return unwrapContract(c.invoker.Call(Hash, "getContract", hash))
+}
+
+// GetContractByID allows to get contract data from its ID.
+func (c *ContractReader) GetContractByID(id int32) (*state.Contract, error) {
+	return unwrapContract(c.invoker.Call(Hash, "getContractById", id))
+}
+
+func unwrapContract(r *result.Invoke, err error) (*state.Contract, error) {
+	itm, err := unwrap.Item(r, err)
 	if err != nil {
 		return nil, err
 	}
@@ -95,6 +123,88 @@ func (c *ContractReader) GetContract(hash util.Uint160) (*state.Contract, error)
 	return res, nil
 }
 
+// GetContractHashes returns an iterator that allows to retrieve all ID-hash
+// mappings for non-native contracts. It depends on the server to provide proper
+// session-based iterator, but can also work with expanded one.
+func (c *ContractReader) GetContractHashes() (*HashesIterator, error) {
+	sess, iter, err := unwrap.SessionIterator(c.invoker.Call(Hash, "getContractHashes"))
+	if err != nil {
+		return nil, err
+	}
+
+	return &HashesIterator{
+		client:   c.invoker,
+		iterator: iter,
+		session:  sess,
+	}, nil
+}
+
+// GetContractHashesExpanded is similar to GetContractHashes (uses the same
+// contract method), but can be useful if the server used doesn't support
+// sessions and doesn't expand iterators. It creates a script that will get num
+// of result items from the iterator right in the VM and return them to you. It's
+// only limited by VM stack and GAS available for RPC invocations.
+func (c *ContractReader) GetContractHashesExpanded(num int) ([]IDHash, error) {
+	arr, err := unwrap.Array(c.invoker.CallAndExpandIterator(Hash, "getContractHashes", num))
+	if err != nil {
+		return nil, err
+	}
+	return itemsToIDHashes(arr)
+}
+
+// Next returns the next set of elements from the iterator (up to num of them).
+// It can return less than num elements in case iterator doesn't have that many
+// or zero elements if the iterator has no more elements or the session is
+// expired.
+func (h *HashesIterator) Next(num int) ([]IDHash, error) {
+	items, err := h.client.TraverseIterator(h.session, &h.iterator, num)
+	if err != nil {
+		return nil, err
+	}
+	return itemsToIDHashes(items)
+}
+
+// Terminate closes the iterator session used by HashesIterator (if it's
+// session-based).
+func (h *HashesIterator) Terminate() error {
+	if h.iterator.ID == nil {
+		return nil
+	}
+	return h.client.TerminateSession(h.session)
+}
+
+func itemsToIDHashes(arr []stackitem.Item) ([]IDHash, error) {
+	res := make([]IDHash, len(arr))
+	for i, itm := range arr {
+		str, ok := itm.Value().([]stackitem.Item)
+		if !ok {
+			return nil, fmt.Errorf("item #%d is not a structure %T", i, itm.Value())
+		}
+		if len(str) != 2 {
+			return nil, fmt.Errorf("item #%d has wrong length", i)
+		}
+		bi, err := str[0].TryBytes()
+		if err != nil {
+			return nil, fmt.Errorf("item #%d has wrong ID: %w", i, err)
+		}
+		if len(bi) != 4 {
+			return nil, fmt.Errorf("item #%d has wrong ID: bad length", i)
+		}
+		id := int32(binary.BigEndian.Uint32(bi))
+		hb, err := str[1].TryBytes()
+		if err != nil {
+			return nil, fmt.Errorf("item #%d has wrong hash: %w", i, err)
+		}
+		u160, err := util.Uint160DecodeBytesBE(hb)
+		if err != nil {
+			return nil, fmt.Errorf("item #%d has wrong hash: %w", i, err)
+		}
+		res[i].ID = id
+		res[i].Hash = u160
+	}
+	return res, nil
+}
+
 // GetMinimumDeploymentFee returns the minimal amount of GAS needed to deploy a
 // contract on the network.
 func (c *ContractReader) GetMinimumDeploymentFee() (*big.Int, error) {

pkg/rpcclient/management/management_test.go

@@ -5,6 +5,7 @@ import (
 	"math/big"
 	"testing"
 
+	"github.com/google/uuid"
 	"github.com/nspcc-dev/neo-go/pkg/core/transaction"
 	"github.com/nspcc-dev/neo-go/pkg/neorpc/result"
 	"github.com/nspcc-dev/neo-go/pkg/smartcontract/manifest"
@@ -43,6 +44,15 @@ func (t *testAct) MakeUnsignedRun(script []byte, attrs []transaction.Attribute)
 func (t *testAct) SendRun(script []byte) (util.Uint256, uint32, error) {
 	return t.txh, t.vub, t.err
 }
+func (t *testAct) CallAndExpandIterator(contract util.Uint160, method string, maxItems int, params ...interface{}) (*result.Invoke, error) {
+	return t.res, t.err
+}
+func (t *testAct) TerminateSession(sessionID uuid.UUID) error {
+	return t.err
+}
+func (t *testAct) TraverseIterator(sessionID uuid.UUID, iterator *result.Iterator, num int) ([]stackitem.Item, error) {
+	return t.res.Stack, t.err
+}
 
 func TestReader(t *testing.T) {
 	ta := new(testAct)
@@ -51,6 +61,8 @@ func TestReader(t *testing.T) {
 	ta.err = errors.New("")
 	_, err := man.GetContract(util.Uint160{1, 2, 3})
 	require.Error(t, err)
+	_, err = man.GetContractByID(1)
+	require.Error(t, err)
 	_, err = man.GetMinimumDeploymentFee()
 	require.Error(t, err)
 	_, err = man.HasMethod(util.Uint160{1, 2, 3}, "method", 0)
@@ -65,6 +77,8 @@ func TestReader(t *testing.T) {
 	}
 	_, err = man.GetContract(util.Uint160{1, 2, 3})
 	require.Error(t, err)
+	_, err = man.GetContractByID(1)
+	require.Error(t, err)
 	fee, err := man.GetMinimumDeploymentFee()
 	require.NoError(t, err)
 	require.Equal(t, big.NewInt(42), fee)
@@ -80,6 +94,8 @@ func TestReader(t *testing.T) {
 	}
 	_, err = man.GetContract(util.Uint160{1, 2, 3})
 	require.Error(t, err)
+	_, err = man.GetContractByID(1)
+	require.Error(t, err)
 	hm, err = man.HasMethod(util.Uint160{1, 2, 3}, "method", 0)
 	require.NoError(t, err)
 	require.False(t, hm)
@@ -92,6 +108,8 @@ func TestReader(t *testing.T) {
 	}
 	_, err = man.GetContract(util.Uint160{1, 2, 3})
 	require.Error(t, err)
+	_, err = man.GetContractByID(1)
+	require.Error(t, err)
 
 	nefFile, _ := nef.NewFile([]byte{1, 2, 3})
 	nefBytes, _ := nefFile.Bytes()
@@ -114,6 +132,109 @@ func TestReader(t *testing.T) {
 	require.Equal(t, int32(1), cs.ID)
 	require.Equal(t, uint16(0), cs.UpdateCounter)
 	require.Equal(t, util.Uint160{1, 2, 3}, cs.Hash)
+	cs2, err := man.GetContractByID(1)
+	require.NoError(t, err)
+	require.Equal(t, cs, cs2)
+}
+
+func TestGetContractHashes(t *testing.T) {
+	ta := &testAct{}
+	man := NewReader(ta)
+
+	ta.err = errors.New("")
+	_, err := man.GetContractHashes()
+	require.Error(t, err)
+	_, err = man.GetContractHashesExpanded(5)
+	require.Error(t, err)
+
+	ta.err = nil
+	iid := uuid.New()
+	ta.res = &result.Invoke{
+		State: "HALT",
+		Stack: []stackitem.Item{
+			stackitem.NewInterop(result.Iterator{
+				ID: &iid,
+			}),
+		},
+	}
+	_, err = man.GetContractHashes()
+	require.Error(t, err)
+
+	// Session-based iterator.
+	sid := uuid.New()
+	ta.res = &result.Invoke{
+		Session: sid,
+		State:   "HALT",
+		Stack: []stackitem.Item{
+			stackitem.NewInterop(result.Iterator{
+				ID: &iid,
+			}),
+		},
+	}
+	iter, err := man.GetContractHashes()
+	require.NoError(t, err)
+
+	ta.res = &result.Invoke{
+		Stack: []stackitem.Item{
+			stackitem.Make([]stackitem.Item{
+				stackitem.Make([]byte{0, 0, 0, 1}),
+				stackitem.Make(util.Uint160{1, 2, 3}.BytesBE()),
+			}),
+		},
+	}
+	vals, err := iter.Next(10)
+	require.NoError(t, err)
+	require.Equal(t, 1, len(vals))
+	require.Equal(t, IDHash{
+		ID:   1,
+		Hash: util.Uint160{1, 2, 3},
+	}, vals[0])
+
+	ta.err = errors.New("")
+	_, err = iter.Next(1)
+	require.Error(t, err)
+
+	err = iter.Terminate()
+	require.Error(t, err)
+
+	// Value-based iterator.
+	ta.err = nil
+	ta.res = &result.Invoke{
+		State: "HALT",
+		Stack: []stackitem.Item{
+			stackitem.NewInterop(result.Iterator{
+				Values: []stackitem.Item{stackitem.NewStruct([]stackitem.Item{
+					stackitem.Make([]byte{0, 0, 0, 1}),
+					stackitem.Make(util.Uint160{1, 2, 3}.BytesBE()),
+				})},
+			}),
+		},
+	}
+	iter, err = man.GetContractHashes()
+	require.NoError(t, err)
+
+	ta.err = errors.New("")
+	err = iter.Terminate()
+	require.NoError(t, err)
+
+	// Expanded
+	ta.err = nil
+	ta.res = &result.Invoke{
+		State: "HALT",
+		Stack: []stackitem.Item{
+			stackitem.Make([]stackitem.Item{stackitem.Make([]stackitem.Item{
+				stackitem.Make([]byte{0, 0, 0, 1}),
+				stackitem.Make(util.Uint160{1, 2, 3}.BytesBE()),
+			})}),
+		},
+	}
+	vals, err = man.GetContractHashesExpanded(5)
+	require.NoError(t, err)
+	require.Equal(t, 1, len(vals))
+	require.Equal(t, IDHash{
+		ID:   1,
+		Hash: util.Uint160{1, 2, 3},
+	}, vals[0])
 }
 
 func TestSetMinimumDeploymentFee(t *testing.T) {
@@ -204,3 +325,31 @@ func TestDeploy(t *testing.T) {
 
 	// Unfortunately, manifest _always_ marshals successfully (or panics).
 }
+
+func TestItemsToIDHashesErrors(t *testing.T) {
+	for name, input := range map[string][]stackitem.Item{
+		"not a struct": {stackitem.Make(1)},
+		"wrong length": {stackitem.Make([]stackitem.Item{})},
+		"wrong id": {stackitem.Make([]stackitem.Item{
+			stackitem.Make([]stackitem.Item{}),
+			stackitem.Make(util.Uint160{1, 2, 3}.BytesBE()),
+		})},
+		"lengthy id": {stackitem.Make([]stackitem.Item{
+			stackitem.Make(util.Uint160{1, 2, 3}.BytesBE()),
+			stackitem.Make(util.Uint160{1, 2, 3}.BytesBE()),
+		})},
+		"not a good hash": {stackitem.Make([]stackitem.Item{
+			stackitem.Make([]byte{0, 0, 0, 1}),
+			stackitem.Make([]stackitem.Item{}),
+		})},
+		"not a good u160 hash": {stackitem.Make([]stackitem.Item{
+			stackitem.Make([]byte{0, 0, 0, 1}),
+			stackitem.Make(util.Uint256{1, 2, 3}.BytesBE()),
+		})},
+	} {
+		t.Run(name, func(t *testing.T) {
+			_, err := itemsToIDHashes(input)
+			require.Error(t, err)
+		})
+	}
+}

pkg/services/rpcsrv/client_test.go

@@ -256,6 +256,9 @@ func TestClientManagementContract(t *testing.T) {
 	cs2, err := c.GetContractStateByHash(gas.Hash)
 	require.NoError(t, err)
 	require.Equal(t, cs2, cs1)
+	cs1, err = manReader.GetContractByID(-6)
+	require.NoError(t, err)
+	require.Equal(t, cs2, cs1)
 
 	ret, err := manReader.HasMethod(gas.Hash, "transfer", 4)
 	require.NoError(t, err)
@@ -275,6 +278,25 @@ func TestClientManagementContract(t *testing.T) {
 	}})
 	require.NoError(t, err)
 
+	ids, err := manReader.GetContractHashesExpanded(10)
+	require.NoError(t, err)
+	ctrs := make([]management.IDHash, 0)
+	for i, s := range []string{testContractHash, verifyContractHash, verifyWithArgsContractHash, nnsContractHash, nfsoContractHash, storageContractHash} {
+		h, err := util.Uint160DecodeStringLE(s)
+		require.NoError(t, err)
+		ctrs = append(ctrs, management.IDHash{ID: int32(i) + 1, Hash: h})
+	}
+	require.Equal(t, ctrs, ids)
+
+	iter, err := manReader.GetContractHashes()
+	require.NoError(t, err)
+	ids, err = iter.Next(3)
+	require.NoError(t, err)
+	require.Equal(t, ctrs[:3], ids)
+	ids, err = iter.Next(10)
+	require.NoError(t, err)
+	require.Equal(t, ctrs[3:], ids)
+
 	man := management.New(act)
 
 	txfee, err := man.SetMinimumDeploymentFeeUnsigned(big.NewInt(1 * 1_0000_0000))