GithubFollow/One-KVM
1
0
Fork 0
mirror of https://github.com/mofeng-git/One-KVM.git synced 2026-09-12 10:34:24 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

refactor: 优化数据库初始化与 WOL 存储

Browse Source
This commit is contained in:
mofeng-git
2026-08-30 21:25:39 +08:00
parent 48fd64319e
commit dbe5672df4
10 changed files with 171 additions and 175 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/atx/mod.rs
View File

@@ -27,7 +27,7 @@ pub use types::{
ActiveLevel, AtxAction, AtxDevices, AtxDriverType, AtxInputBinding, AtxKeyConfig,
AtxOutputBinding, AtxPowerRequest, AtxState, HddStatus, PowerStatus, LCUS_RELAY_MAX_CHANNEL,
};
pub use wol::{list_wol_history, record_wol_history, send_wol};
pub use wol::send_wol;
#[cfg(any(unix, test))]
fn hidraw_uevent_is_usb_relay(uevent: &str) -> bool {

51
src/atx/wol.rs
View File

@@ -7,8 +7,6 @@ use tracing::info;
use crate::error::{AppError, Result};
const WOL_HISTORY_MAX_ENTRIES: i64 = 50;
const MAGIC_PACKET_SIZE: usize = 102;
fn parse_mac_address(mac: &str) -> Result<[u8; 6]> {
@@ -118,55 +116,6 @@ pub fn send_wol(mac_address: &str, interface: Option<&str>) -> Result<()> {
Ok(())
}
pub async fn record_wol_history(pool: &sqlx::Pool<sqlx::Sqlite>, mac_address: &str) -> Result<()> {
sqlx::query(
r#"
INSERT INTO wol_history (mac_address, updated_at)
VALUES (?1, CAST(strftime('%s', 'now') AS INTEGER))
ON CONFLICT(mac_address) DO UPDATE SET
updated_at = excluded.updated_at
"#,
)
.bind(mac_address)
.execute(pool)
.await?;
sqlx::query(
r#"
DELETE FROM wol_history
WHERE mac_address NOT IN (
SELECT mac_address FROM wol_history
ORDER BY updated_at DESC
LIMIT ?1
)
"#,
)
.bind(WOL_HISTORY_MAX_ENTRIES)
.execute(pool)
.await?;
Ok(())
}
pub async fn list_wol_history(
pool: &sqlx::Pool<sqlx::Sqlite>,
limit: usize,
) -> Result<Vec<(String, i64)>> {
let rows = sqlx::query_as(
r#"
SELECT mac_address, updated_at
FROM wol_history
ORDER BY updated_at DESC
LIMIT ?1
"#,
)
.bind(limit as i64)
.fetch_all(pool)
.await?;
Ok(rows)
}
#[cfg(test)]
mod tests {
use super::*;

5
src/auth/two_factor.rs
View File

@@ -227,15 +227,14 @@ impl TwoFactorService {
return Err(AppError::AuthError("Invalid TOTP code".to_string()));
}
let mut transaction = self.pool.begin().await?;
let result =
sqlx::query("INSERT INTO user_totp_credentials (user_id, secret) VALUES (?1, ?2)")
.bind(user_id)
.bind(secret.to_string())
.execute(&mut *transaction)
.execute(&self.pool)
.await;
match result {
Ok(_) => transaction.commit().await?,
Ok(_) => {}
Err(sqlx::Error::Database(error)) if error.is_unique_violation() => {
return Err(AppError::Conflict("TOTP is already enabled".to_string()));
}

16
src/auth/user.rs
View File

@@ -1,7 +1,5 @@
use serde::{Deserialize, Serialize};
use sqlx::{Pool, Sqlite};
use time::format_description::well_known::Rfc3339;
use time::OffsetDateTime;
use uuid::Uuid;
use super::password::{hash_password, verify_password};
@@ -112,12 +110,10 @@ impl UserStore {
}
let password_hash = hash_password(new_password)?;
let now = OffsetDateTime::now_utc();
let result =
sqlx::query("UPDATE users SET password_hash = ?1, updated_at = ?2 WHERE id = ?3")
let result = sqlx::query(
"UPDATE users SET password_hash = ?1, updated_at = datetime('now') WHERE id = ?2",
)
.bind(&password_hash)
.bind(now.format(&Rfc3339).expect("RFC3339 format"))
.bind(user_id)
.execute(&self.pool)
.await?;
@@ -143,10 +139,10 @@ impl UserStore {
return Ok(());
}
let now = OffsetDateTime::now_utc();
let result = sqlx::query("UPDATE users SET username = ?1, updated_at = ?2 WHERE id = ?3")
let result = sqlx::query(
"UPDATE users SET username = ?1, updated_at = datetime('now') WHERE id = ?2",
)
.bind(new_username)
.bind(now.format(&Rfc3339).expect("RFC3339 format"))
.bind(user_id)
.execute(&self.pool)
.await?;

14
src/config/store.rs
View File

@@ -17,13 +17,13 @@ pub struct ConfigStore {
}
impl ConfigStore {
pub fn new(pool: Pool<Sqlite>) -> Result<Self> {
Ok(Self {
pub fn new(pool: Pool<Sqlite>) -> Self {
Self {
pool,
cache: Arc::new(ArcSwap::from_pointee(AppConfig::default())),
change_tx: broadcast::channel(16).0,
write_lock: Arc::new(Mutex::new(())),
})
}
}
pub async fn load(&self) -> Result<()> {
@@ -145,7 +145,7 @@ mod tests {
let db = DatabasePool::new(&db_path).await.unwrap();
db.init_schema().await.unwrap();
let store = ConfigStore::new(db.clone_pool()).unwrap();
let store = ConfigStore::new(db.clone_pool());
store.load().await.unwrap();
let config = store.get();
@@ -163,7 +163,7 @@ mod tests {
assert!(config.initialized);
assert_eq!(config.web.http_port, 9000);
let store2 = ConfigStore::new(db.clone_pool()).unwrap();
let store2 = ConfigStore::new(db.clone_pool());
store2.load().await.unwrap();
let config = store2.get();
assert!(config.initialized);
@@ -176,7 +176,7 @@ mod tests {
let db_path = dir.path().join("test.db");
let db = DatabasePool::new(&db_path).await.unwrap();
db.init_schema().await.unwrap();
let store = ConfigStore::new(db.clone_pool()).unwrap();
let store = ConfigStore::new(db.clone_pool());
store.load().await.unwrap();
sqlx::query("DROP TABLE config")
@@ -210,7 +210,7 @@ mod tests {
.await
.unwrap();
let store = ConfigStore::new(db.clone_pool()).unwrap();
let store = ConfigStore::new(db.clone_pool());
store.load().await.unwrap();
let (persisted,): (String,) =
sqlx::query_as("SELECT value FROM config WHERE key = 'app_config'")

2
src/db/mod.rs
View File

@@ -1,10 +1,12 @@
mod pool;
mod wol_history;
use std::path::Path;
use crate::error::Result;
pub use pool::DatabasePool;
pub use wol_history::WolHistoryStore;
/// Open the application database stored in `data_dir` and ensure its schema exists.
pub async fn open_database_pool(data_dir: &Path) -> Result<DatabasePool> {

147
src/db/pool.rs
View File

@@ -1,4 +1,7 @@
use sqlx::{sqlite::SqlitePoolOptions, Pool, Sqlite};
use sqlx::{
sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePoolOptions, SqliteSynchronous},
Pool, Sqlite,
};
use std::path::Path;
use std::time::Duration;
@@ -15,29 +18,82 @@ impl DatabasePool {
tokio::fs::create_dir_all(parent).await?;
}
let db_url = format!("sqlite:{}?mode=rwc", db_path.display());
let options = SqliteConnectOptions::new()
.filename(db_path)
.create_if_missing(true)
.foreign_keys(true)
.journal_mode(SqliteJournalMode::Wal)
.synchronous(SqliteSynchronous::Normal)
.busy_timeout(Duration::from_secs(5));
let pool = SqlitePoolOptions::new()
.max_connections(4)
.acquire_timeout(Duration::from_secs(5))
.idle_timeout(Duration::from_secs(300))
.connect(&db_url)
.connect_with(options)
.await?;
Ok(Self { pool })
}
pub async fn init_schema(&self) -> Result<()> {
self.create_config_table().await?;
self.create_users_table().await?;
self.create_user_totp_credentials_table().await?;
self.create_api_tokens_table().await?;
self.create_wol_history_table().await?;
// Keep migrations embedded in the binary so deployments do not need an
// extra migrations directory or another runtime dependency.
let mut transaction = self.pool.begin().await?;
sqlx::query(
r#"
CREATE TABLE IF NOT EXISTS schema_migrations (
version INTEGER PRIMARY KEY,
applied_at TEXT NOT NULL DEFAULT (datetime('now'))
)
"#,
)
.execute(&mut *transaction)
.await?;
let current_version: i64 =
sqlx::query_scalar("SELECT COALESCE(MAX(version), 0) FROM schema_migrations")
.fetch_one(&mut *transaction)
.await?;
for (version, statements) in SCHEMA_MIGRATIONS
.iter()
.enumerate()
.map(|(index, statements)| ((index + 1) as i64, *statements))
{
if version <= current_version {
continue;
}
for &statement in statements {
sqlx::query(statement).execute(&mut *transaction).await?;
}
sqlx::query("INSERT INTO schema_migrations (version) VALUES (?1)")
.bind(version)
.execute(&mut *transaction)
.await?;
}
transaction.commit().await?;
Ok(())
}
async fn create_config_table(&self) -> Result<()> {
sqlx::query(
pub fn pool(&self) -> &Pool<Sqlite> {
&self.pool
}
pub fn clone_pool(&self) -> Pool<Sqlite> {
self.pool.clone()
}
pub fn wol_history(&self) -> super::WolHistoryStore {
super::WolHistoryStore::new(self.pool.clone())
}
}
// Each item is one version; statements within a version run atomically.
// New schema changes should be appended as a new item, never edited in place.
const SCHEMA_MIGRATIONS: &[&[&str]] = &[
&[
r#"
CREATE TABLE IF NOT EXISTS config (
key TEXT PRIMARY KEY,
@@ -45,14 +101,6 @@ impl DatabasePool {
updated_at TEXT NOT NULL DEFAULT (datetime('now'))
)
"#,
)
.execute(&self.pool)
.await?;
Ok(())
}
async fn create_users_table(&self) -> Result<()> {
sqlx::query(
r#"
CREATE TABLE IF NOT EXISTS users (
id TEXT PRIMARY KEY,
@@ -62,14 +110,14 @@ impl DatabasePool {
updated_at TEXT NOT NULL DEFAULT (datetime('now'))
)
"#,
r#"
CREATE TABLE IF NOT EXISTS user_totp_credentials (
user_id TEXT PRIMARY KEY,
secret TEXT NOT NULL,
created_at TEXT NOT NULL DEFAULT (datetime('now')),
FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE
)
.execute(&self.pool)
.await?;
Ok(())
}
async fn create_api_tokens_table(&self) -> Result<()> {
sqlx::query(
"#,
r#"
CREATE TABLE IF NOT EXISTS api_tokens (
id TEXT PRIMARY KEY,
@@ -81,56 +129,19 @@ impl DatabasePool {
last_used TEXT
)
"#,
)
.execute(&self.pool)
.await?;
Ok(())
}
async fn create_user_totp_credentials_table(&self) -> Result<()> {
sqlx::query(
r#"
CREATE TABLE IF NOT EXISTS user_totp_credentials (
user_id TEXT PRIMARY KEY,
secret TEXT NOT NULL,
created_at TEXT NOT NULL DEFAULT (datetime('now')),
FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE
)
"#,
)
.execute(&self.pool)
.await?;
Ok(())
}
async fn create_wol_history_table(&self) -> Result<()> {
sqlx::query(
r#"
CREATE TABLE IF NOT EXISTS wol_history (
mac_address TEXT PRIMARY KEY,
updated_at INTEGER NOT NULL
)
"#,
)
.execute(&self.pool)
.await?;
sqlx::query(
r#"
CREATE INDEX IF NOT EXISTS idx_wol_history_updated_at
ON wol_history(updated_at DESC)
"#,
)
.execute(&self.pool)
.await?;
Ok(())
}
pub fn pool(&self) -> &Pool<Sqlite> {
&self.pool
}
pub fn clone_pool(&self) -> Pool<Sqlite> {
self.pool.clone()
}
}
],
&[r#"
CREATE UNIQUE INDEX IF NOT EXISTS idx_users_singleton
ON users ((1))
"#],
];

39
src/db/wol_history.rs Normal file
View File

@@ -0,0 +1,39 @@
use sqlx::{Pool, Sqlite};
use crate::error::Result;
const MAX_ENTRIES: i64 = 50;
#[derive(Clone)]
pub struct WolHistoryStore {
pool: Pool<Sqlite>,
}
impl WolHistoryStore {
pub(crate) fn new(pool: Pool<Sqlite>) -> Self {
Self { pool }
}
pub async fn record(&self, mac_address: &str) -> Result<()> {
let mut tx = self.pool.begin().await?;
sqlx::query("INSERT INTO wol_history (mac_address, updated_at) VALUES (?1, CAST(strftime('%s', 'now') AS INTEGER)) ON CONFLICT(mac_address) DO UPDATE SET updated_at = excluded.updated_at")
.bind(mac_address)
.execute(&mut *tx)
.await?;
sqlx::query("DELETE FROM wol_history WHERE mac_address NOT IN (SELECT mac_address FROM wol_history ORDER BY updated_at DESC LIMIT ?1)")
.bind(MAX_ENTRIES)
.execute(&mut *tx)
.await?;
tx.commit().await?;
Ok(())
}
pub async fn list(&self, limit: usize) -> Result<Vec<(String, i64)>> {
Ok(sqlx::query_as(
"SELECT mac_address, updated_at FROM wol_history ORDER BY updated_at DESC LIMIT ?1",
)
.bind(limit as i64)
.fetch_all(&self.pool)
.await?)
}
}

4
src/runtime/builder.rs
View File

@@ -226,7 +226,7 @@ async fn load_runtime_config(
) -> anyhow::Result<(DatabasePool, ConfigStore, AppConfig)> {
let db = open_database_pool(data_dir).await?;
let config_store = ConfigStore::new(db.clone_pool())?;
let config_store = ConfigStore::new(db.clone_pool());
config_store.load().await?;
let mut config = (*config_store.get()).clone();
config.apply_platform_defaults();
@@ -568,7 +568,7 @@ mod tests {
let data_dir = temp_dir.path().join("data");
let msd_dir = temp_dir.path().join("disabled-msd");
let db = open_database_pool(&data_dir).await.unwrap();
let config_store = ConfigStore::new(db.clone_pool()).unwrap();
let config_store = ConfigStore::new(db.clone_pool());
config_store.load().await.unwrap();
let mut config = (*config_store.get()).clone();
config.msd.enabled = false;

4
src/web/handlers/atx_api.rs
View File

@@ -171,7 +171,7 @@ pub async fn atx_wol(
// Send WOL packet
crate::atx::send_wol(&mac_address, interface)?;
if let Err(error) = crate::atx::record_wol_history(state.db.pool(), &mac_address).await {
if let Err(error) = state.db.wol_history().record(&mac_address).await {
warn!("Failed to persist WOL history: {}", error);
}
@@ -191,7 +191,7 @@ pub async fn atx_wol_history(
.unwrap_or(WOL_HISTORY_DEFAULT_LIMIT)
.clamp(1, WOL_HISTORY_MAX_LIMIT);
let rows = crate::atx::list_wol_history(state.db.pool(), limit).await?;
let rows = state.db.wol_history().list(limit).await?;
let history = rows
.into_iter()