From 2d579597aa8c5d3255b4ad28ae3cd81e084836a0 Mon Sep 17 00:00:00 2001
From: Erik van den Brink <git@erikvandenbrink.com>
Date: Mon, 27 Jun 2022 13:24:50 +0200
Subject: [PATCH] fix getcandidates dos
 https://github.com/neo-project/neo/pull/2686 (#165)

---
 neo3/contracts/native/fungible.py | 32 ++++++++++++++++++++++++++++---
 1 file changed, 29 insertions(+), 3 deletions(-)

diff --git a/neo3/contracts/native/fungible.py b/neo3/contracts/native/fungible.py
index bb1e58d8..e6a6489d 100644
--- a/neo3/contracts/native/fungible.py
+++ b/neo3/contracts/native/fungible.py
@@ -5,6 +5,7 @@
 from neo3 import storage, contracts, vm, settings
 from neo3.core import types, msgrouter, cryptography, serialization, to_script_hash, Size as s, IInteroperable
 from typing import Tuple, List, Dict, Sequence, cast, Optional
+from neo3.contracts.interop import IIterator, StorageIterator
 
 
 class FungibleTokenStorageState(IInteroperable, serialization.ISerializable):
@@ -645,9 +646,28 @@ def vote(self,
     @register("getCandidates", contracts.CallFlags.READ_STATES, cpu_price=1 << 22)
     def get_candidates(self, engine: contracts.ApplicationEngine) -> List[Tuple[cryptography.ECPoint, vm.BigInteger]]:
         """
-        Fetch all registered candidates, convert them to a StackItem and push them onto the evaluation stack.
+        Fetch the first 256 registered candidates, convert them to a StackItem and push them onto the evaluation stack.
+        First meaning the first found, not first registrered.
         """
-        return self._get_candidates(engine.snapshot)
+        return self._get_candidates(engine.snapshot, max_count=256)
+
+    @register("getAllCandidates", contracts.CallFlags.READ_STATES, cpu_price=1 << 22)
+    def get_all_candidates(self, snapshot: storage.Snapshot) -> IIterator:
+        options = (contracts.FindOptions.REMOVE_PREFIX
+                   | contracts.FindOptions.DESERIALIZE_VALUES
+                   | contracts.FindOptions.PICK_FIELD1)
+        it = StorageIterator(snapshot.storages.find(self.key_candidate.to_array()), 1, options)
+        return it
+
+    def get_candidate_vote(self, snapshot: storage.Snapshot, public_key: cryptography.ECPoint) -> vm.BigInteger:
+        si_candidate = snapshot.storages.try_get(self.key_candidate + public_key, read_only=False)
+        if si_candidate is None:
+            return vm.BigInteger(-1)
+        candidate_state = si_candidate.get(_CandidateState)
+        if candidate_state.registered:
+            return candidate_state.votes
+        else:
+            return vm.BigInteger(-1)
 
     @register("getNextBlockValidators", contracts.CallFlags.READ_STATES, cpu_price=1 << 16)
     def get_next_block_validators(self, snapshot: storage.Snapshot) -> List[cryptography.ECPoint]:
@@ -910,8 +930,10 @@ def _compute_committee_members(self, snapshot: storage.Snapshot) -> Dict[cryptog
         return results
 
     def _get_candidates(self,
-                        snapshot: storage.Snapshot) -> \
+                        snapshot: storage.Snapshot,
+                        max_count=-1) -> \
             List[Tuple[cryptography.ECPoint, vm.BigInteger]]:
+        ctr = 0
         self._candidates = []
         policy = contracts.PolicyContract()
         for k, v in snapshot.storages.find(self.key_candidate.to_array()):
@@ -922,6 +944,10 @@ def _get_candidates(self,
                 # take of the CANDIDATE prefix
                 point = cryptography.ECPoint.deserialize_from_bytes(k.key[1:])
                 self._candidates.append((point, candidate.votes))
+                if max_count > 0:
+                    ctr += 1
+                    if ctr == max_count:
+                        break
         self._candidates_dirty = False
         return self._candidates