From 971f9e4185b97dd6a0dcd83b965e027914240a46 Mon Sep 17 00:00:00 2001 From: Hamish Peebles Date: Tue, 9 Apr 2024 10:43:05 +0100 Subject: [PATCH] Prevent users registering twice with the same principal (#5655) --- .../canisters/local_user_index/CHANGELOG.md | 4 +++ .../canisters/local_user_index/api/can.did | 1 + .../api/src/updates/register_user.rs | 1 + .../impl/src/model/local_user_map.rs | 26 ++++++++++++++++++- .../impl/src/updates/register_user.rs | 23 +++++++++++----- 5 files changed, 47 insertions(+), 8 deletions(-) diff --git a/backend/canisters/local_user_index/CHANGELOG.md b/backend/canisters/local_user_index/CHANGELOG.md index 5677070ebc..835937edaa 100644 --- a/backend/canisters/local_user_index/CHANGELOG.md +++ b/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 diff --git a/backend/canisters/local_user_index/api/can.did b/backend/canisters/local_user_index/api/can.did index 366af9c9a0..f08dfe8f94 100644 --- a/backend/canisters/local_user_index/api/can.did +++ b/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; diff --git a/backend/canisters/local_user_index/api/src/updates/register_user.rs b/backend/canisters/local_user_index/api/src/updates/register_user.rs index 15d99399ee..aaa350436d 100644 --- a/backend/canisters/local_user_index/api/src/updates/register_user.rs +++ b/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, diff --git a/backend/canisters/local_user_index/impl/src/model/local_user_map.rs b/backend/canisters/local_user_index/impl/src/model/local_user_map.rs index a4f9f49edc..12fb699d8a 100644 --- a/backend/canisters/local_user_index/impl/src/model/local_user_map.rs +++ b/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, + #[serde(default)] + registration_in_progress: HashMap, } 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 { self.users.iter() } diff --git a/backend/canisters/local_user_index/impl/src/updates/register_user.rs b/backend/canisters/local_user_index/impl/src/updates/register_user.rs index e1910cfba7..3efa0d72d7 100644 --- a/backend/canisters/local_user_index/impl/src/updates/register_user.rs +++ b/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 { 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 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 { const WELCOME_MESSAGES: &[&str] = &[ "Welcome to OpenChat!",