From 90105ef85b99422bb1788dd1d6b38e4ec883f2bd Mon Sep 17 00:00:00 2001
From: Luchuan <luchuan@ngd.neo.org>
Date: Tue, 28 Jul 2020 16:41:23 +0800
Subject: [PATCH] Fix vote (#1792)

* fix vote

* add ut

* throw exception

* fix

* fix ut

Co-authored-by: Tommo-L <luchuan@neo.org>
Co-authored-by: erikzhang <erik@neo.org>
---
 .../SmartContract/Native/Tokens/NeoToken.cs   | 22 +++++++++----------
 .../Native/Tokens/UT_NeoToken.cs              | 17 ++++++++++++++
 2 files changed, 28 insertions(+), 11 deletions(-)

diff --git a/src/neo/SmartContract/Native/Tokens/NeoToken.cs b/src/neo/SmartContract/Native/Tokens/NeoToken.cs
index 4024de5690..aa8ea6da75 100644
--- 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
index 846f03d5cc..13e7c581c4 100644
--- 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
         }