Add database module with connection manager and schema initialization

- Introduced `DatabaseManager` for handling SQLite connections per user with WAL and foreign key support.
- Implemented `initialize_schema` to set up comprehensive database schema covering AI, books, characters, locations, series, and sync tracking.
- Added migration helpers for schema versioning and column updates.
- Structured database module into `connection.rs` and `schema.rs` for clearer organization.
This commit is contained in:
natreex
2026-03-21 09:35:06 -04:00
parent ee4438834c
commit 1f99fe0bdc
6 changed files with 1952 additions and 0 deletions

View File

@@ -0,0 +1,127 @@
use aes::cipher::{BlockDecryptMut, BlockEncryptMut, KeyIvInit};
use base64::{engine::general_purpose::STANDARD as BASE64, Engine};
use hex;
use pbkdf2::pbkdf2_hmac;
use rand::RngCore;
use sha2::{Digest, Sha256, Sha512};
use crate::error::{AppError, AppResult};
const KEY_LENGTH: usize = 32; // 256 bits
const IV_LENGTH: usize = 16; // 128 bits
const SALT_LENGTH: usize = 64;
const PBKDF2_ITERATIONS: u32 = 100_000;
type Aes256CbcEnc = cbc::Encryptor<aes::Aes256>;
type Aes256CbcDec = cbc::Decryptor<aes::Aes256>;
/// Generate a unique encryption key for a user.
/// Key is generated once at first login and stored securely in keyring.
/// Returns Base64 encoded salt+key combination.
pub fn generate_user_encryption_key(user_id: &str) -> AppResult<String> {
let mut salt = vec![0u8; SALT_LENGTH];
rand::rng().fill_bytes(&mut salt);
let mut key = vec![0u8; KEY_LENGTH];
pbkdf2_hmac::<Sha512>(user_id.as_bytes(), &salt, PBKDF2_ITERATIONS, &mut key);
let mut combined = Vec::with_capacity(SALT_LENGTH + KEY_LENGTH);
combined.extend_from_slice(&salt);
combined.extend_from_slice(&key);
Ok(BASE64.encode(&combined))
}
/// Extract the actual encryption key from the stored combined salt+key.
fn extract_key_from_stored(stored_key: &str) -> AppResult<Vec<u8>> {
let combined = BASE64
.decode(stored_key)
.map_err(|e| AppError::Encryption(format!("Invalid base64 key: {}", e)))?;
if combined.len() < SALT_LENGTH + KEY_LENGTH {
return Err(AppError::Encryption("Stored key too short".to_string()));
}
Ok(combined[SALT_LENGTH..SALT_LENGTH + KEY_LENGTH].to_vec())
}
/// Encrypt data with user key (AES-256-CBC).
/// Returns encrypted string with format "iv_hex:encrypted_hex".
/// Compatible with the Fastify/Electron TypeScript implementation.
pub fn encrypt_data_with_user_key(data: &str, user_key: &str) -> AppResult<String> {
let key = extract_key_from_stored(user_key)?;
let mut iv = [0u8; IV_LENGTH];
rand::rng().fill_bytes(&mut iv);
let data_bytes = data.as_bytes();
// PKCS7 padding
let block_size = 16;
let padding_len = block_size - (data_bytes.len() % block_size);
let mut buf = Vec::with_capacity(data_bytes.len() + padding_len);
buf.extend_from_slice(data_bytes);
buf.extend(std::iter::repeat(padding_len as u8).take(padding_len));
let key_array: [u8; 32] = key
.as_slice()
.try_into()
.map_err(|_| AppError::Encryption("Invalid key length".to_string()))?;
let cipher = Aes256CbcEnc::new(&key_array.into(), &iv.into());
let padded_length = buf.len();
cipher
.encrypt_padded_mut::<aes::cipher::block_padding::NoPadding>(&mut buf, padded_length)
.map_err(|encryption_error| AppError::Encryption(format!("Encryption failed: {}", encryption_error)))?;
let iv_hex = hex::encode(iv);
let encrypted_hex = hex::encode(&buf);
Ok(format!("{}:{}", iv_hex, encrypted_hex))
}
/// Decrypt data with user key (AES-256-CBC).
/// Expects encrypted string with format "iv_hex:encrypted_hex".
/// Compatible with the Fastify/Electron TypeScript implementation.
pub fn decrypt_data_with_user_key(encrypted_data: &str, user_key: &str) -> AppResult<String> {
let parts: Vec<&str> = encrypted_data.splitn(2, ':').collect();
if parts.len() != 2 {
return Err(AppError::Encryption(
"Invalid encrypted data format, expected 'iv:data'".to_string(),
));
}
let iv = hex::decode(parts[0])
.map_err(|e| AppError::Encryption(format!("Invalid IV hex: {}", e)))?;
let encrypted_bytes = hex::decode(parts[1])
.map_err(|e| AppError::Encryption(format!("Invalid encrypted hex: {}", e)))?;
let key = extract_key_from_stored(user_key)?;
let key_array: [u8; 32] = key
.as_slice()
.try_into()
.map_err(|_| AppError::Encryption("Invalid key length".to_string()))?;
let iv_array: [u8; 16] = iv
.as_slice()
.try_into()
.map_err(|_| AppError::Encryption("Invalid IV length".to_string()))?;
let mut buf = encrypted_bytes.clone();
let cipher = Aes256CbcDec::new(&key_array.into(), &iv_array.into());
let decrypted = cipher
.decrypt_padded_mut::<aes::cipher::block_padding::Pkcs7>(&mut buf)
.map_err(|e| AppError::Encryption(format!("Decryption failed: {}", e)))?;
String::from_utf8(decrypted.to_vec())
.map_err(|e| AppError::Encryption(format!("Invalid UTF-8 after decryption: {}", e)))
}
/// Hash data using SHA-256 (for non-reversible hashing like titles).
/// Returns hex encoded hash.
pub fn hash_element(data: &str) -> String {
let normalized = data.to_lowercase().trim().to_string();
let mut hasher = Sha256::new();
hasher.update(normalized.as_bytes());
hex::encode(hasher.finalize())
}

