From c7a04e455929207d22c89f70f05e08bf70860844 Mon Sep 17 00:00:00 2001 From: Paul Makles Date: Sun, 11 Jun 2023 09:53:41 +0100 Subject: [PATCH] fix: conflict resolution in migration fix: write new username and display name chore: restrict discriminator search space --- .../admin_migrations/ops/mongodb/scripts.rs | 52 +++++++++++++++---- .../core/database/src/models/users/model.rs | 15 ++++++ crates/quark/src/impl/generic/users/user.rs | 17 +++++- 3 files changed, 71 insertions(+), 13 deletions(-) diff --git a/crates/core/database/src/models/admin_migrations/ops/mongodb/scripts.rs b/crates/core/database/src/models/admin_migrations/ops/mongodb/scripts.rs index 42e9eaef..6b822380 100644 --- a/crates/core/database/src/models/admin_migrations/ops/mongodb/scripts.rs +++ b/crates/core/database/src/models/admin_migrations/ops/mongodb/scripts.rs @@ -1,11 +1,11 @@ -use std::{ops::BitXor, time::Duration}; +use std::{collections::HashSet, ops::BitXor, time::Duration}; use crate::{ mongodb::{ bson::{doc, from_bson, from_document, to_document, Bson, DateTime, Document}, options::FindOptions, }, - MongoDb, + MongoDb, DISCRIMINATOR_SEARCH_SPACE, }; use futures::StreamExt; use rand::seq::SliceRandom; @@ -802,21 +802,38 @@ pub async fn run_migrations(db: &MongoDb, revision: i32) -> i32 { .collect() .await; - let search_space: Vec = (0..9999).map(|v| format!("{:0>4}", v)).collect(); + let search_space: Vec = DISCRIMINATOR_SEARCH_SPACE.iter().cloned().collect(); + let mut claimed: HashSet = HashSet::new(); for i in 0..users.len() { let info = &users[i]; - let discriminator = { + let mut discriminator = { let mut rng = rand::thread_rng(); search_space.choose(&mut rng).unwrap() }; if re_username.is_match(&info.username) { + while claimed.contains(&format!("{}#{}", info.username, discriminator)) { + let new_discriminator = { + let mut rng = rand::thread_rng(); + search_space.choose(&mut rng).unwrap() + }; + + info!( + "Re-rolled {} to {new_discriminator} from {discriminator}", + info.username + ); + + discriminator = new_discriminator; + } + + claimed.insert(format!("{}#{}", info.username, discriminator)); + info!( - "Migrating user \"{}\" to #{} ({i}/{}) - compliant", + "({i}/{}) Migrating user \"{}\" to #{} - compliant", + users.len(), info.username, - discriminator, - users.len() + discriminator ); db.col::("users") @@ -844,11 +861,23 @@ pub async fn run_migrations(db: &MongoDb, revision: i32) -> i32 { sanitised += "_"; } + while claimed.contains(&format!("{}#{}", sanitised, discriminator)) { + let new_discriminator = { + let mut rng = rand::thread_rng(); + search_space.choose(&mut rng).unwrap() + }; + + info!("Re-rolled {sanitised} to {new_discriminator} from {discriminator}"); + discriminator = new_discriminator; + } + + claimed.insert(format!("{}#{}", sanitised, discriminator)); + info!( - "Migrating user \"{}\" to #{} ({i}/{}) - sanitised: \"{}\"", + "({i}/{}) Migrating user \"{}\" to #{} - sanitised: \"{}\"", + users.len(), info.username, discriminator, - users.len(), sanitised ); @@ -859,8 +888,9 @@ pub async fn run_migrations(db: &MongoDb, revision: i32) -> i32 { }, doc! { "$set": { - "username": &info.username, - "discriminator": discriminator + "username": sanitised, + "discriminator": discriminator, + "display_name": &info.username } }, None, diff --git a/crates/core/database/src/models/users/model.rs b/crates/core/database/src/models/users/model.rs index 589589eb..0ea2d08b 100644 --- a/crates/core/database/src/models/users/model.rs +++ b/crates/core/database/src/models/users/model.rs @@ -1,5 +1,8 @@ +use std::collections::HashSet; + use crate::{Database, File}; +use once_cell::sync::Lazy; use revolt_result::{Error, ErrorType, Result}; auto_derived_partial!( @@ -199,3 +202,15 @@ impl User { .await } } + +pub static DISCRIMINATOR_SEARCH_SPACE: Lazy> = Lazy::new(|| { + let mut set = (2..9999) + .map(|v| format!("{:0>4}", v)) + .collect::>(); + + for discrim in [1111, 2222, 3333, 4444, 5555, 6666, 7777, 8888, 9999] { + set.remove(&discrim.to_string()); + } + + set.into_iter().collect() +}); diff --git a/crates/quark/src/impl/generic/users/user.rs b/crates/quark/src/impl/generic/users/user.rs index f73a7adf..2d4c9606 100644 --- a/crates/quark/src/impl/generic/users/user.rs +++ b/crates/quark/src/impl/generic/users/user.rs @@ -8,6 +8,7 @@ use crate::{perms, Database, Error, Result}; use futures::try_join; use impl_ops::impl_op_ex_commutative; +use once_cell::sync::Lazy; use rand::seq::SliceRandom; use revolt_presence::filter_online; use std::collections::HashSet; @@ -202,7 +203,7 @@ impl User { username: &str, preferred: Option, ) -> Result { - let search_space: HashSet = (0..9999).map(|v| format!("{:0>4}", v)).collect(); + let search_space: &HashSet = &DISCRIMINATOR_SEARCH_SPACE_QUARK; let used_discriminators: HashSet = db .fetch_discriminators_in_use(username) .await? @@ -218,7 +219,7 @@ impl User { if let Some(preferred) = preferred { if available_discriminators.contains(&&preferred) { - return Ok(preferred.to_string()); + return Ok(preferred); } } @@ -450,3 +451,15 @@ impl User { } } } + +pub static DISCRIMINATOR_SEARCH_SPACE_QUARK: Lazy> = Lazy::new(|| { + let mut set = (2..9999) + .map(|v| format!("{:0>4}", v)) + .collect::>(); + + for discrim in [1111, 2222, 3333, 4444, 5555, 6666, 7777, 8888, 9999] { + set.remove(&discrim.to_string()); + } + + set.into_iter().collect() +});