From 5bbe30edbc2f6a8f67d550eef56e3823e5bb7f38 Mon Sep 17 00:00:00 2001
From: Paul Makles <me@insrt.uk>
Date: Sun, 11 Jun 2023 09:16:24 +0100
Subject: [PATCH] feat(core/database): migrate to discriminators

---
 Cargo.lock                                    |   9 ++
 crates/core/database/Cargo.toml               |   5 +
 .../admin_migrations/ops/mongodb/scripts.rs   | 108 ++++++++++++++++--
 3 files changed, 114 insertions(+), 8 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 40aec27e..7d10d349 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2867,7 +2867,9 @@ dependencies = [
  "mongodb",
  "nanoid",
  "once_cell",
+ "rand 0.8.5",
  "redis-kiss",
+ "regex",
  "revolt-models",
  "revolt-permissions",
  "revolt-presence",
@@ -2878,6 +2880,7 @@ dependencies = [
  "serde",
  "serde_json",
  "ulid 1.0.0",
+ "unicode-segmentation",
 ]
 
 [[package]]
@@ -4327,6 +4330,12 @@ dependencies = [
  "tinyvec",
 ]
 
+[[package]]
+name = "unicode-segmentation"
+version = "1.10.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1dd624098567895118886609431a7c3b8f516e41d30e0643f03d94592a147e36"
+
 [[package]]
 name = "unicode-xid"
 version = "0.0.4"
diff --git a/crates/core/database/Cargo.toml b/crates/core/database/Cargo.toml
index 6dc17941..08533b18 100644
--- a/crates/core/database/Cargo.toml
+++ b/crates/core/database/Cargo.toml
@@ -29,6 +29,7 @@ revolt-permissions = { version = "0.6.0-rc.2", path = "../permissions", features
 
 # Utility
 log = "0.4"
+rand = "0.8.5"
 ulid = "1.0.0"
 nanoid = "0.4.0"
 once_cell = "1.17"
@@ -46,6 +47,10 @@ redis-kiss = { version = "0.1.4" }
 bson = { optional = true, version = "2.1.0" }
 mongodb = { optional = true, version = "2.1.0", default-features = false }
 
+# Database Migration
+unicode-segmentation = "1.10.1"
+regex = "1"
+
 # Async Language Features
 futures = "0.3.19"
 async-trait = "0.1.51"
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 9fea8f9f..42e9eaef 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
@@ -8,7 +8,9 @@ use crate::{
     MongoDb,
 };
 use futures::StreamExt;
+use rand::seq::SliceRandom;
 use serde::{Deserialize, Serialize};
+use unicode_segmentation::UnicodeSegmentation;
 
 #[derive(Serialize, Deserialize)]
 struct MigrationInfo {
@@ -16,7 +18,7 @@ struct MigrationInfo {
     revision: i32,
 }
 
-pub const LATEST_REVISION: i32 = 24;
+pub const LATEST_REVISION: i32 = 25;
 
 pub async fn migrate_database(db: &MongoDb) {
     let migrations = db.col::<Document>("migrations");
@@ -769,24 +771,114 @@ pub async fn run_migrations(db: &MongoDb, revision: i32) -> i32 {
     }
 
     if revision <= 23 {
-        info!("Running migration [revision 23 / 09-06-2023]: Add collection `channel_webhooks` if not exists, update users index.");
-
-        db.db()
-            .create_collection("channel_webhooks", None)
-            .await
-            .ok();
+        info!("Running migration [revision 23 / 10-06-2023]: Generate discriminators for users.");
 
         db.db()
             .run_command(
                 doc! {
                     "dropIndexes": "users",
-                    "indexes": "username"
+                    "index": "username"
                 },
                 None,
             )
             .await
             .expect("Failed to drop existing username index.");
 
+        #[derive(Serialize, Deserialize)]
+        struct UserInformation {
+            #[serde(rename = "_id")]
+            id: String,
+            username: String,
+        }
+
+        let re_username = regex::Regex::new(r"^(\p{L}|[\d_.-])+$").unwrap();
+
+        let users: Vec<UserInformation> = db
+            .col::<UserInformation>("users")
+            .find(doc! {}, None)
+            .await
+            .unwrap()
+            .map(|doc| doc.expect("id and username"))
+            .collect()
+            .await;
+
+        let search_space: Vec<String> = (0..9999).map(|v| format!("{:0>4}", v)).collect();
+
+        for i in 0..users.len() {
+            let info = &users[i];
+            let discriminator = {
+                let mut rng = rand::thread_rng();
+                search_space.choose(&mut rng).unwrap()
+            };
+
+            if re_username.is_match(&info.username) {
+                info!(
+                    "Migrating user \"{}\" to #{} ({i}/{}) - compliant",
+                    info.username,
+                    discriminator,
+                    users.len()
+                );
+
+                db.col::<UserInformation>("users")
+                    .update_one(
+                        doc! {
+                            "_id": &info.id
+                        },
+                        doc! {
+                            "$set": {
+                                "discriminator": discriminator
+                            }
+                        },
+                        None,
+                    )
+                    .await
+                    .unwrap();
+            } else {
+                let mut sanitised = info
+                    .username
+                    .graphemes(true)
+                    .filter(|s| re_username.is_match(s))
+                    .collect::<String>();
+
+                while sanitised.len() < 2 {
+                    sanitised += "_";
+                }
+
+                info!(
+                    "Migrating user \"{}\" to #{} ({i}/{}) - sanitised: \"{}\"",
+                    info.username,
+                    discriminator,
+                    users.len(),
+                    sanitised
+                );
+
+                db.col::<UserInformation>("users")
+                    .update_one(
+                        doc! {
+                            "_id": &info.id
+                        },
+                        doc! {
+                            "$set": {
+                                "username": &info.username,
+                                "discriminator": discriminator
+                            }
+                        },
+                        None,
+                    )
+                    .await
+                    .unwrap();
+            }
+        }
+    }
+
+    if revision <= 24 {
+        info!("Running migration [revision 24 / 09-06-2023]: Add collection `channel_webhooks` if not exists, update users index.");
+
+        db.db()
+            .create_collection("channel_webhooks", None)
+            .await
+            .ok();
+
         db.db()
             .run_command(
                 doc! {