View File

@@ -0,0 +1,167 @@
use std::collections::HashMap;
use std::fs;
use std::path::PathBuf;
use aes::cipher::{BlockDecryptMut, BlockEncryptMut, KeyIvInit};
use base64::{engine::general_purpose::STANDARD as BASE64, Engine};
use rand::RngCore;
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use crate::error::{AppError, AppResult};
const SERVICE_NAME: &str = "com.eritors.scribe.desktop";
const IV_LENGTH: usize = 16;
type Aes256CbcEnc = cbc::Encryptor<aes::Aes256>;
type Aes256CbcDec = cbc::Decryptor<aes::Aes256>;
// ===== Secure vault: encrypted JSON file (like Electron's safeStorage) =====
#[derive(Serialize, Deserialize, Default)]
struct SecureVault {
token: Option<String>,
encryption_keys: HashMap<String, String>,
last_user_id: Option<String>,
pin_hashes: HashMap<String, String>,
}
fn vault_path() -> PathBuf {
dirs_next::data_dir()
.unwrap_or_else(|| PathBuf::from("."))
.join(SERVICE_NAME)
.join("secure-config.json")
}
fn derive_machine_key() -> [u8; 32] {
let mut hasher = Sha256::new();
hasher.update(SERVICE_NAME.as_bytes());
if let Ok(name) = hostname::get() {
hasher.update(name.as_encoded_bytes());
}
if let Some(dir) = dirs_next::home_dir() {
hasher.update(dir.to_string_lossy().as_bytes());
}
hasher.finalize().into()
}
fn encrypt_vault(data: &[u8], key: &[u8; 32]) -> AppResult<Vec<u8>> {
let mut iv = [0u8; IV_LENGTH];
rand::rng().fill_bytes(&mut iv);
let block_size = 16;
let padding_len = block_size - (data.len() % block_size);
let mut buf = Vec::with_capacity(data.len() + padding_len);
buf.extend_from_slice(data);
buf.extend(std::iter::repeat(padding_len as u8).take(padding_len));
let padded_len = buf.len();
let cipher = Aes256CbcEnc::new(key.into(), &iv.into());
cipher.encrypt_padded_mut::<aes::cipher::block_padding::NoPadding>(&mut buf, padded_len)
.map_err(|e| AppError::Encryption(format!("Vault encrypt failed: {}", e)))?;
let mut result = Vec::with_capacity(IV_LENGTH + buf.len());
result.extend_from_slice(&iv);
result.extend_from_slice(&buf);
Ok(result)
}
fn decrypt_vault(data: &[u8], key: &[u8; 32]) -> AppResult<Vec<u8>> {
if data.len() < IV_LENGTH + 1 {
return Err(AppError::Encryption("Vault data too short".into()));
}
let iv: [u8; 16] = data[..IV_LENGTH].try_into()
.map_err(|_| AppError::Encryption("Invalid IV".into()))?;
let mut buf = data[IV_LENGTH..].to_vec();
let cipher = Aes256CbcDec::new(key.into(), &iv.into());
let decrypted = cipher.decrypt_padded_mut::<aes::cipher::block_padding::Pkcs7>(&mut buf)
.map_err(|e| AppError::Encryption(format!("Vault decrypt failed: {}", e)))?;
Ok(decrypted.to_vec())
}
fn read_vault() -> SecureVault {
let path = vault_path();
let key = derive_machine_key();
match fs::read_to_string(&path) {
Ok(content) => {
if let Ok(raw) = BASE64.decode(content.trim()) {
if let Ok(decrypted) = decrypt_vault(&raw, &key) {
if let Ok(vault) = serde_json::from_slice::<SecureVault>(&decrypted) {
return vault;
}
}
}
SecureVault::default()
}
Err(_) => SecureVault::default(),
}
}
fn write_vault(vault: &SecureVault) -> AppResult<()> {
let path = vault_path();
if let Some(parent) = path.parent() {
fs::create_dir_all(parent).map_err(|e| AppError::Internal(format!("Failed to create vault dir: {}", e)))?;
}
let json = serde_json::to_string(vault)
.map_err(|e| AppError::Internal(format!("Failed to serialize vault: {}", e)))?;
let key = derive_machine_key();
let encrypted = encrypt_vault(json.as_bytes(), &key)?;
let encoded = BASE64.encode(&encrypted);
fs::write(&path, encoded).map_err(|e| AppError::Internal(format!("Failed to write vault: {}", e)))
}
// ===== Public API (same signatures as before) =====
pub fn get_user_encryption_key(user_id: &str) -> AppResult<String> {
let vault = read_vault();
vault.encryption_keys.get(user_id).cloned()
.ok_or_else(|| AppError::Keyring(format!("No encryption key for user {}", user_id)))
}
pub fn set_user_encryption_key(user_id: &str, encryption_key: &str) -> AppResult<()> {
let mut vault = read_vault();
vault.encryption_keys.insert(user_id.to_string(), encryption_key.to_string());
write_vault(&vault)
}
pub fn has_user_encryption_key(user_id: &str) -> bool {
let vault = read_vault();
vault.encryption_keys.contains_key(user_id)
}
pub fn get_token() -> AppResult<Option<String>> {
let vault = read_vault();
Ok(vault.token)
}
pub fn set_token(token: &str) -> AppResult<()> {
let mut vault = read_vault();
vault.token = Some(token.to_string());
write_vault(&vault)
}
pub fn remove_token() -> AppResult<()> {
let mut vault = read_vault();
vault.token = None;
write_vault(&vault)
}
pub fn set_pin_hash(user_id: &str, pin_hash: &str) -> AppResult<()> {
let mut vault = read_vault();
vault.pin_hashes.insert(user_id.to_string(), pin_hash.to_string());
write_vault(&vault)
}
pub fn get_pin_hash(user_id: &str) -> AppResult<Option<String>> {
let vault = read_vault();
Ok(vault.pin_hashes.get(user_id).cloned())
}
pub fn set_last_user_id(user_id: &str) -> AppResult<()> {
let mut vault = read_vault();
vault.last_user_id = Some(user_id.to_string());
write_vault(&vault)
}
pub fn get_last_user_id() -> AppResult<Option<String>> {
let vault = read_vault();
Ok(vault.last_user_id)
}

