From fad061f3d9c5c132cc88734a6ac67e18aea4d388 Mon Sep 17 00:00:00 2001
From: Anna Shaleva <shaleva.ann@nspcc.ru>
Date: Wed, 6 Jul 2022 16:55:51 +0300
Subject: [PATCH] rpc: extend iterator-related client functionality

Create a set of functions that are able to work with both session-based
iterators, default unpacked iterators and client-side unpacked
iterators.
---
 cli/wallet/nep11.go           | 12 ++---
 pkg/rpc/client/helper.go      | 13 ++++++
 pkg/rpc/client/native.go      | 32 ++++++++++++-
 pkg/rpc/client/nep11.go       | 84 ++++++++++++++++++++++++++++++++---
 pkg/rpc/client/rpc.go         |  5 ++-
 pkg/rpc/server/client_test.go | 41 +++++++++++++++--
 6 files changed, 169 insertions(+), 18 deletions(-)

diff --git a/cli/wallet/nep11.go b/cli/wallet/nep11.go
index 1a413dbc6..ab1cb4d5b 100644
--- a/cli/wallet/nep11.go
+++ b/cli/wallet/nep11.go
@@ -118,7 +118,7 @@ func newNEP11Commands() []cli.Command {
 		},
 		{
 			Name:      "ownerOfD",
-			Usage:     "print set of owners of divisible NEP-11 token with the specified ID",
+			Usage:     "print set of owners of divisible NEP-11 token with the specified ID (the default MaxIteratorResultItems will be printed at max)",
 			UsageText: "ownerOfD --rpc-endpoint <node> --timeout <time> --token <hash> --id <token-id>",
 			Action:    printNEP11DOwner,
 			Flags: append([]cli.Flag{
@@ -128,7 +128,7 @@ func newNEP11Commands() []cli.Command {
 		},
 		{
 			Name:      "tokensOf",
-			Usage:     "print list of tokens IDs for the specified NFT owner",
+			Usage:     "print list of tokens IDs for the specified NFT owner (the default MaxIteratorResultItems will be printed at max)",
 			UsageText: "tokensOf --rpc-endpoint <node> --timeout <time> --token <hash> --address <addr>",
 			Action:    printNEP11TokensOf,
 			Flags: append([]cli.Flag{
@@ -138,7 +138,7 @@ func newNEP11Commands() []cli.Command {
 		},
 		{
 			Name:      "tokens",
-			Usage:     "print list of tokens IDs minted by the specified NFT (optional method)",
+			Usage:     "print list of tokens IDs minted by the specified NFT (optional method; the default MaxIteratorResultItems will be printed at max)",
 			UsageText: "tokens --rpc-endpoint <node> --timeout <time> --token <hash>",
 			Action:    printNEP11Tokens,
 			Flags: append([]cli.Flag{
@@ -332,7 +332,7 @@ func printNEP11Owner(ctx *cli.Context, divisible bool) error {
 	}
 
 	if divisible {
-		result, err := c.NEP11DOwnerOf(tokenHash.Uint160(), tokenIDBytes)
+		result, err := c.NEP11DUnpackedOwnerOf(tokenHash.Uint160(), tokenIDBytes)
 		if err != nil {
 			return cli.NewExitError(fmt.Sprintf("failed to call NEP-11 divisible `ownerOf` method: %s", err.Error()), 1)
 		}
@@ -370,7 +370,7 @@ func printNEP11TokensOf(ctx *cli.Context) error {
 		return cli.NewExitError(err, 1)
 	}
 
-	result, err := c.NEP11TokensOf(tokenHash.Uint160(), acc.Uint160())
+	result, err := c.NEP11UnpackedTokensOf(tokenHash.Uint160(), acc.Uint160())
 	if err != nil {
 		return cli.NewExitError(fmt.Sprintf("failed to call NEP-11 `tokensOf` method: %s", err.Error()), 1)
 	}
@@ -396,7 +396,7 @@ func printNEP11Tokens(ctx *cli.Context) error {
 		return cli.NewExitError(err, 1)
 	}
 
-	result, err := c.NEP11Tokens(tokenHash.Uint160())
+	result, err := c.NEP11UnpackedTokens(tokenHash.Uint160())
 	if err != nil {
 		return cli.NewExitError(fmt.Sprintf("failed to call optional NEP-11 `tokens` method: %s", err.Error()), 1)
 	}
diff --git a/pkg/rpc/client/helper.go b/pkg/rpc/client/helper.go
index 7c6d04e09..151f8ac26 100644
--- a/pkg/rpc/client/helper.go
+++ b/pkg/rpc/client/helper.go
@@ -237,3 +237,16 @@ func topIterableFromStack(st []stackitem.Item, resultItemType interface{}) ([]in
 	}
 	return result, nil
 }
+
+// topIteratorFromStack returns the top Iterator from the stack.
+func topIteratorFromStack(st []stackitem.Item) (result.Iterator, error) {
+	index := len(st) - 1 // top stack element is the last in the array
+	if t := st[index].Type(); t != stackitem.InteropT {
+		return result.Iterator{}, fmt.Errorf("expected InteropInterface on stack, got %s", t)
+	}
+	iter, ok := st[index].Value().(result.Iterator)
+	if !ok {
+		return result.Iterator{}, fmt.Errorf("failed to deserialize iterable from interop stackitem: invalid value type (Iterator expected)")
+	}
+	return iter, nil
+}
diff --git a/pkg/rpc/client/native.go b/pkg/rpc/client/native.go
index f154c2633..6dd65af51 100644
--- a/pkg/rpc/client/native.go
+++ b/pkg/rpc/client/native.go
@@ -7,10 +7,12 @@ import (
 	"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/native/noderoles"
 	"github.com/nspcc-dev/neo-go/pkg/crypto/keys"
 	"github.com/nspcc-dev/neo-go/pkg/rpc/client/nns"
+	"github.com/nspcc-dev/neo-go/pkg/rpc/response/result"
 	"github.com/nspcc-dev/neo-go/pkg/smartcontract"
 	"github.com/nspcc-dev/neo-go/pkg/util"
 )
@@ -116,8 +118,34 @@ func (c *Client) NNSIsAvailable(nnsHash util.Uint160, name string) (bool, error)
 	return topBoolFromStack(result.Stack)
 }
 
-// NNSGetAllRecords returns all records for a given name from NNS service.
-func (c *Client) NNSGetAllRecords(nnsHash util.Uint160, name string) ([]nns.RecordState, error) {
+// NNSGetAllRecords returns iterator over records for a given name from NNS service.
+// First return value is the session ID, the second one is Iterator itself, the
+// third one is an error. Use TraverseIterator method to traverse iterator values or
+// TerminateSession to terminate opened iterator session. See TraverseIterator and
+// TerminateSession documentation for more details.
+func (c *Client) NNSGetAllRecords(nnsHash util.Uint160, name string) (uuid.UUID, result.Iterator, error) {
+	res, err := c.InvokeFunction(nnsHash, "getAllRecords", []smartcontract.Parameter{
+		{
+			Type:  smartcontract.StringType,
+			Value: name,
+		},
+	}, nil)
+	if err != nil {
+		return uuid.UUID{}, result.Iterator{}, err
+	}
+	err = getInvocationError(res)
+	if err != nil {
+		return uuid.UUID{}, result.Iterator{}, err
+	}
+
+	iter, err := topIteratorFromStack(res.Stack)
+	return res.Session, iter, err
+}
+
+// NNSUnpackedGetAllRecords returns all records for a given name from NNS service. It differs from
+// NNSGetAllRecords in that no iterator session is used to retrieve values from iterator. Instead,
+// unpacking VM script is created and invoked via `invokescript` JSON-RPC call.
+func (c *Client) NNSUnpackedGetAllRecords(nnsHash util.Uint160, name string) ([]nns.RecordState, error) {
 	result, err := c.InvokeAndPackIteratorResults(nnsHash, "getAllRecords", []smartcontract.Parameter{
 		{
 			Type:  smartcontract.StringType,
diff --git a/pkg/rpc/client/nep11.go b/pkg/rpc/client/nep11.go
index d663c839e..4fe86e1b3 100644
--- a/pkg/rpc/client/nep11.go
+++ b/pkg/rpc/client/nep11.go
@@ -3,9 +3,11 @@ package client
 import (
 	"fmt"
 
+	"github.com/google/uuid"
 	"github.com/nspcc-dev/neo-go/pkg/core/transaction"
 	"github.com/nspcc-dev/neo-go/pkg/encoding/address"
 	"github.com/nspcc-dev/neo-go/pkg/io"
+	"github.com/nspcc-dev/neo-go/pkg/rpc/response/result"
 	"github.com/nspcc-dev/neo-go/pkg/smartcontract"
 	"github.com/nspcc-dev/neo-go/pkg/smartcontract/callflag"
 	"github.com/nspcc-dev/neo-go/pkg/smartcontract/manifest"
@@ -82,8 +84,33 @@ func (c *Client) CreateNEP11TransferTx(acc *wallet.Account, tokenHash util.Uint1
 	}}, cosigners...))
 }
 
-// NEP11TokensOf returns an array of token IDs for the specified owner of the specified NFT token.
-func (c *Client) NEP11TokensOf(tokenHash util.Uint160, owner util.Uint160) ([][]byte, error) {
+// NEP11TokensOf returns iterator over token IDs for the specified owner of the
+// specified NFT token. First return value is the session ID, the second one is
+// Iterator itself, the third one is an error. Use TraverseIterator method to
+// traverse iterator values or TerminateSession to terminate opened iterator
+// session. See TraverseIterator and TerminateSession documentation for more details.
+func (c *Client) NEP11TokensOf(tokenHash util.Uint160, owner util.Uint160) (uuid.UUID, result.Iterator, error) {
+	res, err := c.InvokeFunction(tokenHash, "tokensOf", []smartcontract.Parameter{
+		{
+			Type:  smartcontract.Hash160Type,
+			Value: owner,
+		},
+	}, nil)
+	if err != nil {
+		return uuid.UUID{}, result.Iterator{}, err
+	}
+	err = getInvocationError(res)
+	if err != nil {
+		return uuid.UUID{}, result.Iterator{}, err
+	}
+	iter, err := topIteratorFromStack(res.Stack)
+	return res.Session, iter, err
+}
+
+// NEP11UnpackedTokensOf returns an array of token IDs for the specified owner of the specified NFT token.
+// It differs from NEP11TokensOf in that no iterator session is used to retrieve values from iterator.
+// Instead, unpacking VM script is created and invoked via `invokescript` JSON-RPC call.
+func (c *Client) NEP11UnpackedTokensOf(tokenHash util.Uint160, owner util.Uint160) ([][]byte, error) {
 	result, err := c.InvokeAndPackIteratorResults(tokenHash, "tokensOf", []smartcontract.Parameter{
 		{
 			Type:  smartcontract.Hash160Type,
@@ -159,8 +186,33 @@ func (c *Client) NEP11DBalanceOf(tokenHash, owner util.Uint160, tokenID []byte)
 	return c.nepBalanceOf(tokenHash, owner, tokenID)
 }
 
-// NEP11DOwnerOf returns list of the specified NEP-11 divisible token owners.
-func (c *Client) NEP11DOwnerOf(tokenHash util.Uint160, tokenID []byte) ([]util.Uint160, error) {
+// NEP11DOwnerOf returns iterator over the specified NEP-11 divisible token owners. First return value
+// is the session ID, the second one is Iterator itself, the third one is an error. Use TraverseIterator
+// method to traverse iterator values or TerminateSession to terminate opened iterator session. See
+// TraverseIterator and TerminateSession documentation for more details.
+func (c *Client) NEP11DOwnerOf(tokenHash util.Uint160, tokenID []byte) (uuid.UUID, result.Iterator, error) {
+	res, err := c.InvokeFunction(tokenHash, "ownerOf", []smartcontract.Parameter{
+		{
+			Type:  smartcontract.ByteArrayType,
+			Value: tokenID,
+		},
+	}, nil)
+	sessID := res.Session
+	if err != nil {
+		return sessID, result.Iterator{}, err
+	}
+	err = getInvocationError(res)
+	if err != nil {
+		return sessID, result.Iterator{}, err
+	}
+	arr, err := topIteratorFromStack(res.Stack)
+	return sessID, arr, err
+}
+
+// NEP11DUnpackedOwnerOf returns list of the specified NEP-11 divisible token owners. It differs from
+// NEP11DOwnerOf in that no iterator session is used to retrieve values from iterator. Instead,
+// unpacking VM script is created and invoked via `invokescript` JSON-RPC call.
+func (c *Client) NEP11DUnpackedOwnerOf(tokenHash util.Uint160, tokenID []byte) ([]util.Uint160, error) {
 	result, err := c.InvokeAndPackIteratorResults(tokenHash, "ownerOf", []smartcontract.Parameter{
 		{
 			Type:  smartcontract.ByteArrayType,
@@ -208,8 +260,28 @@ func (c *Client) NEP11Properties(tokenHash util.Uint160, tokenID []byte) (*stack
 	return topMapFromStack(result.Stack)
 }
 
-// NEP11Tokens returns list of the tokens minted by the contract.
-func (c *Client) NEP11Tokens(tokenHash util.Uint160) ([][]byte, error) {
+// NEP11Tokens returns iterator over the tokens minted by the contract. First return
+// value is the session ID, the second one is Iterator itself, the third one is an
+// error. Use TraverseIterator method to traverse iterator values or
+// TerminateSession to terminate opened iterator session. See TraverseIterator and
+// TerminateSession documentation for more details.
+func (c *Client) NEP11Tokens(tokenHash util.Uint160) (uuid.UUID, result.Iterator, error) {
+	res, err := c.InvokeFunction(tokenHash, "tokens", []smartcontract.Parameter{}, nil)
+	if err != nil {
+		return uuid.UUID{}, result.Iterator{}, err
+	}
+	err = getInvocationError(res)
+	if err != nil {
+		return uuid.UUID{}, result.Iterator{}, err
+	}
+	iter, err := topIteratorFromStack(res.Stack)
+	return res.Session, iter, err
+}
+
+// NEP11UnpackedTokens returns list of the tokens minted by the contract. It differs from
+// NEP11Tokens in that no iterator session is used to retrieve values from iterator. Instead,
+// unpacking VM script is created and invoked via `invokescript` JSON-RPC call.
+func (c *Client) NEP11UnpackedTokens(tokenHash util.Uint160) ([][]byte, error) {
 	result, err := c.InvokeAndPackIteratorResults(tokenHash, "tokens", []smartcontract.Parameter{}, nil)
 	if err != nil {
 		return nil, err
diff --git a/pkg/rpc/client/rpc.go b/pkg/rpc/client/rpc.go
index a77aeffc4..7fc0931ee 100644
--- a/pkg/rpc/client/rpc.go
+++ b/pkg/rpc/client/rpc.go
@@ -1151,7 +1151,10 @@ func (c *Client) GetNativeContractHash(name string) (util.Uint160, error) {
 // the specified iterator and session. If result contains no elements, then either
 // Iterator has no elements or session was expired and terminated by the server.
 // If maxItemsCount is non-positive, then the full set of iterator values will be
-// returned using several `traverseiterator` calls if needed.
+// returned using several `traverseiterator` calls if needed. Note that iterator
+// session lifetime is restricted by the RPC-server configuration and is being
+// reset each time iterator is accessed. If session won't be accessed within session
+// expiration time, then it will be terminated by the RPC-server automatically.
 func (c *Client) TraverseIterator(sessionID, iteratorID uuid.UUID, maxItemsCount int) ([]stackitem.Item, error) {
 	var traverseAll bool
 	if maxItemsCount <= 0 {
diff --git a/pkg/rpc/server/client_test.go b/pkg/rpc/server/client_test.go
index c29634a2e..29975afd9 100644
--- a/pkg/rpc/server/client_test.go
+++ b/pkg/rpc/server/client_test.go
@@ -978,7 +978,19 @@ func TestClient_NEP11_D(t *testing.T) {
 		require.EqualValues(t, 80, b)
 	})
 	t.Run("OwnerOf", func(t *testing.T) {
-		b, err := c.NEP11DOwnerOf(nfsoHash, token1ID)
+		sessID, iter, err := c.NEP11DOwnerOf(nfsoHash, token1ID)
+		require.NoError(t, err)
+		items, err := c.TraverseIterator(sessID, *iter.ID, config.DefaultMaxIteratorResultItems)
+		require.NoError(t, err)
+		require.Equal(t, 2, len(items))
+		actual1, err := util.Uint160DecodeBytesBE(items[0].Value().([]byte))
+		require.NoError(t, err)
+		actual0, err := util.Uint160DecodeBytesBE(items[1].Value().([]byte))
+		require.NoError(t, err)
+		require.Equal(t, []util.Uint160{priv1, priv0}, []util.Uint160{actual1, actual0})
+	})
+	t.Run("UnpackedOwnerOf", func(t *testing.T) {
+		b, err := c.NEP11DUnpackedOwnerOf(nfsoHash, token1ID)
 		require.NoError(t, err)
 		require.Equal(t, []util.Uint160{priv1, priv0}, b)
 	})
@@ -1032,7 +1044,26 @@ func TestClient_NNS(t *testing.T) {
 		require.Error(t, err)
 	})
 	t.Run("NNSGetAllRecords, good", func(t *testing.T) {
-		rss, err := c.NNSGetAllRecords(nnsHash, "neo.com")
+		sess, iter, err := c.NNSGetAllRecords(nnsHash, "neo.com")
+		require.NoError(t, err)
+		arr, err := c.TraverseIterator(sess, *iter.ID, config.DefaultMaxIteratorResultItems)
+		require.NoError(t, err)
+		require.Equal(t, 1, len(arr))
+		rs := arr[0].Value().([]stackitem.Item)
+		require.Equal(t, 3, len(rs))
+		actual := nns.RecordState{
+			Name: string(rs[0].Value().([]byte)),
+			Type: nns.RecordType(rs[1].Value().(*big.Int).Int64()),
+			Data: string(rs[2].Value().([]byte)),
+		}
+		require.Equal(t, nns.RecordState{
+			Name: "neo.com",
+			Type: nns.A,
+			Data: "1.2.3.4",
+		}, actual)
+	})
+	t.Run("NNSUnpackedGetAllRecords, good", func(t *testing.T) {
+		rss, err := c.NNSUnpackedGetAllRecords(nnsHash, "neo.com")
 		require.NoError(t, err)
 		require.Equal(t, []nns.RecordState{
 			{
@@ -1043,7 +1074,11 @@ func TestClient_NNS(t *testing.T) {
 		}, rss)
 	})
 	t.Run("NNSGetAllRecords, bad", func(t *testing.T) {
-		_, err := c.NNSGetAllRecords(nnsHash, "neopython.com")
+		_, _, err := c.NNSGetAllRecords(nnsHash, "neopython.com")
+		require.Error(t, err)
+	})
+	t.Run("NNSUnpackedGetAllRecords, bad", func(t *testing.T) {
+		_, err := c.NNSUnpackedGetAllRecords(nnsHash, "neopython.com")
 		require.Error(t, err)
 	})
 }