Fixes GetRegisteredValidators() (neo-project#991)

Tommo-L · Jun 22, 2020 · 343d4da · 343d4da
1 parent 1a1dbaf
commit 343d4da
Show file tree

Hide file tree

Showing 4 changed files with 37 additions and 19 deletions.
diff --git a/neo.UnitTests/SmartContract/Native/Tokens/UT_NeoToken.cs b/neo.UnitTests/SmartContract/Native/Tokens/UT_NeoToken.cs
@@ -128,6 +128,20 @@ public void Check_RegisterValidator()
             ret.Result.Should().BeTrue();
 
             snapshot.Storages.GetChangeSet().Count().Should().Be(keyCount + 1); // New validator
+
+            // Check GetRegisteredValidators
+
+            var validators = NativeContract.NEO.GetRegisteredValidators(snapshot).OrderBy(u => u.PublicKey).ToArray();
+            var check = Blockchain.StandbyValidators.Select(u => u.EncodePoint(true)).ToList();
+            check.Add(point); // Add the new member
+
+            for (int x = 0; x < validators.Length; x++)
+            {
+                Assert.AreEqual(1, check.RemoveAll(u => u.SequenceEqual(validators[x].PublicKey.EncodePoint(true))));
+                Assert.AreEqual(0, validators[x].Votes);
+            }
+
+            Assert.AreEqual(0, check.Count);
         }
 
         [TestMethod]
@@ -357,4 +371,4 @@ internal static void CheckBalance(byte[] account, DataCache<StorageKey, StorageI
             trackable.Item.IsConstant.Should().Be(false);
         }
     }
-}
+}
diff --git a/neo/Ledger/StorageKey.cs b/neo/Ledger/StorageKey.cs
@@ -13,6 +13,25 @@ public class StorageKey : IEquatable<StorageKey>, ISerializable
 
         int ISerializable.Size => ScriptHash.Size + (Key.Length / 16 + 1) * 17;
 
+        internal static byte[] CreateSearchPrefix(UInt160 hash, byte[] prefix)
+        {
+            using (MemoryStream ms = new MemoryStream())
+            {
+                int index = 0;
+                int remain = prefix.Length;
+                while (remain >= 16)
+                {
+                    ms.Write(prefix, index, 16);
+                    ms.WriteByte(16);
+                    index += 16;
+                    remain -= 16;
+                }
+                if (remain > 0)
+                    ms.Write(prefix, index, remain);
+                return hash.ToArray().Concat(ms.ToArray()).ToArray();
+            }
+        }
+
         void ISerializable.Deserialize(BinaryReader reader)
         {
             ScriptHash = reader.ReadSerializable<UInt160>();

diff --git a/neo/SmartContract/InteropService.NEO.cs b/neo/SmartContract/InteropService.NEO.cs
@@ -10,7 +10,6 @@
 using Neo.VM;
 using Neo.VM.Types;
 using System;
-using System.IO;
 using System.Linq;
 using VMArray = Neo.VM.Types.Array;
 
@@ -345,22 +344,7 @@ private static bool Storage_Find(ApplicationEngine engine)
                 StorageContext context = _interface.GetInterface<StorageContext>();
                 if (!CheckStorageContext(engine, context)) return false;
                 byte[] prefix = engine.CurrentContext.EvaluationStack.Pop().GetByteArray();
-                byte[] prefix_key;
-                using (MemoryStream ms = new MemoryStream())
-                {
-                    int index = 0;
-                    int remain = prefix.Length;
-                    while (remain >= 16)
-                    {
-                        ms.Write(prefix, index, 16);
-                        ms.WriteByte(16);
-                        index += 16;
-                        remain -= 16;
-                    }
-                    if (remain > 0)
-                        ms.Write(prefix, index, remain);
-                    prefix_key = context.ScriptHash.ToArray().Concat(ms.ToArray()).ToArray();
-                }
+                byte[] prefix_key = StorageKey.CreateSearchPrefix(context.ScriptHash, prefix);
                 StorageIterator iterator = engine.AddDisposable(new StorageIterator(engine.Snapshot.Storages.Find(prefix_key).Where(p => p.Key.Key.Take(prefix.Length).SequenceEqual(prefix)).GetEnumerator()));
                 engine.CurrentContext.EvaluationStack.Push(StackItem.FromInterface(iterator));
                 return true;

diff --git a/neo/SmartContract/Native/Tokens/NeoToken.cs b/neo/SmartContract/Native/Tokens/NeoToken.cs
@@ -201,7 +201,8 @@ private StackItem GetRegisteredValidators(ApplicationEngine engine, VMArray args
 
         public IEnumerable<(ECPoint PublicKey, BigInteger Votes)> GetRegisteredValidators(Snapshot snapshot)
         {
-            return snapshot.Storages.Find(new[] { Prefix_Validator }).Select(p =>
+            byte[] prefix_key = StorageKey.CreateSearchPrefix(Hash, new[] { Prefix_Validator });
+            return snapshot.Storages.Find(prefix_key).Select(p =>
             (
                 p.Key.Key.Skip(1).ToArray().AsSerializable<ECPoint>(),
                 ValidatorState.FromByteArray(p.Value.Value).Votes