add ContainsKeyPair to Wallet

neo-project · Feb 22, 2025 · 3de0e99 · 3de0e99
1 parent 130a060
commit 3de0e99
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 32 deletions.
diff --git a/src/Neo/Wallets/Wallet.cs b/src/Neo/Wallets/Wallet.cs
@@ -665,7 +665,7 @@ public bool Sign(ContractParametersContext context)
         /// </summary>
         /// <param name="signData">The data to sign.</param>
         /// <param name="publicKey">The public key.</param>
-        /// <returns>The signature, or null if was not found.</returns>
+        /// <returns>The signature, or null if the public key was not found or the corresponding private key is not available.</returns>
         public byte[] Sign(byte[] signData, ECPoint publicKey)
         {
             var account = GetAccount(publicKey);
@@ -675,19 +675,18 @@ public byte[] Sign(byte[] signData, ECPoint publicKey)
         }
 
         /// <summary>
-        /// Gets the index of the first key pair(public key and corresponding private key) that this Wallet owns in the list.
+        /// Checks if the wallet contains the specified public key and the corresponding private key.
+        /// If the wallet has the public key but not the private key, it will return <see langword="false"/>.
         /// </summary>
-        /// <param name="publicKeys">The public key list.</param>
-        /// <returns>The index of the public key in the list; if the public key is not found, return -1.</returns>
-        public int GetMyIndex(IReadOnlyList<ECPoint> publicKeys)
+        /// <param name="publicKey">The public key.</param>
+        /// <returns>
+        /// <see langword="true"/> if the wallet contains the specified public key and the corresponding private key;
+        /// otherwise, <see langword="false"/>.
+        /// </returns>
+        public bool ContainsKeyPair(ECPoint publicKey)
         {
-            for (int index = 0; index < publicKeys.Count; index++)
-            {
-                var account = GetAccount(publicKeys[index]);
-                if (account != null && account.HasKey)
-                    return index;
-            }
-            return -1;
+            var account = GetAccount(publicKey);
+            return account != null && account.HasKey;
         }
 
         /// <summary>

diff --git a/src/Plugins/DBFTPlugin/Consensus/ConsensusContext.cs b/src/Plugins/DBFTPlugin/Consensus/ConsensusContext.cs
@@ -246,16 +246,14 @@ public void Reset(byte viewNumber)
                     }
                 }
 
-                var index = _wallet.GetMyIndex(Validators);
-                if (index >= 0 && index < Validators.Length)
+                _myPublicKey = null;
+                for (int i = 0; i < Validators.Length; i++)
                 {
-                    MyIndex = index;
+                    // ContainsKeyPair may be called multiple times
+                    if (!_wallet.ContainsKeyPair(Validators[i])) continue;
+                    MyIndex = i;
                     _myPublicKey = Validators[MyIndex];
-                }
-                else
-                {
-                    MyIndex = -1;
-                    _myPublicKey = null;
+                    break;
                 }
                 cachedMessages = new Dictionary<UInt256, ConsensusMessage>();
             }

diff --git a/tests/Neo.UnitTests/Wallets/UT_Wallet.cs b/tests/Neo.UnitTests/Wallets/UT_Wallet.cs
@@ -449,30 +449,40 @@ public void TestSign()
         }
 
         [TestMethod]
-        public void TestIndexOf()
+        public void TestContainsKeyPair()
         {
             MyWallet wallet = new();
-            var index = wallet.GetMyIndex([glkey.PublicKey]);
-            Assert.AreEqual(-1, index);
+            var contains = wallet.ContainsKeyPair(glkey.PublicKey);
+            Assert.IsFalse(contains);
 
             wallet.CreateAccount(glkey.PrivateKey);
 
-            index = wallet.GetMyIndex([glkey.PublicKey]);
-            Assert.AreEqual(0, index);
+            contains = wallet.ContainsKeyPair(glkey.PublicKey);
+            Assert.IsTrue(contains);
 
             var key = new byte[32];
             Array.Fill(key, (byte)0x01);
 
             var pair = new KeyPair(key);
-            index = wallet.GetMyIndex([pair.PublicKey]);
-            Assert.AreEqual(-1, index);
-
-            index = wallet.GetMyIndex([pair.PublicKey, glkey.PublicKey]);
-            Assert.AreEqual(1, index);
+            contains = wallet.ContainsKeyPair(pair.PublicKey);
+            Assert.IsFalse(contains);
 
             wallet.CreateAccount(pair.PrivateKey);
-            index = wallet.GetMyIndex([pair.PublicKey, glkey.PublicKey]);
-            Assert.AreEqual(0, index);
+            contains = wallet.ContainsKeyPair(pair.PublicKey);
+            Assert.IsTrue(contains);
+
+            contains = wallet.ContainsKeyPair(glkey.PublicKey);
+            Assert.IsTrue(contains);
+
+            key = new byte[32];
+            Array.Fill(key, (byte)0x02);
+
+            pair = new KeyPair(key);
+            var scriptHash = Contract.CreateSignatureRedeemScript(pair.PublicKey).ToScriptHash();
+            wallet.CreateAccount(scriptHash);
+
+            contains = wallet.ContainsKeyPair(pair.PublicKey);
+            Assert.IsFalse(contains); // no private key
         }
     }
 }