Redo migration table

src/router/gree.rs

@@ -15,7 +15,6 @@ use rsa::pkcs8::DecodePublicKey;
 use crate::router::global;
 use crate::router::userdata;
 use crate::encryption;
-use crate::router::user::{code_to_uid, uid_to_code};
 use crate::sql::SQLite;
 
 lazy_static! {
@@ -231,11 +230,9 @@ pub fn migration_verify(req: HttpRequest, body: String) -> HttpResponse {
     let body = json::parse(&body).unwrap();
     let password = decrypt_transfer_password(&body["migration_password"].to_string());
     
-    let uid = code_to_uid(body["migration_code"].to_string()).parse::<i64>().unwrap_or(0);
+    let user = userdata::user::migration::get_acc_transfer(&body["migration_code"].to_string(), &password);
     
-    let user = userdata::user::migration::get_acc_transfer(uid, &body["migration_code"].to_string(), &password);
+    let resp = if !user["success"].as_bool().unwrap() || user["user_id"] == 0 {
-    
-    let resp = if !user["success"].as_bool().unwrap() || uid == 0 {
         object!{
             result: "ERR",
             messsage: "User Not Found"
@@ -245,7 +242,7 @@ pub fn migration_verify(req: HttpRequest, body: String) -> HttpResponse {
         object!{
             result: "OK",
             src_uuid: user["login_token"].clone(),
-            src_x_uid: uid.to_string(),
+            src_x_uid: user["user_id"].to_string(),
             migration_token: user["login_token"].clone(),
             balance_charge_gem: data_user["gem"]["charge"].to_string(),
             balance_free_gem: data_user["gem"]["free"].to_string(),
@@ -311,12 +308,12 @@ pub fn migration_code(req: HttpRequest) -> HttpResponse {
             uid = uid_str.to_string();
         }
     }
-    
+
     let user = userdata::get_acc(&uid);
-    
+
     let resp = object!{
         result: "OK",
-        migration_code: uid_to_code(user["user"]["id"].to_string())
+        migration_code: userdata::user::migration::get_acc_token(user["user"]["id"].as_i64().unwrap())
     };
     
     send(req, resp)
@@ -334,12 +331,12 @@ pub fn migration_password_register(req: HttpRequest, body: String) -> HttpRespon
             uid = uid_str.to_string();
         }
     }
-    
+
     let user = userdata::get_acc(&uid);
-    let code = uid_to_code(user["user"]["id"].to_string());
+
     let pass = decrypt_transfer_password(&body["migration_password"].to_string());
     
-    userdata::user::migration::save_acc_transfer(&code, &pass);
+    userdata::user::migration::save_acc_transfer(user["user"]["id"].as_i64().unwrap(), &pass);
     
     let resp = object!{
         result: "OK"

src/router/user.rs

@@ -169,41 +169,10 @@ pub fn announcement(req: HttpRequest) -> Option<JsonValue> {
     })
 }
 
-pub fn uid_to_code(uid: String) -> String {
-    //just replace uid with numbers because im too lazy to have a real database and this is close enough anyways
-    uid
-        .replace('1', "A")
-        .replace('2', "G")
-        .replace('3', "W")
-        .replace('4', "Q")
-        .replace('5', "Y")
-        .replace('6', "6")
-        .replace('7', "I")
-        .replace('8', "P")
-        .replace('9', "U")
-        .replace('0', "M")
-        + "7"
-}
-pub fn code_to_uid(code: String) -> String {
-    //just replace uid with numbers because im too lazy to have a real database and this is close enough anyways
-    code
-        .replace('7', "")
-        .replace('A', "1")
-        .replace('G', "2")
-        .replace('W', "3")
-        .replace('Q', "4")
-        .replace('Y', "5")
-        .replace('6', "6")
-        .replace('I', "7")
-        .replace('P', "8")
-        .replace('U', "9")
-        .replace('M', "0")
-}
-
 pub fn get_migration_code(_req: HttpRequest, body: String) -> Option<JsonValue> {
     let body = json::parse(&encryption::decrypt_packet(&body).unwrap()).unwrap();
     
-    let code = uid_to_code(body["user_id"].to_string());
+    let code = userdata::user::migration::get_acc_token(body["user_id"].as_i64()?);
     
     Some(object!{
         "migrationCode": code
@@ -215,9 +184,8 @@ pub fn register_password(req: HttpRequest, body: String) -> Option<JsonValue> {
     let body = json::parse(&encryption::decrypt_packet(&body).unwrap()).unwrap();
     
     let user = userdata::get_acc(&key);
-    let code = uid_to_code(user["user"]["id"].to_string());
     
-    userdata::user::migration::save_acc_transfer(&code, &body["pass"].to_string());
+    userdata::user::migration::save_acc_transfer(user["user"]["id"].as_i64().unwrap(), &body["pass"].to_string());
     
     Some(array![])
 }
@@ -225,18 +193,16 @@ pub fn register_password(req: HttpRequest, body: String) -> Option<JsonValue> {
 pub fn verify_migration_code(_req: HttpRequest, body: String) -> Option<JsonValue> {
     let body = json::parse(&encryption::decrypt_packet(&body).unwrap()).unwrap();
     
-    let uid = code_to_uid(body["migrationCode"].to_string()).parse::<i64>().unwrap_or(0);
+    let user = userdata::user::migration::get_acc_transfer(&body["migrationCode"].to_string(), &body["pass"].to_string());
     
-    let user = userdata::user::migration::get_acc_transfer(uid, &body["migrationCode"].to_string(), &body["pass"].to_string());
+    if !user["success"].as_bool().unwrap() || user["user_id"] == 0 {
-    
-    if !user["success"].as_bool().unwrap() || uid == 0 {
         return None;
     }
     
     let data_user = userdata::get_acc(&user["login_token"].to_string());
     
     Some(object!{
-        "user_id": uid,
+        "user_id": user["user_id"].clone(),
         "uuid": user["login_token"].to_string(),
         "charge": data_user["gem"]["charge"].clone(),
         "free": data_user["gem"]["free"].clone()
@@ -245,11 +211,9 @@ pub fn verify_migration_code(_req: HttpRequest, body: String) -> Option<JsonValu
 pub fn request_migration_code(_req: HttpRequest, body: String) -> Option<JsonValue> {
     let body = json::parse(&encryption::decrypt_packet(&body).unwrap()).unwrap();
     
-    let uid = code_to_uid(body["migrationCode"].to_string()).parse::<i64>().unwrap_or(0);
+    let user = userdata::user::migration::get_acc_transfer(&body["migrationCode"].to_string(), &body["pass"].to_string());
     
-    let user = userdata::user::migration::get_acc_transfer(uid, &body["migrationCode"].to_string(), &body["pass"].to_string());
+    if !user["success"].as_bool().unwrap() || user["user_id"] == 0 {
-    
-    if !user["success"].as_bool().unwrap() || uid == 0 {
         return None;
     }
     

src/router/userdata/mod.rs

@@ -22,15 +22,12 @@ fn get_userdata_database() -> &'static SQLite {
 }
 
 fn setup_tables(conn: &rusqlite::Connection) {
+    user::migration::setup_sql(conn).unwrap();
     conn.execute_batch("
 CREATE TABLE IF NOT EXISTS tokens (
     user_id           BIGINT NOT NULL PRIMARY KEY,
     token             TEXT NOT NULL
 );
-CREATE TABLE IF NOT EXISTS migration (
-    token             TEXT NOT NULL PRIMARY KEY,
-    password          TEXT NOT NULL
-);
 CREATE TABLE IF NOT EXISTS userdata (
     user_id                  BIGINT NOT NULL PRIMARY KEY,
     userdata                 TEXT NOT NULL,
@@ -422,7 +419,7 @@ fn create_webui_token() -> String {
 }
 
 pub fn webui_login(uid: i64, password: &str) -> Result<String, String> {
-    let pass = DATABASE.lock_and_select("SELECT password FROM migration WHERE token=?1", params!(crate::router::user::uid_to_code(uid.to_string()))).unwrap_or_default();
+    let pass = DATABASE.lock_and_select("SELECT password FROM migration WHERE user_id=?1", params!(uid)).unwrap_or_default();
     if !user::migration::verify_password(password, &pass) {
         if acc_exists(uid) && pass.is_empty() {
             return Err(String::from("Migration token not set. Set token in game settings."));
@@ -448,13 +445,12 @@ pub fn webui_import_user(user: JsonValue) -> Result<JsonValue, String> {
     let token = global::create_token();
     
     DATABASE.lock_and_exec("INSERT INTO tokens (user_id, token) VALUES (?1, ?2)", params!(uid, token));
-    let mig = crate::router::user::uid_to_code(uid.to_string());
     
-    user::migration::save_acc_transfer(&mig, &user["password"].to_string());
+    let token = user::migration::save_acc_transfer(uid, &user["password"].to_string());
     
     Ok(object!{
         uid: uid,
-        migration_token: mig
+        migration_token: token
     })
 }
 
@@ -583,7 +579,7 @@ pub fn purge_accounts() -> usize {
         DATABASE.lock_and_exec("DELETE FROM server_data WHERE user_id=?1", params!(user_id));
         DATABASE.lock_and_exec("DELETE FROM webui WHERE user_id=?1", params!(user_id));
         DATABASE.lock_and_exec("DELETE FROM tokens WHERE user_id=?1", params!(user_id));
-        DATABASE.lock_and_exec("DELETE FROM migration WHERE token=?1", params!(crate::router::user::uid_to_code(user_id.to_string())));
+        DATABASE.lock_and_exec("DELETE FROM migration WHERE user_id=?1", params!(user_id));
     }
     DATABASE.lock_and_exec("VACUUM", params!());
     crate::router::gree::vacuum_database();

src/router/userdata/user/migration.rs

@@ -5,26 +5,54 @@ use rand::Rng;
 use sha2::{Digest, Sha256};
 use base64::{Engine as _, engine::general_purpose};
 
-pub fn get_acc_transfer(uid: i64, token: &str, password: &str) -> JsonValue {
+fn generate_token() -> String {
+    let charset = "ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789";
+    let mut rng = rand::rng();
+    let random_string: String = (0..16)
+        .map(|_| {
+            let idx = rng.random_range(0..charset.len());
+            charset.chars().nth(idx).unwrap()
+        })
+        .collect();
+    random_string
+}
+
+pub fn get_acc_transfer(token: &str, password: &str) -> JsonValue {
     let database = userdata::get_userdata_database();
     let data = database.lock_and_select("SELECT password FROM migration WHERE token=?1", params!(token));
     if data.is_err() {
         return object!{success: false};
     }
     if verify_password(password, &data.unwrap()) {
+        let uid: i64 = database.lock_and_select_type("SELECT user_id FROM migration WHERE token=?1", params!(token)).unwrap();
         let login_token = userdata::get_login_token(uid);
         if login_token == String::new() {
             return object!{success: false};
         }
-        return object!{success: true, login_token: login_token};
+        return object!{success: true, login_token: login_token, user_id: uid};
     }
     object!{success: false}
 }
 
-pub fn save_acc_transfer(token: &str, password: &str) {
+pub fn save_acc_transfer(uid: i64, password: &str) -> String {
     let database = userdata::get_userdata_database();
-    database.lock_and_exec("DELETE FROM migration WHERE token=?1", params!(token));
+    let token = if let Ok(value) = database.lock_and_select("SELECT token FROM migration WHERE user_id=?1", params!(uid)) {
-    database.lock_and_exec("INSERT INTO migration (token, password) VALUES (?1, ?2)", params!(token, hash_password(password)));
+        value
+    } else {
+        generate_token()
+    };
+    database.lock_and_exec("DELETE FROM migration WHERE user_id=?1", params!(uid));
+    database.lock_and_exec("INSERT INTO migration (user_id, token, password) VALUES (?1, ?2, ?3)", params!(uid, &token, hash_password(password)));
+    token
+}
+
+pub fn get_acc_token(uid: i64) -> String {
+    let database = userdata::get_userdata_database();
+    if let Ok(value) = database.lock_and_select("SELECT token FROM migration WHERE user_id=?1", params!(uid)) {
+        value
+    } else {
+        save_acc_transfer(uid, "")
+    }
 }
 
 fn generate_salt() -> Vec<u8> {
@@ -35,6 +63,7 @@ fn generate_salt() -> Vec<u8> {
 }
 
 fn hash_password(password: &str) -> String {
+    if password.is_empty() { return String::new(); };
     let salt = &generate_salt();
     let mut hasher = Sha256::new();
     hasher.update(password.as_bytes());
@@ -67,3 +96,62 @@ pub fn verify_password(password: &str, salted_hash: &str) -> bool {
 
     input_hash.as_slice() == hashed_password
 }
+
+pub fn setup_sql(conn: &rusqlite::Connection) -> Result<(), rusqlite::Error> {
+    conn.execute("
+    CREATE TABLE IF NOT EXISTS migration (
+        user_id           BIGINT NOT NULL,
+        token             TEXT NOT NULL,
+        password          TEXT NOT NULL,
+        PRIMARY KEY (user_id, token)
+    );
+    ", [])?;
+    let is_updated = conn.prepare("SELECT user_id FROM migration LIMIT 1;").is_ok();
+    if is_updated { return Ok(()); }
+    println!("Upgrading migration table");
+    conn.execute("DROP TABLE IF EXISTS migration_new;", [])?;
+    conn.execute("
+    CREATE TABLE migration_new (
+        user_id           BIGINT NOT NULL,
+        token             TEXT NOT NULL,
+        password          TEXT NOT NULL,
+        PRIMARY KEY (user_id, token)
+    );
+    ", [])?;
+    
+    let mut stmt = conn.prepare("SELECT token, password FROM migration")?;
+
+    let rows = stmt.query_map([], |row| {
+        Ok((
+            row.get::<_, String>(0)?,
+            row.get::<_, String>(1)?,
+        ))
+    })?;
+
+    let mut insert_new_row = conn.prepare("INSERT INTO migration_new (user_id, token, password) VALUES (?1, ?2, ?3)")?;
+    for row in rows {
+        let (token, password) = row?;
+
+        let user_id = code_to_uid(&token);
+
+        insert_new_row.execute(params![user_id, token, password])?;
+    }
+    conn.execute("DROP TABLE migration;", params!())?;
+    conn.execute("ALTER TABLE migration_new RENAME TO migration;", params!())?;
+    Ok(())
+}
+
+fn code_to_uid(code: &str) -> String {
+    code
+        .replace('7', "")
+        .replace('A', "1")
+        .replace('G', "2")
+        .replace('W', "3")
+        .replace('Q', "4")
+        .replace('Y', "5")
+        .replace('6', "6")
+        .replace('I', "7")
+        .replace('P', "8")
+        .replace('U', "9")
+        .replace('M', "0")
+}

src/sql.rs

@@ -33,6 +33,13 @@ impl SQLite {
             }
         })
     }
+    pub fn lock_and_select_type<T: rusqlite::types::FromSql>(&self, command: &str, args: &[&dyn ToSql]) -> Result<T, rusqlite::Error> {
+        let conn = Connection::open(&self.path).unwrap();
+        let mut stmt = conn.prepare(command)?;
+        stmt.query_row(args, |row| {
+            row.get(0)
+        })
+    }
     pub fn lock_and_select_all(&self, command: &str, args: &[&dyn ToSql]) -> Result<JsonValue, rusqlite::Error> {
         let conn = Connection::open(&self.path).unwrap();
         let mut stmt = conn.prepare(command)?;