Prevent users registering twice with the same principal (#5655)

backend/canisters/local_user_index/CHANGELOG.md

@@ -9,6 +9,10 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 
 - Update `event_store` packages to latest version ([#5593](https://github.com/open-chat-labs/open-chat/pull/5593))
 
+### Fixed
+
+- Prevent users registering twice with the same principal ([#5655](https://github.com/open-chat-labs/open-chat/pull/5655))
+
 ## [[2.0.1113](https://github.com/open-chat-labs/open-chat/releases/tag/v2.0.1113-local_user_index)] - 2024-03-20
 
 ### Changed

backend/canisters/local_user_index/api/can.did

@@ -216,6 +216,7 @@ type RegisterUserResponse = variant {
         user_id : UserId;
         icp_account : AccountIdentifier;
     };
+    RegistrationInProgress;
     AlreadyRegistered;
     UserLimitReached;
     UsernameInvalid;

backend/canisters/local_user_index/api/src/updates/register_user.rs

@@ -14,6 +14,7 @@ pub struct Args {
 #[derive(CandidType, Serialize, Deserialize, Debug)]
 pub enum Response {
     Success(SuccessResult),
+    RegistrationInProgress,
     AlreadyRegistered,
     UserLimitReached,
     UsernameInvalid,

backend/canisters/local_user_index/impl/src/model/local_user_map.rs

@@ -1,16 +1,22 @@
+use candid::Principal;
 use serde::{Deserialize, Serialize};
+use std::collections::hash_map::Entry::{Occupied, Vacant};
 use std::collections::HashMap;
 use types::{BuildVersion, CyclesTopUp, TimestampMillis, UserId};
+use utils::time::MINUTE_IN_MS;
 
 #[derive(Serialize, Deserialize, Default)]
 pub struct LocalUserMap {
     users: HashMap<UserId, LocalUser>,
+    #[serde(default)]
+    registration_in_progress: HashMap<Principal, TimestampMillis>,
 }
 
 impl LocalUserMap {
-    pub fn add(&mut self, user_id: UserId, wasm_version: BuildVersion, now: TimestampMillis) {
+    pub fn add(&mut self, user_id: UserId, principal: Principal, wasm_version: BuildVersion, now: TimestampMillis) {
         let user = LocalUser::new(now, wasm_version);
         self.users.insert(user_id, user);
+        self.registration_in_progress.remove(&principal);
     }
 
     pub fn get(&self, user_id: &UserId) -> Option<&LocalUser> {
@@ -34,6 +40,24 @@ impl LocalUserMap {
         }
     }
 
+    pub fn mark_registration_in_progress(&mut self, principal: Principal, now: TimestampMillis) -> bool {
+        match self.registration_in_progress.entry(principal) {
+            Vacant(e) => {
+                e.insert(now);
+                true
+            }
+            Occupied(mut e) if *e.get() < now.saturating_sub(5 * MINUTE_IN_MS) => {
+                e.insert(now);
+                true
+            }
+            Occupied(_) => false,
+        }
+    }
+
+    pub fn mark_registration_failed(&mut self, principal: &Principal) {
+        self.registration_in_progress.remove(principal);
+    }
+
     pub fn iter(&self) -> impl Iterator<Item = (&UserId, &LocalUser)> {
         self.users.iter()
     }

backend/canisters/local_user_index/impl/src/updates/register_user.rs

@@ -66,9 +66,7 @@ async fn register_user(args: Args) -> Response {
             })
         }
         Err(error) => {
-            if let CreateAndInstallError::InstallFailed(id, ..) = error {
-                mutate_state(|state| state.data.canister_pool.push(id));
-            }
+            mutate_state(|state| rollback(&caller, &error, state));
             InternalError(format!("{error:?}"))
         }
     }
@@ -86,12 +84,16 @@ struct PrepareOk {
 
 fn prepare(args: &Args, state: &mut RuntimeState) -> Result<PrepareOk, Response> {
     let caller = state.env.caller();
-    let mut referral_code = None;
 
     if state.data.global_users.get_by_principal(&caller).is_some() {
         return Err(AlreadyRegistered);
     }
 
+    let now = state.env.now();
+    if !state.data.local_users.mark_registration_in_progress(caller, now) {
+        return Err(RegistrationInProgress);
+    }
+
     let is_from_identity_canister = validate_public_key(
         caller,
         &args.public_key,
@@ -104,8 +106,7 @@ fn prepare(args: &Args, state: &mut RuntimeState) -> Result<PrepareOk, Response>
         return Err(UserLimitReached);
     }
 
-    let now = state.env.now();
-
+    let mut referral_code = None;
     if let Some(code) = &args.referral_code {
         referral_code = match state.data.referral_codes.check(code, now) {
             Ok(r) => Some(r),
@@ -196,7 +197,7 @@ fn commit(
     state: &mut RuntimeState,
 ) {
     let now = state.env.now();
-    state.data.local_users.add(user_id, wasm_version, now);
+    state.data.local_users.add(user_id, principal, wasm_version, now);
     state.data.global_users.add(principal, user_id, false);
 
     state.push_event_to_user_index(UserIndexEvent::UserRegistered(Box::new(UserRegistered {
@@ -248,6 +249,14 @@ fn commit(
     }
 }
 
+fn rollback(principal: &Principal, error: &CreateAndInstallError, state: &mut RuntimeState) {
+    state.data.local_users.mark_registration_failed(principal);
+
+    if let CreateAndInstallError::InstallFailed(id, ..) = error {
+        state.data.canister_pool.push(*id);
+    }
+}
+
 fn welcome_messages() -> Vec<String> {
     const WELCOME_MESSAGES: &[&str] = &[
         "Welcome to OpenChat!",