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; use crate::admin_meta; /// Database backend types #[derive(Debug, Clone, Serialize, Deserialize)] pub enum BackendType { Redb, Sled, Tantivy, // Full-text search backend (no KV storage) // 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 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; // ----- Full-text (Tantivy) minimal RPC endpoints ----- /// Create a new FT index in a Tantivy-backed DB #[method(name = "ftCreate")] async fn ft_create( &self, db_id: u64, index_name: String, schema: Vec<(String, String, Vec)>, ) -> RpcResult; /// Add or replace a document in an FT index #[method(name = "ftAdd")] async fn ft_add( &self, db_id: u64, index_name: String, doc_id: String, score: f64, fields: HashMap, ) -> RpcResult; /// Search an FT index #[method(name = "ftSearch")] async fn ft_search( &self, db_id: u64, index_name: String, query: String, filters: Option>, limit: Option, offset: Option, return_fields: Option>, ) -> RpcResult; /// Delete a document by id from an FT index #[method(name = "ftDel")] async fn ft_del(&self, db_id: u64, index_name: String, doc_id: String) -> RpcResult; /// Get FT index info #[method(name = "ftInfo")] async fn ft_info(&self, db_id: u64, index_name: String) -> RpcResult; /// Drop an FT index #[method(name = "ftDrop")] async fn ft_drop(&self, db_id: u64, index_name: String) -> RpcResult; } /// RPC Server implementation pub struct RpcServerImpl { /// Base directory for database files base_dir: String, /// Managed database servers servers: Arc>>>, /// Default backend type backend: crate::options::BackendType, /// Admin secret used to encrypt DB 0 and authorize admin access admin_secret: String, } impl RpcServerImpl { /// Create a new RPC server instance pub fn new(base_dir: String, backend: crate::options::BackendType, admin_secret: String) -> Self { Self { base_dir, servers: Arc::new(RwLock::new(HashMap::new())), backend, admin_secret, } } /// 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()); } } // Validate existence via admin DB 0 (metadata), not filesystem presence let exists = admin_meta::db_exists(&self.base_dir, self.backend.clone(), &self.admin_secret, db_id) .map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?; if !exists { return Err(jsonrpsee::types::ErrorObjectOwned::owned( -32000, format!("Database {} not found", db_id), None::<()> )); } // Resolve effective backend for this db from admin meta or filesystem; fallback to default let meta_backend = admin_meta::get_database_backend(&self.base_dir, self.backend.clone(), &self.admin_secret, db_id) .ok() .flatten(); let db_path = std::path::PathBuf::from(&self.base_dir).join(format!("{}.db", db_id)); let sniffed_backend = if db_path.exists() { if db_path.is_file() { Some(crate::options::BackendType::Redb) } else if db_path.is_dir() { Some(crate::options::BackendType::Sled) } else { None } } else { None }; let effective_backend = meta_backend.clone().or(sniffed_backend).unwrap_or(self.backend.clone()); if effective_backend != self.backend { eprintln!( "notice: get_or_create_server: db {} backend resolved to {:?} (server default {:?})", db_id, effective_backend, self.backend ); } // If we had to sniff (no meta), persist the resolved backend if meta_backend.is_none() { let _ = admin_meta::set_database_backend(&self.base_dir, self.backend.clone(), &self.admin_secret, db_id, effective_backend.clone()); } // Create server instance with resolved backend let is_tantivy = matches!(effective_backend, crate::options::BackendType::Tantivy); 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: effective_backend.clone(), admin_secret: self.admin_secret.clone(), }; let mut server = Server::new(db_option).await; // Set the selected database to the db_id server.selected_db = db_id; // Lazily open/create physical storage according to admin meta (per-db encryption) // Skip for Tantivy backend (no KV storage to open) if !is_tantivy { let _ = server.current_storage(); } // 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 IDs from admin DB 0 async fn discover_databases(&self) -> Vec { admin_meta::list_dbs(&self.base_dir, self.backend.clone(), &self.admin_secret) .unwrap_or_default() } /// Build database file path for given server/db_id fn db_file_path(&self, server: &Server, db_id: u64) -> std::path::PathBuf { std::path::PathBuf::from(&server.option.dir).join(format!("{}.db", db_id)) } /// Recursively compute size on disk for the database path fn compute_size_on_disk(&self, path: &std::path::Path) -> Option { fn dir_size(p: &std::path::Path) -> u64 { if p.is_file() { std::fs::metadata(p).map(|m| m.len()).unwrap_or(0) } else if p.is_dir() { let mut total = 0u64; if let Ok(read) = std::fs::read_dir(p) { for entry in read.flatten() { total += dir_size(&entry.path()); } } total } else { 0 } } Some(dir_size(path)) } /// Extract created and last access times (secs) from a path, with fallbacks fn get_file_times_secs(path: &std::path::Path) -> (u64, Option) { let now = std::time::SystemTime::now(); let created = std::fs::metadata(path) .and_then(|m| m.created().or_else(|_| m.modified())) .unwrap_or(now) .duration_since(std::time::UNIX_EPOCH) .unwrap_or_default() .as_secs(); let last_access = std::fs::metadata(path) .and_then(|m| m.accessed()) .ok() .and_then(|t| t.duration_since(std::time::UNIX_EPOCH).ok().map(|d| d.as_secs())); (created, last_access) } /// Compose a DatabaseInfo by probing storage and filesystem, with admin meta for access key count async fn build_database_info(&self, db_id: u64, server: &Server) -> DatabaseInfo { // Probe storage to determine encryption state let storage = server.current_storage().ok(); let encrypted = storage.as_ref().map(|s| s.is_encrypted()).unwrap_or(server.option.encrypt); // Get actual key count from storage let key_count = storage.as_ref() .and_then(|s| s.dbsize().ok()) .map(|count| count as u64); // Get database name from admin meta let name = admin_meta::get_database_name(&self.base_dir, self.backend.clone(), &self.admin_secret, db_id) .ok() .flatten(); // Compute size on disk and timestamps from the DB file path let db_path = self.db_file_path(server, db_id); let size_on_disk = self.compute_size_on_disk(&db_path); let (created_at, last_access) = Self::get_file_times_secs(&db_path); let backend = match server.option.backend { crate::options::BackendType::Redb => BackendType::Redb, crate::options::BackendType::Sled => BackendType::Sled, crate::options::BackendType::Tantivy => BackendType::Tantivy, }; DatabaseInfo { id: db_id, name, backend, encrypted, redis_version: Some("7.0".to_string()), storage_path: Some(server.option.dir.clone()), size_on_disk, key_count, created_at, last_access, } } } #[jsonrpsee::core::async_trait] impl RpcServer for RpcServerImpl { async fn create_database( &self, backend: BackendType, config: DatabaseConfig, encryption_key: Option, ) -> RpcResult { // Allocate new ID via admin DB 0 let db_id = admin_meta::allocate_next_id(&self.base_dir, self.backend.clone(), &self.admin_secret) .map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?; // Persist per-db encryption key in admin DB 0 if provided if let Some(ref key) = encryption_key { admin_meta::set_enc_key(&self.base_dir, self.backend.clone(), &self.admin_secret, db_id, key) .map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?; } // Persist database name if provided if let Some(ref name) = config.name { admin_meta::set_database_name(&self.base_dir, self.backend.clone(), &self.admin_secret, db_id, name) .map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?; } // Ensure base dir exists if let Err(e) = std::fs::create_dir_all(&self.base_dir) { return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, format!("Failed to ensure base dir: {}", e), None::<()>)); } // Map RPC backend to options backend and persist it in admin meta for this db id let opt_backend = match backend { BackendType::Redb => crate::options::BackendType::Redb, BackendType::Sled => crate::options::BackendType::Sled, BackendType::Tantivy => crate::options::BackendType::Tantivy, }; admin_meta::set_database_backend(&self.base_dir, self.backend.clone(), &self.admin_secret, db_id, opt_backend.clone()) .map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?; // Create server instance using base_dir, chosen backend and admin secret let is_tantivy_new = matches!(opt_backend, crate::options::BackendType::Tantivy); let option = DBOption { dir: self.base_dir.clone(), port: 0, // Not used for RPC-managed databases debug: false, encryption_key: None, // per-db key is stored in admin DB 0 encrypt: false, // encryption decided per-db at open time backend: opt_backend.clone(), admin_secret: self.admin_secret.clone(), }; let mut server = Server::new(option).await; server.selected_db = db_id; // Initialize storage to create physical .db with proper encryption from admin meta // Skip for Tantivy backend (no KV storage to initialize) if !is_tantivy_new { let _ = server.current_storage(); } // Store the server in cache 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 { // 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 { if let Ok(server) = self.get_or_create_server(db_id).await { // Build accurate info from storage/meta/fs let info = self.build_database_info(db_id, &server).await; 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?; // Build accurate info from storage/meta/fs let info = self.build_database_info(db_id, &server).await; Ok(info) } 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) } // ----- Full-text (Tantivy) minimal RPC endpoints ----- async fn ft_create( &self, db_id: u64, index_name: String, schema: Vec<(String, String, Vec)>, ) -> RpcResult { let server = self.get_or_create_server(db_id).await?; if db_id == 0 { return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "FT not allowed on DB 0", None::<()>)); } if !matches!(server.option.backend, crate::options::BackendType::Tantivy) { return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "DB backend is not Tantivy", None::<()>)); } let proto = crate::search_cmd::ft_create_cmd(&*server, index_name, schema) .await .map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?; match proto { crate::protocol::Protocol::Error(msg) => { Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, msg, None::<()>)) } _ => Ok(true), } } async fn ft_add( &self, db_id: u64, index_name: String, doc_id: String, score: f64, fields: HashMap, ) -> RpcResult { let server = self.get_or_create_server(db_id).await?; if db_id == 0 { return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "FT not allowed on DB 0", None::<()>)); } if !matches!(server.option.backend, crate::options::BackendType::Tantivy) { return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "DB backend is not Tantivy", None::<()>)); } let proto = crate::search_cmd::ft_add_cmd(&*server, index_name, doc_id, score, fields) .await .map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?; match proto { crate::protocol::Protocol::Error(msg) => { Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, msg, None::<()>)) } _ => Ok(true), } } async fn ft_search( &self, db_id: u64, index_name: String, query: String, filters: Option>, limit: Option, offset: Option, return_fields: Option>, ) -> RpcResult { let server = self.get_or_create_server(db_id).await?; if db_id == 0 { return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "FT not allowed on DB 0", None::<()>)); } if !matches!(server.option.backend, crate::options::BackendType::Tantivy) { return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "DB backend is not Tantivy", None::<()>)); } let proto = crate::search_cmd::ft_search_cmd( &*server, index_name, query, filters.unwrap_or_default(), limit, offset, return_fields, ) .await .map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?; match proto { crate::protocol::Protocol::Error(msg) => { Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, msg, None::<()>)) } _ => Ok(serde_json::json!({ "resp": proto.encode() })), } } async fn ft_del(&self, db_id: u64, index_name: String, doc_id: String) -> RpcResult { let server = self.get_or_create_server(db_id).await?; if db_id == 0 { return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "FT not allowed on DB 0", None::<()>)); } if !matches!(server.option.backend, crate::options::BackendType::Tantivy) { return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "DB backend is not Tantivy", None::<()>)); } let proto = crate::search_cmd::ft_del_cmd(&*server, index_name, doc_id) .await .map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?; match proto { crate::protocol::Protocol::Error(msg) => { Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, msg, None::<()>)) } crate::protocol::Protocol::SimpleString(s) => Ok(s == "1"), _ => Ok(false), } } async fn ft_info(&self, db_id: u64, index_name: String) -> RpcResult { let server = self.get_or_create_server(db_id).await?; if db_id == 0 { return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "FT not allowed on DB 0", None::<()>)); } if !matches!(server.option.backend, crate::options::BackendType::Tantivy) { return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "DB backend is not Tantivy", None::<()>)); } let proto = crate::search_cmd::ft_info_cmd(&*server, index_name) .await .map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?; match proto { crate::protocol::Protocol::Error(msg) => { Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, msg, None::<()>)) } _ => Ok(serde_json::json!({ "resp": proto.encode() })), } } async fn ft_drop(&self, db_id: u64, index_name: String) -> RpcResult { let server = self.get_or_create_server(db_id).await?; if db_id == 0 { return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "FT not allowed on DB 0", None::<()>)); } if !matches!(server.option.backend, crate::options::BackendType::Tantivy) { return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "DB backend is not Tantivy", None::<()>)); } let proto = crate::search_cmd::ft_drop_cmd(&*server, index_name) .await .map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?; match proto { crate::protocol::Protocol::Error(msg) => { Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, msg, None::<()>)) } crate::protocol::Protocol::SimpleString(s) => Ok(s.eq_ignore_ascii_case("OK")), _ => Ok(false), } } async fn add_access_key(&self, db_id: u64, key: String, permissions: String) -> RpcResult { 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::<()> )), }; admin_meta::add_access_key(&self.base_dir, self.backend.clone(), &self.admin_secret, db_id, &key, perms) .map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?; Ok(true) } async fn delete_access_key(&self, db_id: u64, key_hash: String) -> RpcResult { let ok = admin_meta::delete_access_key(&self.base_dir, self.backend.clone(), &self.admin_secret, db_id, &key_hash) .map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?; Ok(ok) } async fn list_access_keys(&self, db_id: u64) -> RpcResult> { let pairs = admin_meta::list_access_keys(&self.base_dir, self.backend.clone(), &self.admin_secret, db_id) .map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?; let keys: Vec = pairs.into_iter().map(|(hash, perm, ts)| AccessKeyInfo { hash, permissions: perm, created_at: ts, }).collect(); Ok(keys) } async fn set_database_public(&self, db_id: u64, public: bool) -> RpcResult { admin_meta::set_database_public(&self.base_dir, self.backend.clone(), &self.admin_secret, db_id, public) .map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?; Ok(true) } }