From bfda60c683d9747015c927057b181602a6afb382 Mon Sep 17 00:00:00 2001
From: Evgenii Stratonikov <evgeniy@nspcc.ru>
Date: Thu, 6 Aug 2020 15:13:21 +0300
Subject: [PATCH] native: add missing witness checks to
 NEO.(Un)registerCandidate

---
 pkg/core/native/native_neo.go | 18 ++++++++++++++++--
 1 file changed, 16 insertions(+), 2 deletions(-)

diff --git a/pkg/core/native/native_neo.go b/pkg/core/native/native_neo.go
index 237479c06..b5fab4515 100644
--- a/pkg/core/native/native_neo.go
+++ b/pkg/core/native/native_neo.go
@@ -223,7 +223,14 @@ func (n *NEO) unclaimedGas(ic *interop.Context, args []stackitem.Item) stackitem
 }
 
 func (n *NEO) registerCandidate(ic *interop.Context, args []stackitem.Item) stackitem.Item {
-	err := n.RegisterCandidateInternal(ic, toPublicKey(args[0]))
+	pub := toPublicKey(args[0])
+	ok, err := runtime.CheckKeyedWitness(ic, pub)
+	if err != nil {
+		panic(err)
+	} else if !ok {
+		return stackitem.NewBool(false)
+	}
+	err = n.RegisterCandidateInternal(ic, pub)
 	return stackitem.NewBool(err == nil)
 }
 
@@ -243,7 +250,14 @@ func (n *NEO) RegisterCandidateInternal(ic *interop.Context, pub *keys.PublicKey
 }
 
 func (n *NEO) unregisterCandidate(ic *interop.Context, args []stackitem.Item) stackitem.Item {
-	err := n.UnregisterCandidateInternal(ic, toPublicKey(args[0]))
+	pub := toPublicKey(args[0])
+	ok, err := runtime.CheckKeyedWitness(ic, pub)
+	if err != nil {
+		panic(err)
+	} else if !ok {
+		return stackitem.NewBool(false)
+	}
+	err = n.UnregisterCandidateInternal(ic, pub)
 	return stackitem.NewBool(err == nil)
 }