Fix vote (#1792)

* fix vote * add ut * throw exception * fix * fix ut Co-authored-by: Tommo-L <[email protected]> Co-authored-by: erikzhang <[email protected]>
neo-project · Jul 28, 2020 · 90105ef · 90105ef
1 parent 13f303d
commit 90105ef
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 11 deletions.
diff --git a/src/neo/SmartContract/Native/Tokens/NeoToken.cs b/src/neo/SmartContract/Native/Tokens/NeoToken.cs
@@ -141,10 +141,15 @@ private bool UnregisterCandidate(ApplicationEngine engine, ECPoint pubkey)
         private bool Vote(ApplicationEngine engine, UInt160 account, ECPoint voteTo)
         {
             if (!engine.CheckWitnessInternal(account)) return false;
-            StorageKey key_account = CreateStorageKey(Prefix_Account).Add(account);
-            if (engine.Snapshot.Storages.TryGet(key_account) is null) return false;
-            StorageItem storage_account = engine.Snapshot.Storages.GetAndChange(key_account);
-            NeoAccountState state_account = storage_account.GetInteroperable<NeoAccountState>();
+            NeoAccountState state_account = engine.Snapshot.Storages.GetAndChange(CreateStorageKey(Prefix_Account).Add(account))?.GetInteroperable<NeoAccountState>();
+            if (state_account is null) return false;
+            CandidateState validator_new = null;
+            if (voteTo != null)
+            {
+                validator_new = engine.Snapshot.Storages.GetAndChange(CreateStorageKey(Prefix_Candidate).Add(voteTo))?.GetInteroperable<CandidateState>();
+                if (validator_new is null) return false;
+                if (!validator_new.Registered) return false;
+            }
             if (state_account.VoteTo is null ^ voteTo is null)
             {
                 StorageItem item = engine.Snapshot.Storages.GetAndChange(CreateStorageKey(Prefix_VotersCount));
@@ -163,14 +168,9 @@ private bool Vote(ApplicationEngine engine, UInt160 account, ECPoint voteTo)
                     engine.Snapshot.Storages.Delete(key);
             }
             state_account.VoteTo = voteTo;
-            if (voteTo != null)
+            if (validator_new != null)
             {
-                StorageKey key = CreateStorageKey(Prefix_Candidate).Add(voteTo);
-                if (engine.Snapshot.Storages.TryGet(key) is null) return false;
-                StorageItem storage_validator = engine.Snapshot.Storages.GetAndChange(key);
-                CandidateState state_validator = storage_validator.GetInteroperable<CandidateState>();
-                if (!state_validator.Registered) return false;
-                state_validator.Votes += state_account.Balance;
+                validator_new.Votes += state_account.Balance;
             }
             return true;
         }

diff --git a/tests/neo.UnitTests/SmartContract/Native/Tokens/UT_NeoToken.cs b/tests/neo.UnitTests/SmartContract/Native/Tokens/UT_NeoToken.cs
@@ -72,6 +72,23 @@ public void Check_Vote()
             ret.Result.Should().BeFalse();
             ret.State.Should().BeTrue();
 
+            // no registered
+
+            var accountState = snapshot.Storages.TryGet(CreateStorageKey(20, from)).GetInteroperable<NeoAccountState>();
+            accountState.VoteTo = null;
+            ret = Check_Vote(snapshot, from, ECCurve.Secp256r1.G.ToArray(), true);
+            ret.Result.Should().BeFalse();
+            ret.State.Should().BeTrue();
+            accountState.VoteTo.Should().BeNull();
+
+            // normal case
+
+            snapshot.Storages.Add(CreateStorageKey(33, ECCurve.Secp256r1.G.ToArray()), new StorageItem(new CandidateState()));
+            ret = Check_Vote(snapshot, from, ECCurve.Secp256r1.G.ToArray(), true);
+            ret.Result.Should().BeTrue();
+            ret.State.Should().BeTrue();
+            accountState.VoteTo.Should().Be(ECCurve.Secp256r1.G);
+
             // TODO: More votes tests
         }