[Add] GetStorage for ContractState (neo-project#3615)

* [Add] GetStorage for ContractState * Add more parameters and unit tests * Changed names * Update src/Neo/Extensions/SmartContract/ContractStateExtensions.cs Co-authored-by: nan01ab <[email protected]> * Update src/Neo/Extensions/SmartContract/ContractStateExtensions.cs Co-authored-by: nan01ab <[email protected]> * Update src/Neo/Extensions/SmartContract/ContractStateExtensions.cs Co-authored-by: nan01ab <[email protected]> * Updated RPC server to use `FindStorage` * updated namespace * Fixed errors --------- Co-authored-by: Shargon <[email protected]> Co-authored-by: Jimmy <[email protected]> Co-authored-by: nan01ab <[email protected]>
nan01ab · Dec 27, 2024 · f65c193 · f65c193
1 parent 2c521a0
commit f65c193
Show file tree

Hide file tree

Showing 3 changed files with 154 additions and 3 deletions.
diff --git a/src/Neo/Extensions/SmartContract/ContractStateExtensions.cs b/src/Neo/Extensions/SmartContract/ContractStateExtensions.cs
@@ -0,0 +1,88 @@
+// Copyright (C) 2015-2024 The Neo Project.
+//
+// ContractStateExtensions.cs file belongs to the neo project and is free
+// software distributed under the MIT software license, see the
+// accompanying file LICENSE in the main directory of the
+// repository or http://www.opensource.org/licenses/mit-license.php
+// for more details.
+//
+// Redistribution and use in source and binary forms with or without
+// modifications are permitted.
+
+using Neo.Persistence;
+using Neo.SmartContract;
+using Neo.SmartContract.Native;
+using System;
+using System.Collections.Generic;
+
+namespace Neo.Extensions
+{
+    public static class ContractStateExtensions
+    {
+        /// <summary>
+        /// Get Storage value by storage map key.
+        /// </summary>
+        /// <param name="contractState"></param>
+        /// <param name="snapshot">Snapshot of the database.</param>
+        /// <param name="storageKey">Key in the storage map.</param>
+        /// <returns>Storage value of the item.</returns>
+        /// <exception cref="ArgumentNullException"><paramref name="contractState"/> or <paramref name="snapshot"/> is null</exception>
+        public static StorageItem GetStorage(this ContractState contractState, DataCache snapshot, byte[] storageKey)
+        {
+            if (contractState is null)
+                throw new ArgumentNullException(nameof(contractState));
+
+            if (snapshot is null)
+                throw new ArgumentNullException(nameof(snapshot));
+
+            if (storageKey is null)
+                storageKey = [];
+
+            return snapshot.TryGet(StorageKey.CreateSearchPrefix(contractState.Id, storageKey));
+        }
+
+        /// <summary>
+        /// All storage items stored in the given contract.
+        /// </summary>
+        /// <param name="contractState"></param>
+        /// <param name="snapshot">Snapshot of the database.</param>
+        /// <param name="prefix">Prefix of the key.</param>
+        /// <param name="seekDirection"></param>
+        /// <returns>All storage of the given contract.</returns>
+        /// <exception cref="ArgumentNullException"><paramref name="contractState"/> or <paramref name="snapshot"/> is null</exception>
+        public static IEnumerable<(StorageKey Key, StorageItem Value)> FindStorage(this ContractState contractState, DataCache snapshot, byte[] prefix = null, SeekDirection seekDirection = SeekDirection.Forward)
+        {
+            if (contractState is null)
+                throw new ArgumentNullException(nameof(contractState));
+
+            if (snapshot is null)
+                throw new ArgumentNullException(nameof(snapshot));
+
+            if (prefix is null)
+                prefix = [];
+
+            return snapshot.Find(StorageKey.CreateSearchPrefix(contractState.Id, prefix), seekDirection);
+        }
+
+        /// <summary>
+        /// All storage items stored in the given contract.
+        /// </summary>
+        /// <param name="contractManagement"></param>
+        /// <param name="snapshot">Snapshot of the database.</param>
+        /// <param name="prefix">Prefix of the key.</param>
+        /// <param name="contractId">Id of the contract.</param>
+        /// <param name="seekDirection"></param>
+        /// <returns>All storage of the given contract.</returns>
+        /// <exception cref="ArgumentNullException"><paramref name="snapshot"/> is null</exception>
+        public static IEnumerable<(StorageKey Key, StorageItem Value)> FindContractStorage(this ContractManagement contractManagement, DataCache snapshot, int contractId, byte[] prefix = null, SeekDirection seekDirection = SeekDirection.Forward)
+        {
+            if (snapshot is null)
+                throw new ArgumentNullException(nameof(snapshot));
+
+            if (prefix is null)
+                prefix = [];
+
+            return snapshot.Find(StorageKey.CreateSearchPrefix(contractId, prefix), seekDirection);
+        }
+    }
+}
diff --git a/src/Plugins/RpcServer/RpcServer.Blockchain.cs b/src/Plugins/RpcServer/RpcServer.Blockchain.cs
@@ -10,7 +10,6 @@
 // modifications are permitted.
 
 using Neo.Extensions;
-using Neo.IO;
 using Neo.Json;
 using Neo.Network.P2P.Payloads;
 using Neo.Plugins.RpcServer.Model;
@@ -265,14 +264,13 @@ protected internal virtual JToken FindStorage(ContractNameOrHashOrId contractNam
             }
 
             byte[] prefix = Result.Ok_Or(() => Convert.FromBase64String(base64KeyPrefix), RpcError.InvalidParams.WithData($"Invalid Base64 string{base64KeyPrefix}"));
-            byte[] prefix_key = StorageKey.CreateSearchPrefix(id, prefix);
 
             JObject json = new();
             JArray jarr = new();
             int pageSize = settings.FindStoragePageSize;
             int i = 0;
 
-            using (var iter = snapshot.Find(prefix_key).Skip(count: start).GetEnumerator())
+            using (var iter = NativeContract.ContractManagement.FindContractStorage(snapshot, id, prefix).Skip(count: start).GetEnumerator())
             {
                 var hasMore = false;
                 while (iter.MoveNext())

diff --git a/tests/Neo.UnitTests/Extensions/UT_ContractStateExtensions.cs b/tests/Neo.UnitTests/Extensions/UT_ContractStateExtensions.cs
@@ -0,0 +1,65 @@
+// Copyright (C) 2015-2024 The Neo Project.
+//
+// UT_ContractStateExtensions.cs file belongs to the neo project and is free
+// software distributed under the MIT software license, see the
+// accompanying file LICENSE in the main directory of the
+// repository or http://www.opensource.org/licenses/mit-license.php
+// for more details.
+//
+// Redistribution and use in source and binary forms with or without
+// modifications are permitted.
+
+using Microsoft.VisualStudio.TestTools.UnitTesting;
+using Neo.Extensions;
+using Neo.SmartContract;
+using Neo.SmartContract.Native;
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Text;
+using System.Threading.Tasks;
+
+namespace Neo.UnitTests.Extensions
+{
+    [TestClass]
+    public class UT_ContractStateExtensions
+    {
+        private NeoSystem system;
+
+        [TestInitialize]
+        public void Initialize()
+        {
+            system = TestBlockchain.TheNeoSystem;
+        }
+
+        [TestCleanup]
+        public void Clean()
+        {
+            TestBlockchain.ResetStore();
+        }
+
+        [TestMethod]
+        public void TestGetStorage()
+        {
+            var contractStorage = NativeContract.ContractManagement.FindContractStorage(system.StoreView, NativeContract.NEO.Id);
+
+            Assert.IsNotNull(contractStorage);
+
+            var neoContract = NativeContract.ContractManagement.GetContractById(system.StoreView, NativeContract.NEO.Id);
+
+            contractStorage = neoContract.FindStorage(system.StoreView);
+
+            Assert.IsNotNull(contractStorage);
+
+            contractStorage = neoContract.FindStorage(system.StoreView, [20]);
+
+            Assert.IsNotNull(contractStorage);
+
+            UInt160 address = "0x9f8f056a53e39585c7bb52886418c7bed83d126b";
+            var item = neoContract.GetStorage(system.StoreView, [20, .. address.ToArray()]);
+
+            Assert.IsNotNull(item);
+            Assert.AreEqual(100_000_000, item.GetInteroperable<AccountState>().Balance);
+        }
+    }
+}