use std::collections::HashMap; use std::sync::Arc; use tokio::sync::RwLock; use jsonrpsee::{core::RpcResult, proc_macros::rpc}; use serde::{Deserialize, Serialize}; use sha2::{Digest, Sha256}; use crate::server::Server; use crate::options::DBOption; /// Database backend types #[derive(Debug, Clone, Serialize, Deserialize)] pub enum BackendType { Redb, Sled, // Future: InMemory, Custom(String) } /// Database configuration #[derive(Debug, Clone, Serialize, Deserialize)] pub struct DatabaseConfig { pub name: Option, pub storage_path: Option, pub max_size: Option, pub redis_version: Option, } /// Database information returned by metadata queries #[derive(Debug, Clone, Serialize, Deserialize)] pub struct DatabaseInfo { pub id: u64, pub name: Option, pub backend: BackendType, pub encrypted: bool, pub redis_version: Option, pub storage_path: Option, pub size_on_disk: Option, pub key_count: Option, pub created_at: u64, pub last_access: Option, } /// Access permissions for database keys #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] pub enum Permissions { Read, ReadWrite, } /// Access key information #[derive(Debug, Clone, Serialize, Deserialize)] pub struct AccessKey { pub hash: String, pub permissions: Permissions, pub created_at: u64, } /// Database metadata containing access keys #[derive(Debug, Clone, Serialize, Deserialize)] pub struct DatabaseMeta { pub public: bool, pub keys: HashMap, } /// Access key information returned by RPC #[derive(Debug, Clone, Serialize, Deserialize)] pub struct AccessKeyInfo { pub hash: String, pub permissions: Permissions, pub created_at: u64, } /// Hash a plaintext key using SHA-256 pub fn hash_key(key: &str) -> String { let mut hasher = Sha256::new(); hasher.update(key.as_bytes()); format!("{:x}", hasher.finalize()) } /// RPC trait for HeroDB management #[rpc(server, client, namespace = "herodb")] pub trait Rpc { /// Create a new database with specified configuration #[method(name = "createDatabase")] async fn create_database( &self, backend: BackendType, config: DatabaseConfig, encryption_key: Option, ) -> RpcResult; /// Set encryption for an existing database (write-only key) #[method(name = "setEncryption")] async fn set_encryption(&self, db_id: u64, encryption_key: String) -> RpcResult; /// List all managed databases #[method(name = "listDatabases")] async fn list_databases(&self) -> RpcResult>; /// Get detailed information about a specific database #[method(name = "getDatabaseInfo")] async fn get_database_info(&self, db_id: u64) -> RpcResult; /// Delete a database #[method(name = "deleteDatabase")] async fn delete_database(&self, db_id: u64) -> RpcResult; /// Get server statistics #[method(name = "getServerStats")] async fn get_server_stats(&self) -> RpcResult>; /// Add an access key to a database #[method(name = "addAccessKey")] async fn add_access_key(&self, db_id: u64, key: String, permissions: String) -> RpcResult; /// Delete an access key from a database #[method(name = "deleteAccessKey")] async fn delete_access_key(&self, db_id: u64, key_hash: String) -> RpcResult; /// List all access keys for a database #[method(name = "listAccessKeys")] async fn list_access_keys(&self, db_id: u64) -> RpcResult>; /// Set database public/private status #[method(name = "setDatabasePublic")] async fn set_database_public(&self, db_id: u64, public: bool) -> RpcResult; } /// RPC Server implementation pub struct RpcServerImpl { /// Base directory for database files base_dir: String, /// Managed database servers servers: Arc>>>, /// Next unencrypted database ID to assign next_unencrypted_id: Arc>, /// Next encrypted database ID to assign next_encrypted_id: Arc>, /// Default backend type backend: crate::options::BackendType, /// Encryption keys for databases encryption_keys: Arc>>>, } impl RpcServerImpl { /// Create a new RPC server instance pub fn new(base_dir: String, backend: crate::options::BackendType) -> Self { Self { base_dir, servers: Arc::new(RwLock::new(HashMap::new())), next_unencrypted_id: Arc::new(RwLock::new(0)), next_encrypted_id: Arc::new(RwLock::new(10)), backend, encryption_keys: Arc::new(RwLock::new(HashMap::new())), } } /// Get or create a server instance for the given database ID async fn get_or_create_server(&self, db_id: u64) -> Result, jsonrpsee::types::ErrorObjectOwned> { // Check if server already exists { let servers = self.servers.read().await; if let Some(server) = servers.get(&db_id) { return Ok(server.clone()); } } // Check if database file exists let db_path = std::path::PathBuf::from(&self.base_dir).join(format!("{}.db", db_id)); if !db_path.exists() { return Err(jsonrpsee::types::ErrorObjectOwned::owned( -32000, format!("Database {} not found", db_id), None::<()> )); } // Create server instance with default options let db_option = DBOption { dir: self.base_dir.clone(), port: 0, // Not used for RPC-managed databases debug: false, encryption_key: None, encrypt: false, backend: self.backend.clone(), }; let mut server = Server::new(db_option).await; // Set the selected database to the db_id for proper file naming server.selected_db = db_id; // Store the server let mut servers = self.servers.write().await; servers.insert(db_id, Arc::new(server.clone())); Ok(Arc::new(server)) } /// Discover existing database files in the base directory async fn discover_databases(&self) -> Vec { let mut db_ids = Vec::new(); if let Ok(entries) = std::fs::read_dir(&self.base_dir) { for entry in entries.flatten() { if let Ok(file_name) = entry.file_name().into_string() { // Check if it's a database file (ends with .db) if file_name.ends_with(".db") { // Extract database ID from filename (e.g., "11.db" -> 11) if let Some(id_str) = file_name.strip_suffix(".db") { if let Ok(db_id) = id_str.parse::() { db_ids.push(db_id); } } } } } } db_ids } /// Get the next available database ID async fn get_next_db_id(&self, is_encrypted: bool) -> u64 { if is_encrypted { let mut id = self.next_encrypted_id.write().await; let current_id = *id; *id += 1; current_id } else { let mut id = self.next_unencrypted_id.write().await; let current_id = *id; *id += 1; current_id } } /// Load database metadata from file (static version) pub async fn load_meta_static(base_dir: &str, db_id: u64) -> Result { let meta_path = std::path::PathBuf::from(base_dir).join(format!("{}_meta.json", db_id)); // If meta file doesn't exist, create and persist default if !meta_path.exists() { let default_meta = DatabaseMeta { public: true, keys: HashMap::new(), }; // Persist default metadata to disk Self::save_meta_static(base_dir, db_id, &default_meta).await?; return Ok(default_meta); } // Read file as UTF-8 JSON let json_str = std::fs::read_to_string(&meta_path) .map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned( -32000, format!("Failed to read meta file: {}", e), None::<()> ))?; serde_json::from_str(&json_str) .map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned( -32000, format!("Failed to parse meta JSON: {}", e), None::<()> )) } /// Load database metadata from file async fn load_meta(&self, db_id: u64) -> Result { let meta_path = std::path::PathBuf::from(&self.base_dir).join(format!("{}_meta.json", db_id)); // If meta file doesn't exist, create and persist default if !meta_path.exists() { let default_meta = DatabaseMeta { public: true, keys: HashMap::new(), }; self.save_meta(db_id, &default_meta).await?; return Ok(default_meta); } // Read file as UTF-8 JSON (meta files are always plain JSON) let json_str = std::fs::read_to_string(&meta_path) .map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned( -32000, format!("Failed to read meta file: {}", e), None::<()> ))?; serde_json::from_str(&json_str) .map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned( -32000, format!("Failed to parse meta JSON: {}", e), None::<()> )) } /// Save database metadata to file (static version) pub async fn save_meta_static(base_dir: &str, db_id: u64, meta: &DatabaseMeta) -> Result<(), jsonrpsee::types::ErrorObjectOwned> { let meta_path = std::path::PathBuf::from(base_dir).join(format!("{}_meta.json", db_id)); let json_str = serde_json::to_string_pretty(meta) .map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned( -32000, format!("Failed to serialize meta: {}", e), None::<()> ))?; std::fs::write(&meta_path, json_str) .map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned( -32000, format!("Failed to write meta file: {}", e), None::<()> ))?; Ok(()) } /// Save database metadata to file async fn save_meta(&self, db_id: u64, meta: &DatabaseMeta) -> Result<(), jsonrpsee::types::ErrorObjectOwned> { let meta_path = std::path::PathBuf::from(&self.base_dir).join(format!("{}_meta.json", db_id)); let json_str = serde_json::to_string_pretty(meta) .map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned( -32000, format!("Failed to serialize meta: {}", e), None::<()> ))?; // Meta files are always stored as plain JSON (even when data DB is encrypted) std::fs::write(&meta_path, json_str) .map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned( -32000, format!("Failed to write meta file: {}", e), None::<()> ))?; Ok(()) } } #[jsonrpsee::core::async_trait] impl RpcServer for RpcServerImpl { async fn create_database( &self, backend: BackendType, config: DatabaseConfig, encryption_key: Option, ) -> RpcResult { let db_id = self.get_next_db_id(encryption_key.is_some()).await; // Handle both Redb and Sled backends match backend { BackendType::Redb | BackendType::Sled => { // Create database directory let db_dir = if let Some(path) = &config.storage_path { std::path::PathBuf::from(path) } else { std::path::PathBuf::from(&self.base_dir).join(format!("rpc_db_{}", db_id)) }; // Ensure directory exists std::fs::create_dir_all(&db_dir) .map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned( -32000, format!("Failed to create directory: {}", e), None::<()> ))?; // Create DB options let encrypt = encryption_key.is_some(); let option = DBOption { dir: db_dir.to_string_lossy().to_string(), port: 0, // Not used for RPC-managed databases debug: false, encryption_key: encryption_key.clone(), encrypt, backend: match backend { BackendType::Redb => crate::options::BackendType::Redb, BackendType::Sled => crate::options::BackendType::Sled, }, }; // Create server instance let mut server = Server::new(option).await; // Set the selected database to the db_id for proper file naming server.selected_db = db_id; // Initialize the storage to create the database file let _ = server.current_storage(); // Store the encryption key { let mut keys = self.encryption_keys.write().await; keys.insert(db_id, encryption_key.clone()); } // Initialize meta file let meta = DatabaseMeta { public: true, keys: HashMap::new(), }; self.save_meta(db_id, &meta).await?; // Store the server let mut servers = self.servers.write().await; servers.insert(db_id, Arc::new(server)); Ok(db_id) } } } async fn set_encryption(&self, db_id: u64, _encryption_key: String) -> RpcResult { // Note: In a real implementation, we'd need to modify the existing database // For now, return false as encryption can only be set during creation let _servers = self.servers.read().await; // TODO: Implement encryption setting for existing databases Ok(false) } async fn list_databases(&self) -> RpcResult> { let db_ids = self.discover_databases().await; let mut result = Vec::new(); for db_id in db_ids { // Try to get or create server for this database if let Ok(server) = self.get_or_create_server(db_id).await { let backend = match server.option.backend { crate::options::BackendType::Redb => BackendType::Redb, crate::options::BackendType::Sled => BackendType::Sled, }; let info = DatabaseInfo { id: db_id, name: None, // TODO: Store name in server metadata backend, encrypted: server.option.encrypt, redis_version: Some("7.0".to_string()), // Default Redis compatibility storage_path: Some(server.option.dir.clone()), size_on_disk: None, // TODO: Calculate actual size key_count: None, // TODO: Get key count from storage created_at: std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) .unwrap() .as_secs(), last_access: None, }; result.push(info); } } Ok(result) } async fn get_database_info(&self, db_id: u64) -> RpcResult { let server = self.get_or_create_server(db_id).await?; let backend = match server.option.backend { crate::options::BackendType::Redb => BackendType::Redb, crate::options::BackendType::Sled => BackendType::Sled, }; Ok(DatabaseInfo { id: db_id, name: None, backend, encrypted: server.option.encrypt, redis_version: Some("7.0".to_string()), storage_path: Some(server.option.dir.clone()), size_on_disk: None, key_count: None, created_at: std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) .unwrap() .as_secs(), last_access: None, }) } async fn delete_database(&self, db_id: u64) -> RpcResult { let mut servers = self.servers.write().await; if let Some(_server) = servers.remove(&db_id) { // Clean up database files let db_path = std::path::PathBuf::from(&self.base_dir).join(format!("{}.db", db_id)); if db_path.exists() { if db_path.is_dir() { std::fs::remove_dir_all(&db_path).ok(); } else { std::fs::remove_file(&db_path).ok(); } } Ok(true) } else { Ok(false) } } async fn get_server_stats(&self) -> RpcResult> { let db_ids = self.discover_databases().await; let mut stats = HashMap::new(); stats.insert("total_databases".to_string(), serde_json::json!(db_ids.len())); stats.insert("uptime".to_string(), serde_json::json!( std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) .unwrap() .as_secs() )); Ok(stats) } async fn add_access_key(&self, db_id: u64, key: String, permissions: String) -> RpcResult { let mut meta = self.load_meta(db_id).await?; let perms = match permissions.to_lowercase().as_str() { "read" => Permissions::Read, "readwrite" => Permissions::ReadWrite, _ => return Err(jsonrpsee::types::ErrorObjectOwned::owned( -32000, "Invalid permissions: use 'read' or 'readwrite'", None::<()> )), }; let hash = hash_key(&key); let access_key = AccessKey { hash: hash.clone(), permissions: perms, created_at: std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) .unwrap() .as_secs(), }; meta.keys.insert(hash, access_key); self.save_meta(db_id, &meta).await?; Ok(true) } async fn delete_access_key(&self, db_id: u64, key_hash: String) -> RpcResult { let mut meta = self.load_meta(db_id).await?; if meta.keys.remove(&key_hash).is_some() { // If no keys left, make database public if meta.keys.is_empty() { meta.public = true; } self.save_meta(db_id, &meta).await?; Ok(true) } else { Ok(false) } } async fn list_access_keys(&self, db_id: u64) -> RpcResult> { let meta = self.load_meta(db_id).await?; let keys: Vec = meta.keys.values() .map(|k| AccessKeyInfo { hash: k.hash.clone(), permissions: k.permissions.clone(), created_at: k.created_at, }) .collect(); Ok(keys) } async fn set_database_public(&self, db_id: u64, public: bool) -> RpcResult { let mut meta = self.load_meta(db_id).await?; meta.public = public; self.save_meta(db_id, &meta).await?; Ok(true) } }