From 32b5f5b68cd9280d72ea3e0b4fde3d3a611d811b Mon Sep 17 00:00:00 2001 From: Ethan O'Brien <77750390+ethanaobrien@users.noreply.github.com> Date: Sat, 4 May 2024 14:37:32 -0500 Subject: [PATCH] Setup tables in a single init function --- src/router/clear_rate.rs | 15 +++++++-------- src/router/gree.rs | 15 +++++++++------ src/router/userdata/mod.rs | 26 +++++--------------------- src/sql.rs | 9 +++++---- 4 files changed, 26 insertions(+), 39 deletions(-) diff --git a/src/router/clear_rate.rs b/src/router/clear_rate.rs index 4785ed4..1253738 100644 --- a/src/router/clear_rate.rs +++ b/src/router/clear_rate.rs @@ -9,7 +9,7 @@ use crate::encryption; use crate::sql::SQLite; lazy_static! { - static ref DATABASE: SQLite = SQLite::new("live_statistics.db"); + static ref DATABASE: SQLite = SQLite::new("live_statistics.db", setup_tables); } pub struct Live { @@ -24,8 +24,8 @@ pub struct Live { pub master_pass: i64, } -fn create_store() { - DATABASE.lock_and_exec("CREATE TABLE IF NOT EXISTS lives ( +fn setup_tables(conn: &SQLite) { + conn.lock_and_exec("CREATE TABLE IF NOT EXISTS lives ( live_id INT NOT NULL PRIMARY KEY, normal_failed BIGINT NOT NULL, normal_pass BIGINT NOT NULL, @@ -36,13 +36,13 @@ fn create_store() { master_failed BIGINT NOT NULL, master_pass BIGINT NOT NULL )", params!()); -} - -fn update_live_score(id: i64, uid: i64, score: i64) { - DATABASE.lock_and_exec("CREATE TABLE IF NOT EXISTS scores ( + conn.lock_and_exec("CREATE TABLE IF NOT EXISTS scores ( live_id INT NOT NULL PRIMARY KEY, score_data TEXT NOT NULL )", params!()); +} + +fn update_live_score(id: i64, uid: i64, score: i64) { if uid == 0 || score == 0 { return; } @@ -181,7 +181,6 @@ fn get_json() -> JsonValue { } fn get_clearrate_json() -> JsonValue { - create_store(); loop { match CACHED_DATA.lock() { Ok(mut result) => { diff --git a/src/router/gree.rs b/src/router/gree.rs index aea5c88..8edc889 100644 --- a/src/router/gree.rs +++ b/src/router/gree.rs @@ -20,7 +20,15 @@ use openssl::hash::MessageDigest; use openssl::sign::Verifier; lazy_static! { - static ref DATABASE: SQLite = SQLite::new("gree.db"); + static ref DATABASE: SQLite = SQLite::new("gree.db", setup_tables); +} + +fn setup_tables(conn: &SQLite) { + conn.create_store_v2("CREATE TABLE IF NOT EXISTS users ( + cert TEXT NOT NULL, + uuid TEXT NOT NULL, + user_id BIGINT NOT NULL PRIMARY KEY + )"); } fn update_cert(uid: i64, cert: &str) { @@ -37,11 +45,6 @@ fn update_cert(uid: i64, cert: &str) { DATABASE.lock_and_exec("UPDATE users SET cert=?1 WHERE user_id=?2", params!(cert, uid)); } fn create_acc(cert: &str) -> String { - DATABASE.create_store_v2("CREATE TABLE IF NOT EXISTS users ( - cert TEXT NOT NULL, - uuid TEXT NOT NULL, - user_id BIGINT NOT NULL PRIMARY KEY - )"); let uuid = global::create_token(); let user = userdata::get_acc(&uuid); let user_id = user["user"]["id"].as_i64().unwrap(); diff --git a/src/router/userdata/mod.rs b/src/router/userdata/mod.rs index c121a7e..e073878 100644 --- a/src/router/userdata/mod.rs +++ b/src/router/userdata/mod.rs @@ -8,26 +8,22 @@ use base64::{Engine as _, engine::general_purpose}; use crate::sql::SQLite; lazy_static! { - static ref DATABASE: SQLite = SQLite::new("userdata.db"); + static ref DATABASE: SQLite = SQLite::new("userdata.db", setup_tables); static ref NEW_USER: JsonValue = { json::parse(include_str!("new_user.json")).unwrap() }; } -fn create_token_store() { - DATABASE.create_store_v2("CREATE TABLE IF NOT EXISTS tokens ( +fn setup_tables(conn: &SQLite) { + conn.create_store_v2("CREATE TABLE IF NOT EXISTS tokens ( user_id BIGINT NOT NULL PRIMARY KEY, token TEXT NOT NULL )"); -} -fn create_migration_store() { - DATABASE.create_store_v2("CREATE TABLE IF NOT EXISTS migration ( + conn.create_store_v2("CREATE TABLE IF NOT EXISTS migration ( token TEXT NOT NULL PRIMARY KEY, password TEXT NOT NULL )"); -} -fn create_users_store() { - DATABASE.create_store_v2("CREATE TABLE IF NOT EXISTS users ( + conn.create_store_v2("CREATE TABLE IF NOT EXISTS users ( user_id BIGINT NOT NULL PRIMARY KEY, userdata TEXT NOT NULL, userhome TEXT NOT NULL, @@ -40,7 +36,6 @@ fn create_users_store() { } fn acc_exists(uid: i64) -> bool { - create_users_store(); DATABASE.lock_and_select("SELECT user_id FROM users WHERE user_id=?1", params!(uid)).is_ok() } fn get_key(auth_key: &str) -> i64 { @@ -74,8 +69,6 @@ fn generate_uid() -> i64 { } fn create_acc(uid: i64, login: &str) { - create_users_store(); - let mut new_user = NEW_USER.clone(); new_user["user"]["id"] = uid.into(); new_user["stamina"]["last_updated_time"] = global::timestamp().into(); @@ -91,13 +84,11 @@ fn create_acc(uid: i64, login: &str) { 0 )); - create_token_store(); 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 = DATABASE.lock_and_select("SELECT user_id FROM tokens WHERE token = ?1;", params!(token)); if !data.is_ok() { return 0; @@ -108,7 +99,6 @@ fn get_uid(token: &str) -> i64 { // Needed by gree pub fn get_login_token(uid: i64) -> String { - create_token_store(); let data = DATABASE.lock_and_select("SELECT token FROM tokens WHERE user_id=?1", params!(uid)); if !data.is_ok() { return String::new(); @@ -249,7 +239,6 @@ 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 = DATABASE.lock_and_select("SELECT password FROM migration WHERE token=?1", params!(token)); if !data.is_ok() { return object!{success: false}; @@ -265,7 +254,6 @@ pub fn get_acc_transfer(uid: i64, token: &str, password: &str) -> JsonValue { } pub fn save_acc_transfer(token: &str, password: &str) { - create_migration_store(); 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))); } @@ -396,7 +384,6 @@ fn create_webui_token() -> String { pub fn webui_login(uid: i64, password: &str) -> Result { create_webui_store(); - create_migration_store(); 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 == "" { @@ -414,9 +401,6 @@ pub fn webui_login(uid: i64, password: &str) -> Result { pub fn webui_import_user(user: JsonValue) -> Result { let mut user = user; - create_webui_store(); - create_migration_store(); - create_token_store(); let uid = user["userdata"]["user"]["id"].as_i64().unwrap(); if acc_exists(uid) { return Err(String::from("User already exists")); diff --git a/src/sql.rs b/src/sql.rs index 199593e..f3b32de 100644 --- a/src/sql.rs +++ b/src/sql.rs @@ -9,14 +9,15 @@ pub struct SQLite { } impl SQLite { - pub fn new(path: &str) -> SQLite { + pub fn new(path: &str, setup: fn(&SQLite)) -> SQLite { let conn = Connection::open(path).unwrap(); conn.execute("PRAGMA foreign_keys = ON;", ()).unwrap(); - - SQLite { + let instance = SQLite { engine: Mutex::new(conn), sleep_duration: 10 - } + }; + setup(&instance); + instance } pub fn lock_and_exec(&self, command: &str, args: &[&dyn ToSql]) { loop {