From e32b298072736c3e93bfb908ac832819219a8962 Mon Sep 17 00:00:00 2001
From: Hamish Peebles <hamishpeebles@gmail.com>
Date: Thu, 21 Nov 2024 22:34:44 +0000
Subject: [PATCH] Enforce a limit of 100 registrations

---
 canister/can.did                             |  1 +
 canister/src/jobs/check_for_new_nns_votes.rs |  3 +-
 canister/src/jobs/process_votes.rs           |  3 +-
 canister/src/lib.rs                          |  1 +
 canister/src/queries/list_neuron_pairs.rs    |  2 +-
 canister/src/state.rs                        | 14 ++-------
 canister/src/updates/register_neuron_pair.rs | 31 ++++++++++++++++++--
 7 files changed, 39 insertions(+), 16 deletions(-)

diff --git a/canister/can.did b/canister/can.did
index eb08d62..e2faf1c 100644
--- a/canister/can.did
+++ b/canister/can.did
@@ -21,6 +21,7 @@ type RegisterNeuronPairError = variant {
   ErrorCallingGovernanceCanister : record { int32; text };
   NotPermittedToVote;
   AlreadyRegistered;
+  RegistrationLimitExceeded : nat32;
   GovernanceError : record { int32; text };
 };
 type Result = variant { Ok : nat64; Err : RegisterNeuronPairError };
diff --git a/canister/src/jobs/check_for_new_nns_votes.rs b/canister/src/jobs/check_for_new_nns_votes.rs
index 0566f9e..10d6a81 100644
--- a/canister/src/jobs/check_for_new_nns_votes.rs
+++ b/canister/src/jobs/check_for_new_nns_votes.rs
@@ -16,7 +16,8 @@ async fn run() {
     log("Checking for new NNS votes");
 
     let futures: Vec<_> = state::mutate(|s| {
-        s.iter_neuron_pairs()
+        s.neuron_pairs()
+            .values()
             .map(|p| run_single(p.id(), s.nns_governance_canister_id(), p.nns_neuron_id()))
             .collect()
     });
diff --git a/canister/src/jobs/process_votes.rs b/canister/src/jobs/process_votes.rs
index 3e2dfb6..1578964 100644
--- a/canister/src/jobs/process_votes.rs
+++ b/canister/src/jobs/process_votes.rs
@@ -103,7 +103,8 @@ async fn process_vote(vote: VoteToProcess) {
         }
         VoteToProcess::PendingWtnVote(pair_id, wtn_vote) => {
             let Some((canister_id, neuron_id)) = state::read(|s| {
-                s.neuron_pair(pair_id)
+                s.neuron_pairs()
+                    .get(&pair_id)
                     .map(|p| (s.wtn_governance_canister_id(), p.wtn_neuron_id()))
             }) else {
                 return;
diff --git a/canister/src/lib.rs b/canister/src/lib.rs
index 321764b..40f1803 100644
--- a/canister/src/lib.rs
+++ b/canister/src/lib.rs
@@ -73,6 +73,7 @@ struct RegisterNeuronPairArgs {
 enum RegisterNeuronPairError {
     AlreadyRegistered,
     NotPermittedToVote,
+    RegistrationLimitExceeded(u32),
     GovernanceError(i32, String),
     ErrorCallingGovernanceCanister(i32, String),
 }
diff --git a/canister/src/queries/list_neuron_pairs.rs b/canister/src/queries/list_neuron_pairs.rs
index 44e9890..1d4d1b8 100644
--- a/canister/src/queries/list_neuron_pairs.rs
+++ b/canister/src/queries/list_neuron_pairs.rs
@@ -3,5 +3,5 @@ use ic_cdk::query;
 
 #[query]
 fn list_neuron_pairs() -> Vec<NeuronPairPublic> {
-    state::read(|s| s.list_neuron_pairs())
+    state::read(|s| s.neuron_pairs().values().map(|pair| pair.into()).collect())
 }
diff --git a/canister/src/state.rs b/canister/src/state.rs
index 847b3a2..2855a14 100644
--- a/canister/src/state.rs
+++ b/canister/src/state.rs
@@ -1,6 +1,6 @@
 use crate::logs::log;
 use crate::neuron_pair::NeuronPair;
-use crate::{InitArgs, NeuronPairPublic, NnsVote, VoteToProcess, WtnVote};
+use crate::{InitArgs, NnsVote, VoteToProcess, WtnVote};
 use ic_principal::Principal;
 use serde::{Deserialize, Serialize};
 use std::cell::RefCell;
@@ -84,10 +84,6 @@ impl State {
         self.wtn_protocol_canister_id
     }
 
-    pub fn list_neuron_pairs(&self) -> Vec<NeuronPairPublic> {
-        self.neuron_pairs.values().map(|p| p.into()).collect()
-    }
-
     pub fn register_neuron_pair(
         &mut self,
         caller: Principal,
@@ -120,12 +116,8 @@ impl State {
         }
     }
 
-    pub fn neuron_pair(&self, pair_id: u64) -> Option<&NeuronPair> {
-        self.neuron_pairs.get(&pair_id)
-    }
-
-    pub fn iter_neuron_pairs(&self) -> impl Iterator<Item = &NeuronPair> {
-        self.neuron_pairs.values()
+    pub fn neuron_pairs(&self) -> &BTreeMap<u64, NeuronPair> {
+        &self.neuron_pairs
     }
 
     pub fn record_nns_vote(&mut self, pair_id: u64, vote: NnsVote) {
diff --git a/canister/src/updates/register_neuron_pair.rs b/canister/src/updates/register_neuron_pair.rs
index b99b426..72249b6 100644
--- a/canister/src/updates/register_neuron_pair.rs
+++ b/canister/src/updates/register_neuron_pair.rs
@@ -6,13 +6,19 @@ use ic_principal::Principal;
 use serde::{Deserialize, Serialize};
 
 const REGISTER_VOTE_PERMISSION: i32 = 4;
+const REGISTRATIONS_LIMIT: u32 = 100;
 
 #[update]
 async fn register_neuron_pair(
     args: RegisterNeuronPairArgs,
 ) -> Result<u64, RegisterNeuronPairError> {
-    let caller = ic_cdk::caller();
-    let wtn_governance_canister = state::read(|s| s.wtn_governance_canister_id());
+    let PrepareSuccess {
+        caller,
+        wtn_governance_canister,
+    } = match prepare() {
+        Ok(success) => success,
+        Err(error) => return Err(error),
+    };
 
     match call_get_neuron(wtn_governance_canister, args.wtn_neuron_id).await {
         Ok(response) => match response.0.result.unwrap() {
@@ -49,6 +55,27 @@ async fn register_neuron_pair(
     }
 }
 
+struct PrepareSuccess {
+    caller: Principal,
+    wtn_governance_canister: Principal,
+}
+
+fn prepare() -> Result<PrepareSuccess, RegisterNeuronPairError> {
+    state::read(|s| {
+        if s.neuron_pairs().len() >= REGISTRATIONS_LIMIT as usize {
+            Err(RegisterNeuronPairError::RegistrationLimitExceeded(
+                REGISTRATIONS_LIMIT,
+            ))
+        } else {
+            Ok(s.wtn_governance_canister_id())
+        }
+    })
+    .map(|wtn_governance_canister| PrepareSuccess {
+        caller: ic_cdk::caller(),
+        wtn_governance_canister,
+    })
+}
+
 async fn call_get_neuron(
     governance_canister: Principal,
     neuron_id: [u8; 32],