Remove duplicate database related code

This commit is contained in:
Ethan O'Brien
2024-05-04 14:03:19 -05:00
parent 5c21bb6e2d
commit 83571e03c2
5 changed files with 188 additions and 314 deletions

View File

@ -1,5 +1,4 @@
use rusqlite::{Connection, params, ToSql};
use std::sync::{Mutex, MutexGuard};
use rusqlite::params;
use lazy_static::lazy_static;
use json::{JsonValue, array, object};
use crate::router::global;
@ -7,109 +6,29 @@ use uuid::Uuid;
use rand::Rng;
use sha2::{Digest, Sha256};
use base64::{Engine as _, engine::general_purpose};
use crate::sql::SQLite;
lazy_static! {
static ref ENGINE: Mutex<Option<Connection>> = Mutex::new(None);
static ref DATABASE: SQLite = SQLite::new("userdata.db");
static ref NEW_USER: JsonValue = {
json::parse(include_str!("new_user.json")).unwrap()
};
}
fn init(engine: &mut MutexGuard<'_, Option<Connection>>) {
let conn = Connection::open("userdata.db").unwrap();
conn.execute("PRAGMA foreign_keys = ON;", ()).unwrap();
engine.replace(conn);
}
fn lock_and_exec(command: &str, args: &[&dyn ToSql]) {
loop {
match ENGINE.lock() {
Ok(mut result) => {
if result.is_none() {
init(&mut result);
}
let conn = result.as_ref().unwrap();
conn.execute(command, args).unwrap();
return;
}
Err(_) => {
std::thread::sleep(std::time::Duration::from_millis(15));
}
}
}
}
fn lock_and_select(command: &str, args: &[&dyn ToSql]) -> Result<String, rusqlite::Error> {
loop {
match ENGINE.lock() {
Ok(mut result) => {
if result.is_none() {
init(&mut result);
}
let conn = result.as_ref().unwrap();
let mut stmt = conn.prepare(command)?;
return stmt.query_row(args, |row| {
match row.get::<usize, i64>(0) {
Ok(val) => Ok(val.to_string()),
Err(_) => row.get(0)
}
});
}
Err(_) => {
std::thread::sleep(std::time::Duration::from_millis(15));
}
}
}
}
fn lock_and_select_all(command: &str, args: &[&dyn ToSql]) -> Result<JsonValue, rusqlite::Error> {
loop {
match ENGINE.lock() {
Ok(mut result) => {
if result.is_none() {
init(&mut result);
}
let conn = result.as_ref().unwrap();
let mut stmt = conn.prepare(command)?;
let map = stmt.query_map(args, |row| {
match row.get::<usize, i64>(0) {
Ok(val) => Ok(val.to_string()),
Err(_) => row.get(0)
}
})?;
let mut rv = array![];
for val in map {
let res = val?;
match res.clone().parse::<i64>() {
Ok(v) => rv.push(v).unwrap(),
Err(_) => rv.push(res).unwrap()
};
}
return Ok(rv);
}
Err(_) => {
std::thread::sleep(std::time::Duration::from_millis(15));
}
}
}
}
fn create_store_v2(table: &str) {
lock_and_exec(table, params!());
}
fn create_token_store() {
create_store_v2("CREATE TABLE IF NOT EXISTS tokens (
DATABASE.create_store_v2("CREATE TABLE IF NOT EXISTS tokens (
user_id BIGINT NOT NULL PRIMARY KEY,
token TEXT NOT NULL
)");
}
fn create_migration_store() {
create_store_v2("CREATE TABLE IF NOT EXISTS migration (
DATABASE.create_store_v2("CREATE TABLE IF NOT EXISTS migration (
token TEXT NOT NULL PRIMARY KEY,
password TEXT NOT NULL
)");
}
fn create_users_store() {
create_store_v2("CREATE TABLE IF NOT EXISTS users (
DATABASE.create_store_v2("CREATE TABLE IF NOT EXISTS users (
user_id BIGINT NOT NULL PRIMARY KEY,
userdata TEXT NOT NULL,
userhome TEXT NOT NULL,
@ -123,7 +42,7 @@ fn create_users_store() {
fn acc_exists(uid: i64) -> bool {
create_users_store();
lock_and_select("SELECT user_id FROM users WHERE user_id=?1", params!(uid)).is_ok()
DATABASE.lock_and_select("SELECT user_id FROM users WHERE user_id=?1", params!(uid)).is_ok()
}
fn get_key(auth_key: &str) -> i64 {
let uid = get_uid(&auth_key);
@ -140,7 +59,7 @@ fn get_key(auth_key: &str) -> i64 {
key
}
fn uid_exists(uid: i64) -> bool {
let data = lock_and_select("SELECT user_id FROM users WHERE user_id=?1", params!(uid));
let data = DATABASE.lock_and_select("SELECT user_id FROM users WHERE user_id=?1", params!(uid));
data.is_ok()
}
@ -162,7 +81,7 @@ fn create_acc(uid: i64, login: &str) {
new_user["user"]["id"] = uid.into();
new_user["stamina"]["last_updated_time"] = global::timestamp().into();
lock_and_exec("INSERT INTO users (user_id, userdata, userhome, missions, loginbonus, sifcards, friends, friend_request_disabled) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)", params!(
DATABASE.lock_and_exec("INSERT INTO users (user_id, userdata, userhome, missions, loginbonus, sifcards, friends, friend_request_disabled) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)", params!(
uid,
json::stringify(new_user),
include_str!("new_user_home.json"),
@ -174,13 +93,13 @@ fn create_acc(uid: i64, login: &str) {
));
create_token_store();
lock_and_exec("DELETE FROM tokens WHERE token=?1", params!(login));
lock_and_exec("INSERT INTO tokens (user_id, token) VALUES (?1, ?2)", params!(uid, login));
DATABASE.lock_and_exec("DELETE FROM tokens WHERE token=?1", params!(login));
DATABASE.lock_and_exec("INSERT INTO tokens (user_id, token) VALUES (?1, ?2)", params!(uid, login));
}
fn get_uid(token: &str) -> i64 {
create_token_store();
let data = lock_and_select("SELECT user_id FROM tokens WHERE token = ?1;", params!(token));
let data = DATABASE.lock_and_select("SELECT user_id FROM tokens WHERE token = ?1;", params!(token));
if !data.is_ok() {
return 0;
}
@ -190,7 +109,7 @@ fn get_uid(token: &str) -> i64 {
fn get_login_token(uid: i64) -> String {
create_token_store();
let data = lock_and_select("SELECT token FROM tokens WHERE user_id=?1", params!(uid));
let data = DATABASE.lock_and_select("SELECT token FROM tokens WHERE user_id=?1", params!(uid));
if !data.is_ok() {
return String::new();
}
@ -200,7 +119,7 @@ fn get_login_token(uid: i64) -> String {
fn get_data(auth_key: &str, row: &str) -> JsonValue {
let key = get_key(&auth_key);
let result = lock_and_select(&format!("SELECT {} FROM users WHERE user_id=?1", row), params!(key));
let result = DATABASE.lock_and_select(&format!("SELECT {} FROM users WHERE user_id=?1", row), params!(key));
json::parse(&result.unwrap()).unwrap()
}
@ -268,11 +187,11 @@ pub fn get_acc_friends(auth_key: &str) -> JsonValue {
pub fn save_data(auth_key: &str, row: &str, data: JsonValue) {
let key = get_key(&auth_key);
lock_and_exec(&format!("UPDATE users SET {}=?1 WHERE user_id=?2", row), params!(json::stringify(data), key));
DATABASE.lock_and_exec(&format!("UPDATE users SET {}=?1 WHERE user_id=?2", row), params!(json::stringify(data), key));
}
pub fn save_acc(auth_key: &str, data: JsonValue) {
lock_and_exec("UPDATE users SET friend_request_disabled=?1 WHERE user_id=?2", params!(data["user"]["friend_request_disabled"].as_i32().unwrap(), get_key(&auth_key)));
DATABASE.lock_and_exec("UPDATE users SET friend_request_disabled=?1 WHERE user_id=?2", params!(data["user"]["friend_request_disabled"].as_i32().unwrap(), get_key(&auth_key)));
save_data(auth_key, "userdata", data);
}
pub fn save_acc_home(auth_key: &str, data: JsonValue) {
@ -331,7 +250,7 @@ fn verify_password(password: &str, salted_hash: &str) -> bool {
pub fn get_acc_transfer(uid: i64, token: &str, password: &str) -> JsonValue {
create_migration_store();
let data = lock_and_select("SELECT password FROM migration WHERE token=?1", params!(token));
let data = DATABASE.lock_and_select("SELECT password FROM migration WHERE token=?1", params!(token));
if !data.is_ok() {
return object!{success: false};
}
@ -347,8 +266,8 @@ pub fn get_acc_transfer(uid: i64, token: &str, password: &str) -> JsonValue {
pub fn save_acc_transfer(token: &str, password: &str) {
create_migration_store();
lock_and_exec("DELETE FROM migration WHERE token=?1", params!(token));
lock_and_exec("INSERT INTO migration (token, password) VALUES (?1, ?2)", params!(token, hash_password(password)));
DATABASE.lock_and_exec("DELETE FROM migration WHERE token=?1", params!(token));
DATABASE.lock_and_exec("INSERT INTO migration (token, password) VALUES (?1, ?2)", params!(token, hash_password(password)));
}
pub fn get_name_and_rank(uid: i64) -> JsonValue {
@ -366,7 +285,7 @@ pub fn get_name_and_rank(uid: i64) -> JsonValue {
user_rank: 1
}
}
let result = lock_and_select("SELECT userdata FROM users WHERE user_id=?1", params!(uid));
let result = DATABASE.lock_and_select("SELECT userdata FROM users WHERE user_id=?1", params!(uid));
let data = json::parse(&result.unwrap()).unwrap();
return object!{
@ -386,7 +305,7 @@ pub fn get_acc_from_uid(uid: i64) -> JsonValue {
if uid == 0 || !acc_exists(uid) {
return object!{"error": true}
}
let result = lock_and_select("SELECT userdata FROM users WHERE user_id=?1", params!(uid));
let result = DATABASE.lock_and_select("SELECT userdata FROM users WHERE user_id=?1", params!(uid));
json::parse(&result.unwrap()).unwrap()
}
@ -396,11 +315,11 @@ pub fn friend_request(uid: i64, requestor: i64) {
return;
}
let uid = get_uid(&login_token);
let friends = lock_and_select("SELECT friends FROM users WHERE user_id=?1", params!(uid));
let friends = DATABASE.lock_and_select("SELECT friends FROM users WHERE user_id=?1", params!(uid));
let mut friends = json::parse(&friends.unwrap()).unwrap();
if !friends["pending_user_id_list"].contains(requestor) {
friends["pending_user_id_list"].push(requestor).unwrap();
lock_and_exec("UPDATE users SET friends=?1 WHERE user_id=?2", params!(json::stringify(friends), uid));
DATABASE.lock_and_exec("UPDATE users SET friends=?1 WHERE user_id=?2", params!(json::stringify(friends), uid));
}
}
@ -410,7 +329,7 @@ pub fn friend_request_approve(uid: i64, requestor: i64, accepted: bool, key: &st
return;
}
let uid = get_uid(&login_token);
let friends = lock_and_select("SELECT friends FROM users WHERE user_id=?1", params!(uid));
let friends = DATABASE.lock_and_select("SELECT friends FROM users WHERE user_id=?1", params!(uid));
let mut friends = json::parse(&friends.unwrap()).unwrap();
let index = friends[key].members().into_iter().position(|r| *r.to_string() == requestor.to_string());
if !index.is_none() {
@ -423,7 +342,7 @@ pub fn friend_request_approve(uid: i64, requestor: i64, accepted: bool, key: &st
if accepted && !friends["friend_user_id_list"].contains(requestor) {
friends["friend_user_id_list"].push(requestor).unwrap();
}
lock_and_exec("UPDATE users SET friends=?1 WHERE user_id=?2", params!(json::stringify(friends), uid));
DATABASE.lock_and_exec("UPDATE users SET friends=?1 WHERE user_id=?2", params!(json::stringify(friends), uid));
}
pub fn friend_request_disabled(uid: i64) -> bool {
@ -432,7 +351,7 @@ pub fn friend_request_disabled(uid: i64) -> bool {
return true;
}
let uid = get_uid(&login_token);
let user = lock_and_select("SELECT userdata FROM users WHERE user_id=?1", params!(uid));
let user = DATABASE.lock_and_select("SELECT userdata FROM users WHERE user_id=?1", params!(uid));
let user = json::parse(&user.unwrap()).unwrap();
user["user"]["friend_request_disabled"].to_string() == "1"
}
@ -443,24 +362,24 @@ pub fn friend_remove(uid: i64, requestor: i64) {
return;
}
let uid = get_uid(&login_token);
let friends = lock_and_select("SELECT friends FROM users WHERE user_id=?1", params!(uid));
let friends = DATABASE.lock_and_select("SELECT friends FROM users WHERE user_id=?1", params!(uid));
let mut friends = json::parse(&friends.unwrap()).unwrap();
let index = friends["friend_user_id_list"].members().into_iter().position(|r| *r.to_string() == requestor.to_string());
if !index.is_none() {
friends["friend_user_id_list"].array_remove(index.unwrap());
}
lock_and_exec("UPDATE users SET friends=?1 WHERE user_id=?2", params!(json::stringify(friends), uid));
DATABASE.lock_and_exec("UPDATE users SET friends=?1 WHERE user_id=?2", params!(json::stringify(friends), uid));
}
pub fn get_random_uids(count: i32) -> JsonValue {
if count <= 0 {
return array![];
}
lock_and_select_all(&format!("SELECT user_id FROM users WHERE friend_request_disabled=?1 ORDER BY RANDOM() LIMIT {}", count), params!(0)).unwrap()
DATABASE.lock_and_select_all(&format!("SELECT user_id FROM users WHERE friend_request_disabled=?1 ORDER BY RANDOM() LIMIT {}", count), params!(0)).unwrap()
}
fn create_webui_store() {
create_store_v2("CREATE TABLE IF NOT EXISTS webui (
DATABASE.create_store_v2("CREATE TABLE IF NOT EXISTS webui (
user_id BIGINT NOT NULL PRIMARY KEY,
token TEXT NOT NULL,
last_login BIGINT NOT NULL
@ -469,7 +388,7 @@ fn create_webui_store() {
fn create_webui_token() -> String {
let token = format!("{}", Uuid::new_v4());
if lock_and_select("SELECT user_id FROM webui WHERE token=?1", params!(token)).is_ok() {
if DATABASE.lock_and_select("SELECT user_id FROM webui WHERE token=?1", params!(token)).is_ok() {
return create_webui_token();
}
token
@ -478,7 +397,7 @@ fn create_webui_token() -> String {
pub fn webui_login(uid: i64, password: &str) -> Result<String, String> {
create_webui_store();
create_migration_store();
let pass = lock_and_select("SELECT password FROM migration WHERE token=?1", params!(crate::router::user::uid_to_code(uid.to_string()))).unwrap_or(String::new());
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(String::new());
if !verify_password(password, &pass) {
if acc_exists(uid) && pass == "" {
return Err(String::from("Migration token not set. Set token in game settings."));
@ -488,8 +407,8 @@ pub fn webui_login(uid: i64, password: &str) -> Result<String, String> {
let new_token = create_webui_token();
lock_and_exec("DELETE FROM webui WHERE user_id=?1", params!(uid));
lock_and_exec("INSERT INTO webui (user_id, token, last_login) VALUES (?1, ?2, ?3)", params!(uid, new_token, global::timestamp()));
DATABASE.lock_and_exec("DELETE FROM webui WHERE user_id=?1", params!(uid));
DATABASE.lock_and_exec("INSERT INTO webui (user_id, token, last_login) VALUES (?1, ?2, ?3)", params!(uid, new_token, global::timestamp()));
Ok(new_token)
}
@ -509,7 +428,7 @@ pub fn webui_import_user(user: JsonValue) -> Result<JsonValue, String> {
user["sif_cards"] = array![];
}
lock_and_exec("INSERT INTO users (user_id, userdata, userhome, missions, loginbonus, sifcards, friends, friend_request_disabled) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)", params!(
DATABASE.lock_and_exec("INSERT INTO users (user_id, userdata, userhome, missions, loginbonus, sifcards, friends, friend_request_disabled) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)", params!(
uid,
json::stringify(user["userdata"].clone()),
json::stringify(user["home"].clone()),
@ -527,7 +446,7 @@ pub fn webui_import_user(user: JsonValue) -> Result<JsonValue, String> {
token = format!("{}", Uuid::new_v4());
}
lock_and_exec("INSERT INTO tokens (user_id, token) VALUES (?1, ?2)", params!(uid, 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());
save_acc_transfer(&mig, &user["password"].to_string());
@ -539,7 +458,7 @@ pub fn webui_import_user(user: JsonValue) -> Result<JsonValue, String> {
}
fn webui_login_token(token: &str) -> Option<String> {
let uid = lock_and_select("SELECT user_id FROM webui WHERE token=?1", params!(token)).unwrap_or(String::new());
let uid = DATABASE.lock_and_select("SELECT user_id FROM webui WHERE token=?1", params!(token)).unwrap_or(String::new());
if uid == String::new() || token == "" {
return None;
}
@ -547,14 +466,14 @@ fn webui_login_token(token: &str) -> Option<String> {
if uid == 0 {
return None;
}
let last_login = lock_and_select("SELECT last_login FROM webui WHERE user_id=?1", params!(uid)).unwrap_or(String::new()).parse::<i64>().unwrap_or(0);
let last_login = DATABASE.lock_and_select("SELECT last_login FROM webui WHERE user_id=?1", params!(uid)).unwrap_or(String::new()).parse::<i64>().unwrap_or(0);
let limit = 24 * 60 * 60; //1 day
//Expired token
if (global::timestamp() as i64) > last_login + limit {
lock_and_exec("DELETE FROM webui WHERE user_id=?1", params!(uid));
DATABASE.lock_and_exec("DELETE FROM webui WHERE user_id=?1", params!(uid));
return None;
}
let login_token = lock_and_select("SELECT token FROM tokens WHERE user_id=?1", params!(uid)).unwrap_or(String::new());
let login_token = DATABASE.lock_and_select("SELECT token FROM tokens WHERE user_id=?1", params!(uid)).unwrap_or(String::new());
if login_token == String::new() {
return None;
}
@ -595,5 +514,5 @@ pub fn webui_start_loginbonus(bonus_id: i64, token: &str) -> JsonValue {
}
pub fn webui_logout(token: &str) {
lock_and_exec("DELETE FROM webui WHERE token=?1", params!(token));
DATABASE.lock_and_exec("DELETE FROM webui WHERE token=?1", params!(token));
}