WIP: adding access management control to db instances

src/rpc.rs

@@ -3,6 +3,7 @@ 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;
@@ -39,6 +40,43 @@ pub struct DatabaseInfo {
     pub last_access: Option<u64>,
 }
 
+/// 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<String, AccessKey>,
+}
+
+/// 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 {
@@ -70,6 +108,22 @@ pub trait Rpc {
     /// Get server statistics
     #[method(name = "getServerStats")]
     async fn get_server_stats(&self) -> RpcResult<HashMap<String, serde_json::Value>>;
+
+    /// Add an access key to a database
+    #[method(name = "addAccessKey")]
+    async fn add_access_key(&self, db_id: u64, key: String, permissions: String) -> RpcResult<bool>;
+
+    /// Delete an access key from a database
+    #[method(name = "deleteAccessKey")]
+    async fn delete_access_key(&self, db_id: u64, key_hash: String) -> RpcResult<bool>;
+
+    /// List all access keys for a database
+    #[method(name = "listAccessKeys")]
+    async fn list_access_keys(&self, db_id: u64) -> RpcResult<Vec<AccessKeyInfo>>;
+
+    /// Set database public/private status
+    #[method(name = "setDatabasePublic")]
+    async fn set_database_public(&self, db_id: u64, public: bool) -> RpcResult<bool>;
 }
 
 /// RPC Server implementation
@@ -84,6 +138,8 @@ pub struct RpcServerImpl {
     next_encrypted_id: Arc<RwLock<u64>>,
     /// Default backend type
     backend: crate::options::BackendType,
+    /// Encryption keys for databases
+    encryption_keys: Arc<RwLock<HashMap<u64, Option<String>>>>,
 }
 
 impl RpcServerImpl {
@@ -95,6 +151,7 @@ impl RpcServerImpl {
             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())),
         }
     }
 
@@ -177,6 +234,166 @@ impl RpcServerImpl {
             current_id
         }
     }
+
+    /// Load database metadata from file (static version)
+    pub async fn load_meta_static(base_dir: &str, db_id: u64) -> Result<DatabaseMeta, jsonrpsee::types::ErrorObjectOwned> {
+        let meta_path = std::path::PathBuf::from(base_dir).join(format!("{}_meta.json", db_id));
+
+        // If meta file doesn't exist, return default
+        if !meta_path.exists() {
+            return Ok(DatabaseMeta {
+                public: true,
+                keys: HashMap::new(),
+            });
+        }
+
+        // Read file
+        let content = std::fs::read(&meta_path)
+            .map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(
+                -32000,
+                format!("Failed to read meta file: {}", e),
+                None::<()>
+            ))?;
+
+        let json_str = String::from_utf8(content)
+            .map_err(|_| jsonrpsee::types::ErrorObjectOwned::owned(
+                -32000,
+                "Invalid UTF-8 in meta file",
+                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<DatabaseMeta, jsonrpsee::types::ErrorObjectOwned> {
+        let meta_path = std::path::PathBuf::from(&self.base_dir).join(format!("{}_meta.json", db_id));
+
+        // If meta file doesn't exist, create 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 and potentially decrypt
+        let content = std::fs::read(&meta_path)
+            .map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(
+                -32000,
+                format!("Failed to read meta file: {}", e),
+                None::<()>
+            ))?;
+
+        let json_str = if db_id >= 10 {
+            // Encrypted database, decrypt meta
+            if let Some(key) = self.encryption_keys.read().await.get(&db_id).and_then(|k| k.as_ref()) {
+                use crate::crypto::CryptoFactory;
+                let crypto = CryptoFactory::new(key.as_bytes());
+                String::from_utf8(crypto.decrypt(&content)
+                    .map_err(|_| jsonrpsee::types::ErrorObjectOwned::owned(
+                        -32000,
+                        "Failed to decrypt meta file",
+                        None::<()>
+                    ))?)
+                    .map_err(|_| jsonrpsee::types::ErrorObjectOwned::owned(
+                        -32000,
+                        "Invalid UTF-8 in decrypted meta",
+                        None::<()>
+                    ))?
+            } else {
+                return Err(jsonrpsee::types::ErrorObjectOwned::owned(
+                    -32000,
+                    "Encryption key not found for encrypted database",
+                    None::<()>
+                ));
+            }
+        } else {
+            String::from_utf8(content)
+                .map_err(|_| jsonrpsee::types::ErrorObjectOwned::owned(
+                    -32000,
+                    "Invalid UTF-8 in meta file",
+                    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(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(meta)
+            .map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(
+                -32000,
+                format!("Failed to serialize meta: {}", e),
+                None::<()>
+            ))?;
+
+        if db_id >= 10 {
+            // Encrypted database, encrypt meta
+            if let Some(key) = self.encryption_keys.read().await.get(&db_id).and_then(|k| k.as_ref()) {
+                use crate::crypto::CryptoFactory;
+                let crypto = CryptoFactory::new(key.as_bytes());
+                let encrypted = crypto.encrypt(json_str.as_bytes());
+                std::fs::write(&meta_path, encrypted)
+                    .map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(
+                        -32000,
+                        format!("Failed to write encrypted meta file: {}", e),
+                        None::<()>
+                    ))?;
+            } else {
+                return Err(jsonrpsee::types::ErrorObjectOwned::owned(
+                    -32000,
+                    "Encryption key not found for encrypted database",
+                    None::<()>
+                ));
+            }
+        } else {
+            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]
@@ -213,7 +430,7 @@ impl RpcServer for RpcServerImpl {
                     dir: db_dir.to_string_lossy().to_string(),
                     port: 0, // Not used for RPC-managed databases
                     debug: false,
-                    encryption_key,
+                    encryption_key: encryption_key.clone(),
                     encrypt,
                     backend: match backend {
                         BackendType::Redb => crate::options::BackendType::Redb,
@@ -230,6 +447,19 @@ impl RpcServer for RpcServerImpl {
                 // 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));
@@ -339,4 +569,66 @@ impl RpcServer for RpcServerImpl {
 
         Ok(stats)
     }
+
+    async fn add_access_key(&self, db_id: u64, key: String, permissions: String) -> RpcResult<bool> {
+        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<bool> {
+        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<Vec<AccessKeyInfo>> {
+        let meta = self.load_meta(db_id).await?;
+        let keys: Vec<AccessKeyInfo> = 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<bool> {
+        let mut meta = self.load_meta(db_id).await?;
+        meta.public = public;
+        self.save_meta(db_id, &meta).await?;
+        Ok(true)
+    }
 }