View File

@@ -0,0 +1,2 @@
pub mod encryption;
pub mod key_manager;

View File

@@ -0,0 +1,62 @@
use rusqlite::Connection;
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::{Arc, Mutex};
use crate::error::{AppError, AppResult};
pub struct DatabaseManager {
connections: HashMap<String, Connection>,
base_path: PathBuf,
}
impl DatabaseManager {
pub fn new(base_path: PathBuf) -> Self {
Self {
connections: HashMap::new(),
base_path,
}
}
pub fn get_db_path(&self, user_id: &str) -> PathBuf {
self.base_path.join(format!("eritors-local-{}.db", user_id))
}
pub fn initialize(&mut self, user_id: &str) -> AppResult<()> {
if self.connections.contains_key(user_id) {
return Ok(());
}
let db_path = self.get_db_path(user_id);
if let Some(parent) = db_path.parent() {
std::fs::create_dir_all(parent)?;
}
let conn = Connection::open(&db_path)?;
conn.execute_batch("PRAGMA journal_mode=WAL;")?;
conn.execute_batch("PRAGMA foreign_keys=ON;")?;
self.connections.insert(user_id.to_string(), conn);
Ok(())
}
pub fn get_connection(&self, user_id: &str) -> AppResult<&Connection> {
self.connections
.get(user_id)
.ok_or_else(|| AppError::Database(rusqlite::Error::InvalidParameterName(
format!("No database connection for user: {}", user_id),
)))
}
pub fn close(&mut self, user_id: &str) {
self.connections.remove(user_id);
}
}
pub type DbManager = Arc<Mutex<DatabaseManager>>;
pub fn create_db_manager(base_path: PathBuf) -> DbManager {
Arc::new(Mutex::new(DatabaseManager::new(base_path)))
}

2
src-tauri/src/db/mod.rs Normal file
View File

@@ -0,0 +1,2 @@
pub mod connection;
pub mod schema;

1592
src-tauri/src/db/schema.rs Normal file

File diff suppressed because it is too large Load Diff