From 9bdd8151af437d629b6fd093df8d53822016c32e Mon Sep 17 00:00:00 2001
From: Anna Shaleva <shaleva.ann@nspcc.ru>
Date: Wed, 6 Jul 2022 17:08:53 +0300
Subject: [PATCH] rpc: restrict (*Client).TraverseIterator with single RPC call

Do not unwrap the whole set of iterator values even on demand.
---
 pkg/rpc/client/rpc.go         | 38 +++++++++++++++--------------------
 pkg/rpc/server/client_test.go |  7 +------
 2 files changed, 17 insertions(+), 28 deletions(-)

diff --git a/pkg/rpc/client/rpc.go b/pkg/rpc/client/rpc.go
index 7fc0931ee..a35043f43 100644
--- a/pkg/rpc/client/rpc.go
+++ b/pkg/rpc/client/rpc.go
@@ -1150,36 +1150,30 @@ func (c *Client) GetNativeContractHash(name string) (util.Uint160, error) {
 // TraverseIterator returns a set of iterator values (maxItemsCount at max) for
 // 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. 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.
+// If maxItemsCount is non-positive, then config.DefaultMaxIteratorResultItems
+// iterator values will be returned using single `traverseiterator` call.
+// 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 {
 		maxItemsCount = config.DefaultMaxIteratorResultItems
-		traverseAll = true
 	}
 	var (
-		result []stackitem.Item
 		params = request.NewRawParams(sessionID.String(), iteratorID.String(), maxItemsCount)
+		resp   []json.RawMessage
 	)
-	for {
-		var resp []json.RawMessage
-		if err := c.performRequest("traverseiterator", params, &resp); err != nil {
-			return nil, err
-		}
-		for i, iBytes := range resp {
-			itm, err := stackitem.FromJSONWithTypes(iBytes)
-			if err != nil {
-				return nil, fmt.Errorf("failed to unmarshal %d-th iterator value: %w", i, err)
-			}
-			result = append(result, itm)
-		}
-		if len(resp) < maxItemsCount || !traverseAll {
-			break
+	if err := c.performRequest("traverseiterator", params, &resp); err != nil {
+		return nil, err
+	}
+	result := make([]stackitem.Item, len(resp))
+	for i, iBytes := range resp {
+		itm, err := stackitem.FromJSONWithTypes(iBytes)
+		if err != nil {
+			return nil, fmt.Errorf("failed to unmarshal %d-th iterator value: %w", i, err)
 		}
+		result[i] = itm
 	}
 
 	return result, nil
diff --git a/pkg/rpc/server/client_test.go b/pkg/rpc/server/client_test.go
index 29975afd9..9b45ecfaa 100644
--- a/pkg/rpc/server/client_test.go
+++ b/pkg/rpc/server/client_test.go
@@ -1163,12 +1163,7 @@ func TestClient_IteratorSessions(t *testing.T) {
 
 		set, err := c.TraverseIterator(sID, iID, -1)
 		require.NoError(t, err)
-		require.Equal(t, storageItemsCount, len(set))
-
-		// No more items should be left.
-		set, err = c.TraverseIterator(sID, iID, -1)
-		require.NoError(t, err)
-		require.Equal(t, 0, len(set))
+		require.Equal(t, config.DefaultMaxIteratorResultItems, len(set))
 	})
 
 	t.Run("traverse, concurrent access", func(t *testing.T) {