use super::*;
use async_trait::async_trait;
use std::collections::HashMap;
use tokio::sync::RwLock;
#[cfg(feature = "sqlite")]
use sqlx::Row;
#[async_trait]
pub trait PermissionStorage: Send + Sync {
async fn store_persistent_permission(
&self,
permission: GrantedPermission,
) -> Result<(), PermissionError>;
async fn get_persistent_permissions(
&self,
session_id: &str,
) -> Result<Vec<GrantedPermission>, PermissionError>;
async fn revoke_persistent_permission(
&self,
permission_id: Uuid,
) -> Result<(), PermissionError>;
async fn clear_expired_permissions(&self) -> Result<u64, PermissionError>;
}
pub struct InMemoryPermissionStorage {
permissions: RwLock<HashMap<String, Vec<GrantedPermission>>>,
}
impl InMemoryPermissionStorage {
pub fn new() -> Self {
Self {
permissions: RwLock::new(HashMap::new()),
}
}
}
#[async_trait]
impl PermissionStorage for InMemoryPermissionStorage {
async fn store_persistent_permission(
&self,
permission: GrantedPermission,
) -> Result<(), PermissionError> {
let mut permissions = self.permissions.write().await;
permissions
.entry(permission.request.session_id.clone())
.or_insert_with(Vec::new)
.push(permission);
Ok(())
}
async fn get_persistent_permissions(
&self,
session_id: &str,
) -> Result<Vec<GrantedPermission>, PermissionError> {
let permissions = self.permissions.read().await;
Ok(permissions.get(session_id).cloned().unwrap_or_default())
}
async fn revoke_persistent_permission(
&self,
permission_id: Uuid,
) -> Result<(), PermissionError> {
let mut permissions = self.permissions.write().await;
for session_perms in permissions.values_mut() {
session_perms.retain(|p| p.request.id != permission_id);
}
Ok(())
}
async fn clear_expired_permissions(&self) -> Result<u64, PermissionError> {
let mut permissions = self.permissions.write().await;
let mut cleared = 0u64;
let now = chrono::Utc::now();
for session_perms in permissions.values_mut() {
let original_len = session_perms.len();
session_perms.retain(|p| {
if let Some(expires_at) = p.expires_at {
now <= expires_at
} else {
true
}
});
cleared += (original_len - session_perms.len()) as u64;
}
Ok(cleared)
}
}
#[cfg(feature = "sqlite")]
pub struct SqlitePermissionStorage {
pool: sqlx::SqlitePool,
}
#[cfg(feature = "sqlite")]
impl SqlitePermissionStorage {
pub async fn new(database_url: &str) -> Result<Self, PermissionError> {
let pool = sqlx::SqlitePool::connect(database_url).await?;
sqlx::query(r#"
CREATE TABLE IF NOT EXISTS permissions (
id TEXT PRIMARY KEY,
session_id TEXT NOT NULL,
tool_name TEXT NOT NULL,
permission_type TEXT NOT NULL,
description TEXT NOT NULL,
action TEXT NOT NULL,
path TEXT,
params TEXT NOT NULL,
granted_at TEXT NOT NULL,
expires_at TEXT,
persistent BOOLEAN NOT NULL DEFAULT FALSE
)
"#)
.execute(&pool)
.await?;
sqlx::query(r#"
CREATE INDEX IF NOT EXISTS idx_permissions_session
ON permissions(session_id)
"#)
.execute(&pool)
.await?;
Ok(Self { pool })
}
}
#[cfg(feature = "sqlite")]
#[async_trait]
impl PermissionStorage for SqlitePermissionStorage {
async fn store_persistent_permission(
&self,
permission: GrantedPermission,
) -> Result<(), PermissionError> {
sqlx::query(r#"
INSERT INTO permissions (
id, session_id, tool_name, permission_type, description,
action, path, params, granted_at, expires_at, persistent
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
"#)
.bind(permission.request.id.to_string())
.bind(&permission.request.session_id)
.bind(&permission.request.tool_name)
.bind(serde_json::to_string(&permission.request.permission)?)
.bind(&permission.request.description)
.bind(&permission.request.action)
.bind(permission.request.path.as_ref().map(|p| p.to_string_lossy().to_string()))
.bind(serde_json::to_string(&permission.request.params)?)
.bind(permission.granted_at.to_rfc3339())
.bind(permission.expires_at.map(|t| t.to_rfc3339()))
.bind(permission.persistent)
.execute(&self.pool)
.await?;
Ok(())
}
async fn get_persistent_permissions(
&self,
session_id: &str,
) -> Result<Vec<GrantedPermission>, PermissionError> {
let rows = sqlx::query(r#"
SELECT id, session_id, tool_name, permission_type, description,
action, path, params, granted_at, expires_at, persistent
FROM permissions
WHERE session_id = ? AND persistent = TRUE
"#)
.bind(session_id)
.fetch_all(&self.pool)
.await?;
let mut permissions = Vec::new();
for row in rows {
let permission_request = PermissionRequest {
id: Uuid::parse_str(&row.get::<String, _>("id"))
.map_err(|e| PermissionError::Integration(format!("Invalid UUID: {}", e)))?,
session_id: row.get("session_id"),
tool_name: row.get("tool_name"),
permission: serde_json::from_str(&row.get::<String, _>("permission_type"))?,
description: row.get("description"),
action: row.get("action"),
path: row.get::<Option<String>, _>("path").map(PathBuf::from),
params: serde_json::from_str(&row.get::<String, _>("params"))?,
created_at: chrono::DateTime::parse_from_rfc3339(&row.get::<String, _>("granted_at"))
.map_err(|e| PermissionError::Integration(format!("Invalid datetime: {}", e)))?
.with_timezone(&chrono::Utc),
};
let granted_permission = GrantedPermission {
request: permission_request,
granted_at: chrono::DateTime::parse_from_rfc3339(&row.get::<String, _>("granted_at"))
.map_err(|e| PermissionError::Integration(format!("Invalid datetime: {}", e)))?
.with_timezone(&chrono::Utc),
expires_at: row.get::<Option<String>, _>("expires_at")
.map(|s| chrono::DateTime::parse_from_rfc3339(&s))
.transpose()
.map_err(|e| PermissionError::Integration(format!("Invalid datetime: {}", e)))?
.map(|dt| dt.with_timezone(&chrono::Utc)),
persistent: row.get("persistent"),
};
permissions.push(granted_permission);
}
Ok(permissions)
}
async fn revoke_persistent_permission(
&self,
permission_id: Uuid,
) -> Result<(), PermissionError> {
sqlx::query("DELETE FROM permissions WHERE id = ?")
.bind(permission_id.to_string())
.execute(&self.pool)
.await?;
Ok(())
}
async fn clear_expired_permissions(&self) -> Result<u64, PermissionError> {
let result = sqlx::query(r#"
DELETE FROM permissions
WHERE expires_at IS NOT NULL AND expires_at < ?
"#)
.bind(chrono::Utc::now().to_rfc3339())
.execute(&self.pool)
.await?;
Ok(result.rows_affected())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_in_memory_storage() {
let storage = InMemoryPermissionStorage::new();
let request = create_permission_request(
"test-session",
"file_tool",
Permission::FileRead,
"Test permission",
"read",
Some(PathBuf::from("/tmp/test.txt")),
serde_json::json!({}),
);
let granted = GrantedPermission {
request,
granted_at: chrono::Utc::now(),
expires_at: None,
persistent: true,
};
storage.store_persistent_permission(granted.clone()).await.unwrap();
let permissions = storage.get_persistent_permissions("test-session").await.unwrap();
assert_eq!(permissions.len(), 1);
assert_eq!(permissions[0].request.tool_name, "file_tool");
}
#[tokio::test]
async fn test_clear_expired_permissions() {
let storage = InMemoryPermissionStorage::new();
let request = create_permission_request(
"test-session",
"file_tool",
Permission::FileRead,
"Test permission",
"read",
None,
serde_json::json!({}),
);
let expired = GrantedPermission {
request,
granted_at: chrono::Utc::now() - chrono::Duration::hours(1),
expires_at: Some(chrono::Utc::now() - chrono::Duration::minutes(30)),
persistent: true,
};
storage.store_persistent_permission(expired).await.unwrap();
let cleared = storage.clear_expired_permissions().await.unwrap();
assert_eq!(cleared, 1);
let permissions = storage.get_persistent_permissions("test-session").await.unwrap();
assert_eq!(permissions.len(), 0);
}
}