From ffd53c46505493888b7d5413cc5bf3ddf872e625 Mon Sep 17 00:00:00 2001 From: Aaron Longwell Date: Thu, 15 Jan 2026 17:13:46 -0800 Subject: [PATCH] feat(sqlite): add fine-grained permission configuration for tools Add support for restricting SQLite tools to specific databases and tables with separate read/write permissions. This enables tighter security controls when exposing database tools to agents. - Add SqliteConfig with builder pattern for configuring permissions - Support AllowList/DenyList modes for table-level read/write access - Add SQL parser using sqlparser crate to extract table operations - Create ConfiguredXxxTool wrappers that enforce permissions - Add factory functions: tools_for_database(), read_only_tools_for_tables() - Validate empty deny lists at build() time with Result-based error handling --- Cargo.toml | 1 + mixtape-tools/Cargo.toml | 3 +- mixtape-tools/src/sqlite/config.rs | 385 ++++++++++++++++++++++ mixtape-tools/src/sqlite/configured.rs | 431 +++++++++++++++++++++++++ mixtape-tools/src/sqlite/error.rs | 16 + mixtape-tools/src/sqlite/mod.rs | 207 ++++++++++++ mixtape-tools/src/sqlite/query/mod.rs | 8 +- mixtape-tools/src/sqlite/sql_parser.rs | 306 ++++++++++++++++++ 8 files changed, 1352 insertions(+), 5 deletions(-) create mode 100644 mixtape-tools/src/sqlite/config.rs create mode 100644 mixtape-tools/src/sqlite/configured.rs create mode 100644 mixtape-tools/src/sqlite/sql_parser.rs diff --git a/Cargo.toml b/Cargo.toml index 70686cc..9a9e762 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -59,6 +59,7 @@ aws-smithy-runtime-api = { version = "1.7", features = ["client", "http-1x"] } # Database rusqlite = { version = "0", features = ["bundled"] } +sqlparser = "0" # Utilities base64 = "0.22" diff --git a/mixtape-tools/Cargo.toml b/mixtape-tools/Cargo.toml index 299666f..fb08f01 100644 --- a/mixtape-tools/Cargo.toml +++ b/mixtape-tools/Cargo.toml @@ -20,7 +20,7 @@ edit = [] search = [] fetch = [] aws = [] -sqlite = ["dep:rusqlite", "dep:base64", "dep:sha2", "dep:hex"] +sqlite = ["dep:rusqlite", "dep:base64", "dep:sha2", "dep:hex", "dep:sqlparser"] [dependencies] mixtape-core.workspace = true @@ -67,6 +67,7 @@ rusqlite = { workspace = true, optional = true } base64 = { workspace = true, optional = true } sha2 = { workspace = true, optional = true } hex = { workspace = true, optional = true } +sqlparser = { workspace = true, optional = true } [target.'cfg(unix)'.dependencies] nix.workspace = true diff --git a/mixtape-tools/src/sqlite/config.rs b/mixtape-tools/src/sqlite/config.rs new file mode 100644 index 0000000..e30844f --- /dev/null +++ b/mixtape-tools/src/sqlite/config.rs @@ -0,0 +1,385 @@ +//! Configuration types for SQLite tool permissions +//! +//! This module provides configuration for restricting SQLite tool access +//! to specific databases and tables. + +use std::collections::HashSet; +use thiserror::Error; + +/// Errors that can occur when building a [`SqliteConfig`]. +#[derive(Debug, Error)] +pub enum ConfigError { + /// An empty deny list was specified, which allows all tables. + /// Use [`TablePermissionMode::AllowAll`] instead. + #[error( + "empty deny list for {operation} permissions allows all tables - use AllowAll or omit" + )] + EmptyDenyList { + /// Whether this was for "read" or "write" permissions + operation: &'static str, + }, +} + +/// Table permission mode - either allow specific tables or deny specific tables +#[derive(Debug, Clone, Default)] +pub enum TablePermissionMode { + /// Allow all tables (default behavior) + #[default] + AllowAll, + /// Only allow access to specified tables + AllowList(HashSet), + /// Deny access to specified tables, allow all others + DenyList(HashSet), +} + +impl TablePermissionMode { + /// Check if access to a table is allowed + pub fn is_allowed(&self, table: &str) -> bool { + match self { + TablePermissionMode::AllowAll => true, + TablePermissionMode::AllowList(allowed) => allowed.contains(table), + TablePermissionMode::DenyList(denied) => !denied.contains(table), + } + } + + /// Create an allow-list permission + pub fn allow(tables: I) -> Self + where + I: IntoIterator, + S: Into, + { + TablePermissionMode::AllowList(tables.into_iter().map(Into::into).collect()) + } + + /// Create a deny-list permission. + /// + /// Note: Empty deny lists are rejected at [`SqliteConfigBuilder::build()`] time + /// since they allow all tables, which is confusing. Use [`TablePermissionMode::AllowAll`] instead. + pub fn deny(tables: I) -> Self + where + I: IntoIterator, + S: Into, + { + TablePermissionMode::DenyList(tables.into_iter().map(Into::into).collect()) + } + + /// Returns true if this is an empty deny list (which allows everything). + fn is_empty_deny_list(&self) -> bool { + matches!(self, TablePermissionMode::DenyList(tables) if tables.is_empty()) + } +} + +/// Table permissions for read and write operations +#[derive(Debug, Clone, Default)] +pub struct TablePermissions { + /// Tables allowed/denied for read operations (SELECT) + pub read: TablePermissionMode, + + /// Tables allowed/denied for write operations (INSERT, UPDATE, DELETE) + pub write: TablePermissionMode, +} + +impl TablePermissions { + /// Create permissions that allow all operations + pub fn allow_all() -> Self { + Self::default() + } + + /// Create read-only permissions for specific tables + pub fn read_only(tables: I) -> Self + where + I: IntoIterator, + S: Into, + { + Self { + read: TablePermissionMode::AllowList(tables.into_iter().map(Into::into).collect()), + write: TablePermissionMode::AllowList(HashSet::new()), // Empty allow list denies all + } + } +} + +/// Configuration for SQLite tools with permission constraints +#[derive(Debug, Clone, Default)] +pub struct SqliteConfig { + /// Restrict tool to a specific database path. + /// When set, the tool will ignore any db_path in the input and always use this path. + pub db_path: Option, + + /// Table-level permissions + pub table_permissions: TablePermissions, +} + +impl SqliteConfig { + /// Create a new configuration with no restrictions + pub fn new() -> Self { + Self::default() + } + + /// Create a builder for constructing configuration + pub fn builder() -> SqliteConfigBuilder { + SqliteConfigBuilder::default() + } + + /// Check if a table is allowed for read access + pub fn can_read(&self, table: &str) -> bool { + self.table_permissions.read.is_allowed(table) + } + + /// Check if a table is allowed for write access + pub fn can_write(&self, table: &str) -> bool { + self.table_permissions.write.is_allowed(table) + } + + /// Get the effective database path (configured path takes precedence) + pub fn effective_db_path(&self, input_path: Option) -> Option { + self.db_path.clone().or(input_path) + } +} + +/// Builder for SqliteConfig +#[derive(Debug, Clone, Default)] +pub struct SqliteConfigBuilder { + db_path: Option, + table_permissions: TablePermissions, +} + +impl SqliteConfigBuilder { + /// Restrict to a specific database path + pub fn db_path(mut self, path: impl Into) -> Self { + self.db_path = Some(path.into()); + self + } + + /// Set read permissions + pub fn read_tables(mut self, mode: TablePermissionMode) -> Self { + self.table_permissions.read = mode; + self + } + + /// Set write permissions + pub fn write_tables(mut self, mode: TablePermissionMode) -> Self { + self.table_permissions.write = mode; + self + } + + /// Allow only specific tables for reading + pub fn allow_read(mut self, tables: I) -> Self + where + I: IntoIterator, + S: Into, + { + self.table_permissions.read = TablePermissionMode::allow(tables); + self + } + + /// Deny specific tables for reading + pub fn deny_read(mut self, tables: I) -> Self + where + I: IntoIterator, + S: Into, + { + self.table_permissions.read = TablePermissionMode::deny(tables); + self + } + + /// Allow only specific tables for writing + pub fn allow_write(mut self, tables: I) -> Self + where + I: IntoIterator, + S: Into, + { + self.table_permissions.write = TablePermissionMode::allow(tables); + self + } + + /// Deny specific tables for writing + pub fn deny_write(mut self, tables: I) -> Self + where + I: IntoIterator, + S: Into, + { + self.table_permissions.write = TablePermissionMode::deny(tables); + self + } + + /// Deny all write operations + pub fn read_only(mut self) -> Self { + // An empty allow list denies everything + self.table_permissions.write = TablePermissionMode::AllowList(HashSet::new()); + self + } + + /// Build the configuration. + /// + /// # Errors + /// + /// Returns [`ConfigError::EmptyDenyList`] if either read or write permissions + /// use an empty deny list, since that allows all tables (use `AllowAll` instead). + pub fn build(self) -> Result { + // Validate: empty deny lists are confusing since they allow everything + if self.table_permissions.read.is_empty_deny_list() { + return Err(ConfigError::EmptyDenyList { operation: "read" }); + } + if self.table_permissions.write.is_empty_deny_list() { + return Err(ConfigError::EmptyDenyList { operation: "write" }); + } + + Ok(SqliteConfig { + db_path: self.db_path, + table_permissions: self.table_permissions, + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_table_permission_mode_allow_all() { + let mode = TablePermissionMode::AllowAll; + assert!(mode.is_allowed("users")); + assert!(mode.is_allowed("secrets")); + assert!(mode.is_allowed("anything")); + } + + #[test] + fn test_table_permission_mode_allow_list() { + let mode = TablePermissionMode::allow(["users", "orders"]); + assert!(mode.is_allowed("users")); + assert!(mode.is_allowed("orders")); + assert!(!mode.is_allowed("secrets")); + assert!(!mode.is_allowed("admin")); + } + + #[test] + fn test_table_permission_mode_deny_list() { + let mode = TablePermissionMode::deny(["secrets", "admin_logs"]); + assert!(mode.is_allowed("users")); + assert!(mode.is_allowed("orders")); + assert!(!mode.is_allowed("secrets")); + assert!(!mode.is_allowed("admin_logs")); + } + + #[test] + fn test_empty_deny_list_returns_error_read() { + let result = SqliteConfig::builder() + .deny_read(Vec::::new()) + .build(); + + assert!(result.is_err()); + let err = result.unwrap_err(); + assert!(err.to_string().contains("empty deny list")); + assert!(err.to_string().contains("read")); + } + + #[test] + fn test_empty_deny_list_returns_error_write() { + let result = SqliteConfig::builder() + .deny_write(Vec::::new()) + .build(); + + assert!(result.is_err()); + let err = result.unwrap_err(); + assert!(err.to_string().contains("empty deny list")); + assert!(err.to_string().contains("write")); + } + + #[test] + fn test_sqlite_config_builder_db_path() { + let config = SqliteConfig::builder() + .db_path("/data/app.db") + .build() + .unwrap(); + + assert_eq!(config.db_path, Some("/data/app.db".to_string())); + } + + #[test] + fn test_sqlite_config_builder_allow_read() { + let config = SqliteConfig::builder() + .allow_read(["users", "products"]) + .build() + .unwrap(); + + assert!(config.can_read("users")); + assert!(config.can_read("products")); + assert!(!config.can_read("secrets")); + } + + #[test] + fn test_sqlite_config_builder_deny_write() { + let config = SqliteConfig::builder() + .deny_write(["audit_log"]) + .build() + .unwrap(); + + assert!(config.can_write("users")); + assert!(!config.can_write("audit_log")); + } + + #[test] + fn test_sqlite_config_builder_read_only() { + let config = SqliteConfig::builder().read_only().build().unwrap(); + + assert!(config.can_read("users")); + assert!(!config.can_write("users")); + assert!(!config.can_write("anything")); + } + + #[test] + fn test_effective_db_path_config_takes_precedence() { + let config = SqliteConfig::builder() + .db_path("/configured/path.db") + .build() + .unwrap(); + + // Configured path should override input path + assert_eq!( + config.effective_db_path(Some("/input/path.db".to_string())), + Some("/configured/path.db".to_string()) + ); + + // Configured path should be used when input is None + assert_eq!( + config.effective_db_path(None), + Some("/configured/path.db".to_string()) + ); + } + + #[test] + fn test_effective_db_path_falls_back_to_input() { + let config = SqliteConfig::new(); + + // Without configured path, input path should be used + assert_eq!( + config.effective_db_path(Some("/input/path.db".to_string())), + Some("/input/path.db".to_string()) + ); + + // Without either, should be None + assert_eq!(config.effective_db_path(None), None); + } + + #[test] + fn test_combined_permissions() { + let config = SqliteConfig::builder() + .db_path("/data/app.db") + .allow_read(["users", "products", "orders"]) + .allow_write(["orders"]) + .build() + .unwrap(); + + // Read permissions + assert!(config.can_read("users")); + assert!(config.can_read("products")); + assert!(config.can_read("orders")); + assert!(!config.can_read("secrets")); + + // Write permissions + assert!(!config.can_write("users")); + assert!(!config.can_write("products")); + assert!(config.can_write("orders")); + assert!(!config.can_write("secrets")); + } +} diff --git a/mixtape-tools/src/sqlite/configured.rs b/mixtape-tools/src/sqlite/configured.rs new file mode 100644 index 0000000..b45f0fe --- /dev/null +++ b/mixtape-tools/src/sqlite/configured.rs @@ -0,0 +1,431 @@ +//! Configured tool wrappers with permission support +//! +//! This module provides wrapper tools that enforce database path and table +//! permission restrictions configured at tool creation time. + +use crate::prelude::*; +use crate::sqlite::config::SqliteConfig; +use crate::sqlite::error::SqliteToolError; +use crate::sqlite::query::{ + BulkInsertInput, BulkInsertTool, ReadQueryInput, ReadQueryTool, SchemaQueryInput, + SchemaQueryTool, WriteQueryInput, WriteQueryTool, +}; +use crate::sqlite::sql_parser::extract_table_operations; +use std::sync::Arc; + +/// Validate query permissions against the configuration +fn validate_query(config: &SqliteConfig, sql: &str) -> Result<(), SqliteToolError> { + let ops = + extract_table_operations(sql).map_err(|e| SqliteToolError::InvalidQuery(e.to_string()))?; + + // Check read permissions + for table in &ops.read { + if !config.can_read(table) { + return Err(SqliteToolError::PermissionDenied { + operation: "read".to_string(), + table: table.clone(), + }); + } + } + + // Check write permissions + for table in &ops.write { + if !config.can_write(table) { + return Err(SqliteToolError::PermissionDenied { + operation: "write".to_string(), + table: table.clone(), + }); + } + } + + Ok(()) +} + +/// A ReadQueryTool with permission configuration +pub struct ConfiguredReadQueryTool { + config: Arc, + inner: ReadQueryTool, +} + +impl ConfiguredReadQueryTool { + /// Create a new configured read query tool + pub fn new(config: SqliteConfig) -> Self { + Self { + config: Arc::new(config), + inner: ReadQueryTool, + } + } + + /// Create with a shared config (useful when multiple tools share permissions) + pub fn with_shared_config(config: Arc) -> Self { + Self { + config, + inner: ReadQueryTool, + } + } +} + +impl Tool for ConfiguredReadQueryTool { + type Input = ReadQueryInput; + + fn name(&self) -> &str { + self.inner.name() + } + + fn description(&self) -> &str { + self.inner.description() + } + + async fn execute(&self, mut input: Self::Input) -> Result { + input.db_path = self.config.effective_db_path(input.db_path); + validate_query(&self.config, &input.query)?; + self.inner.execute(input).await + } +} + +/// A WriteQueryTool with permission configuration +pub struct ConfiguredWriteQueryTool { + config: Arc, + inner: WriteQueryTool, +} + +impl ConfiguredWriteQueryTool { + /// Create a new configured write query tool + pub fn new(config: SqliteConfig) -> Self { + Self { + config: Arc::new(config), + inner: WriteQueryTool, + } + } + + /// Create with a shared config + pub fn with_shared_config(config: Arc) -> Self { + Self { + config, + inner: WriteQueryTool, + } + } +} + +impl Tool for ConfiguredWriteQueryTool { + type Input = WriteQueryInput; + + fn name(&self) -> &str { + self.inner.name() + } + + fn description(&self) -> &str { + self.inner.description() + } + + async fn execute(&self, mut input: Self::Input) -> Result { + input.db_path = self.config.effective_db_path(input.db_path); + validate_query(&self.config, &input.query)?; + self.inner.execute(input).await + } +} + +/// A SchemaQueryTool with permission configuration +pub struct ConfiguredSchemaQueryTool { + config: Arc, + inner: SchemaQueryTool, +} + +impl ConfiguredSchemaQueryTool { + /// Create a new configured schema query tool + pub fn new(config: SqliteConfig) -> Self { + Self { + config: Arc::new(config), + inner: SchemaQueryTool, + } + } + + /// Create with a shared config + pub fn with_shared_config(config: Arc) -> Self { + Self { + config, + inner: SchemaQueryTool, + } + } +} + +impl Tool for ConfiguredSchemaQueryTool { + type Input = SchemaQueryInput; + + fn name(&self) -> &str { + self.inner.name() + } + + fn description(&self) -> &str { + self.inner.description() + } + + async fn execute(&self, mut input: Self::Input) -> Result { + input.db_path = self.config.effective_db_path(input.db_path); + validate_query(&self.config, &input.query)?; + self.inner.execute(input).await + } +} + +/// A BulkInsertTool with permission configuration +pub struct ConfiguredBulkInsertTool { + config: Arc, + inner: BulkInsertTool, +} + +impl ConfiguredBulkInsertTool { + /// Create a new configured bulk insert tool + pub fn new(config: SqliteConfig) -> Self { + Self { + config: Arc::new(config), + inner: BulkInsertTool, + } + } + + /// Create with a shared config + pub fn with_shared_config(config: Arc) -> Self { + Self { + config, + inner: BulkInsertTool, + } + } +} + +impl Tool for ConfiguredBulkInsertTool { + type Input = BulkInsertInput; + + fn name(&self) -> &str { + self.inner.name() + } + + fn description(&self) -> &str { + self.inner.description() + } + + async fn execute(&self, mut input: Self::Input) -> Result { + input.db_path = self.config.effective_db_path(input.db_path); + if !self.config.can_write(&input.table) { + return Err(SqliteToolError::PermissionDenied { + operation: "write".to_string(), + table: input.table.clone(), + } + .into()); + } + self.inner.execute(input).await + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::sqlite::test_utils::TestDatabase; + + #[tokio::test] + async fn test_configured_read_tool_allows_permitted_table() { + let db = TestDatabase::with_schema( + "CREATE TABLE users (id INTEGER, name TEXT); + INSERT INTO users VALUES (1, 'Alice');", + ) + .await; + + let config = SqliteConfig::builder() + .db_path(db.key()) + .allow_read(["users"]) + .build() + .unwrap(); + + let tool = ConfiguredReadQueryTool::new(config); + let result = tool + .execute(ReadQueryInput::new("SELECT * FROM users")) + .await; + + assert!(result.is_ok()); + } + + #[tokio::test] + async fn test_configured_read_tool_denies_unpermitted_table() { + let db = TestDatabase::with_schema( + "CREATE TABLE secrets (id INTEGER, data TEXT); + INSERT INTO secrets VALUES (1, 'secret data');", + ) + .await; + + let config = SqliteConfig::builder() + .db_path(db.key()) + .allow_read(["users"]) // secrets not in allow list + .build() + .unwrap(); + + let tool = ConfiguredReadQueryTool::new(config); + let result = tool + .execute(ReadQueryInput::new("SELECT * FROM secrets")) + .await; + + assert!(result.is_err()); + let err_msg = result.unwrap_err().to_string(); + assert!(err_msg.contains("Permission denied")); + assert!(err_msg.contains("secrets")); + } + + #[tokio::test] + async fn test_configured_write_tool_allows_permitted_table() { + let db = TestDatabase::with_schema("CREATE TABLE orders (id INTEGER, amount REAL);").await; + + let config = SqliteConfig::builder() + .db_path(db.key()) + .allow_write(["orders"]) + .build() + .unwrap(); + + let tool = ConfiguredWriteQueryTool::new(config); + let result = tool + .execute(WriteQueryInput { + query: "INSERT INTO orders VALUES (1, 99.99)".to_string(), + params: vec![], + db_path: None, + }) + .await; + + assert!(result.is_ok()); + } + + #[tokio::test] + async fn test_configured_write_tool_denies_unpermitted_table() { + let db = TestDatabase::with_schema("CREATE TABLE audit_log (id INTEGER, msg TEXT);").await; + + let config = SqliteConfig::builder() + .db_path(db.key()) + .deny_write(["audit_log"]) + .build() + .unwrap(); + + let tool = ConfiguredWriteQueryTool::new(config); + let result = tool + .execute(WriteQueryInput { + query: "INSERT INTO audit_log VALUES (1, 'hacked')".to_string(), + params: vec![], + db_path: None, + }) + .await; + + assert!(result.is_err()); + let err_msg = result.unwrap_err().to_string(); + assert!(err_msg.contains("Permission denied")); + assert!(err_msg.contains("audit_log")); + } + + #[tokio::test] + async fn test_configured_tool_uses_configured_db_path() { + let db = TestDatabase::with_schema("CREATE TABLE test (id INTEGER);").await; + + let config = SqliteConfig::builder().db_path(db.key()).build().unwrap(); + + let tool = ConfiguredReadQueryTool::new(config); + + // Input provides a different path, but configured path should be used + let result = tool + .execute(ReadQueryInput::new("SELECT * FROM test").db_path("/nonexistent/path.db")) + .await; + + // Should succeed because it uses the configured path, not the input path + assert!(result.is_ok()); + } + + #[tokio::test] + async fn test_configured_bulk_insert_checks_permissions() { + let db = TestDatabase::with_schema("CREATE TABLE users (id INTEGER, name TEXT);").await; + + let config = SqliteConfig::builder() + .db_path(db.key()) + .deny_write(["users"]) + .build() + .unwrap(); + + let tool = ConfiguredBulkInsertTool::new(config); + let mut record = serde_json::Map::new(); + record.insert("id".to_string(), serde_json::json!(1)); + record.insert("name".to_string(), serde_json::json!("Alice")); + + let result = tool + .execute(BulkInsertInput { + table: "users".to_string(), + data: vec![record], + batch_size: 1000, + db_path: None, + }) + .await; + + assert!(result.is_err()); + let err_msg = result.unwrap_err().to_string(); + assert!(err_msg.contains("Permission denied")); + } + + #[tokio::test] + async fn test_read_only_config_blocks_all_writes() { + let db = TestDatabase::with_schema("CREATE TABLE data (id INTEGER);").await; + + let config = SqliteConfig::builder() + .db_path(db.key()) + .read_only() + .build() + .unwrap(); + + let tool = ConfiguredWriteQueryTool::new(config); + let result = tool + .execute(WriteQueryInput { + query: "INSERT INTO data VALUES (1)".to_string(), + params: vec![], + db_path: None, + }) + .await; + + assert!(result.is_err()); + } + + #[tokio::test] + async fn test_shared_config_between_tools() { + let db = TestDatabase::with_schema( + "CREATE TABLE users (id INTEGER, name TEXT); + CREATE TABLE orders (id INTEGER, user_id INTEGER);", + ) + .await; + + let config = Arc::new( + SqliteConfig::builder() + .db_path(db.key()) + .allow_read(["users", "orders"]) + .allow_write(["orders"]) + .build() + .unwrap(), + ); + + let read_tool = ConfiguredReadQueryTool::with_shared_config(config.clone()); + let write_tool = ConfiguredWriteQueryTool::with_shared_config(config.clone()); + + // Read from users should work + assert!(read_tool + .execute(ReadQueryInput::new("SELECT * FROM users")) + .await + .is_ok()); + + // Write to users should fail + assert!(write_tool + .execute(WriteQueryInput { + query: "INSERT INTO users VALUES (1, 'test')".to_string(), + params: vec![], + db_path: None, + }) + .await + .is_err()); + + // Write to orders should work + assert!(write_tool + .execute(WriteQueryInput { + query: "INSERT INTO orders VALUES (1, 1)".to_string(), + params: vec![], + db_path: None, + }) + .await + .is_ok()); + } +} diff --git a/mixtape-tools/src/sqlite/error.rs b/mixtape-tools/src/sqlite/error.rs index bb41998..0ab0e4b 100644 --- a/mixtape-tools/src/sqlite/error.rs +++ b/mixtape-tools/src/sqlite/error.rs @@ -35,6 +35,10 @@ pub enum SqliteToolError { #[error("Transaction error: {0}")] TransactionError(String), + /// Permission denied for table operation + #[error("Permission denied: cannot {operation} table '{table}'")] + PermissionDenied { operation: String, table: String }, + /// Path validation or filesystem error #[error("Path error: {0}")] PathError(String), @@ -131,6 +135,18 @@ mod tests { assert_eq!(err.to_string(), "Transaction error: no active transaction"); } + #[test] + fn test_permission_denied_display() { + let err = SqliteToolError::PermissionDenied { + operation: "write".to_string(), + table: "secrets".to_string(), + }; + assert_eq!( + err.to_string(), + "Permission denied: cannot write table 'secrets'" + ); + } + #[test] fn test_path_error_display() { let err = SqliteToolError::PathError("invalid path".to_string()); diff --git a/mixtape-tools/src/sqlite/mod.rs b/mixtape-tools/src/sqlite/mod.rs index f554662..cdd53c4 100644 --- a/mixtape-tools/src/sqlite/mod.rs +++ b/mixtape-tools/src/sqlite/mod.rs @@ -85,6 +85,57 @@ //! .await?; //! ``` //! +//! # Fine-Grained Permissions +//! +//! For tighter control, use configured tools that restrict access to specific +//! databases and tables. These tools validate SQL queries before execution and +//! enforce table-level read/write permissions. +//! +//! ## Lock Tools to a Specific Database +//! +//! ```rust,ignore +//! use mixtape_tools::sqlite; +//! +//! // Tools will only access this database, ignoring any db_path in input +//! let agent = Agent::builder() +//! .add_tools(sqlite::tools_for_database("/data/app.db")) +//! .build() +//! .await?; +//! ``` +//! +//! ## Read-Only Access to Specific Tables +//! +//! ```rust,ignore +//! use mixtape_tools::sqlite; +//! +//! // Can only SELECT from these tables, all writes blocked +//! let agent = Agent::builder() +//! .add_tools(sqlite::read_only_tools_for_tables( +//! "/data/app.db", +//! ["users", "products", "orders"] +//! )) +//! .build() +//! .await?; +//! ``` +//! +//! ## Custom Table Permissions +//! +//! ```rust,ignore +//! use mixtape_tools::sqlite::{SqliteConfig, tools_with_config}; +//! +//! let config = SqliteConfig::builder() +//! .db_path("/data/app.db") +//! .allow_read(["users", "products", "orders", "analytics"]) +//! .allow_write(["analytics"]) // Can only write to analytics +//! .deny_read(["secrets"]) // Block access to secrets table +//! .build()?; +//! +//! let agent = Agent::builder() +//! .add_tools(tools_with_config(config)) +//! .build() +//! .await?; +//! ``` +//! //! # Tool Categories //! //! ## Database Management (Safe) @@ -122,12 +173,15 @@ //! - `sqlite_export_migrations` - Export migrations for transfer (Safe) //! - `sqlite_import_migrations` - Import migrations as pending (Destructive) +pub mod config; +pub mod configured; pub mod database; pub mod error; pub mod maintenance; pub mod manager; pub mod migration; pub mod query; +mod sql_parser; pub mod table; #[cfg(test)] pub mod test_utils; @@ -135,6 +189,13 @@ pub mod transaction; pub mod types; // Re-export commonly used items +pub use config::{ + ConfigError, SqliteConfig, SqliteConfigBuilder, TablePermissionMode, TablePermissions, +}; +pub use configured::{ + ConfiguredBulkInsertTool, ConfiguredReadQueryTool, ConfiguredSchemaQueryTool, + ConfiguredWriteQueryTool, +}; pub use database::{CloseDatabaseTool, DatabaseInfoTool, ListDatabasesTool, OpenDatabaseTool}; pub use error::SqliteToolError; pub use maintenance::{BackupDatabaseTool, ExportSchemaTool, VacuumDatabaseTool}; @@ -211,6 +272,118 @@ pub fn all_tools() -> Vec> { tools } +// ============================================================================= +// Configured tool factory functions +// ============================================================================= + +use std::sync::Arc; + +/// Create a set of SQLite tools restricted to a specific database. +/// +/// The returned tools will only access the specified database, ignoring +/// any `db_path` parameter in tool inputs. +/// +/// # Example +/// +/// ```rust,ignore +/// use mixtape_tools::sqlite; +/// +/// let agent = Agent::builder() +/// .add_tools(sqlite::tools_for_database("/data/app.db")) +/// .build() +/// .await?; +/// ``` +pub fn tools_for_database(db_path: impl Into) -> Vec> { + let config = Arc::new( + SqliteConfig::builder() + .db_path(db_path) + .build() + .expect("tools_for_database: config build should never fail"), + ); + + vec![ + box_tool(ConfiguredReadQueryTool::with_shared_config(config.clone())), + box_tool(ConfiguredWriteQueryTool::with_shared_config(config.clone())), + box_tool(ConfiguredSchemaQueryTool::with_shared_config( + config.clone(), + )), + box_tool(ConfiguredBulkInsertTool::with_shared_config(config)), + ] +} + +/// Create read-only tools for specific tables in a specific database. +/// +/// The returned tools can only read from the specified tables. +/// All write operations will be denied. +/// +/// # Example +/// +/// ```rust,ignore +/// use mixtape_tools::sqlite; +/// +/// let agent = Agent::builder() +/// .add_tools(sqlite::read_only_tools_for_tables( +/// "/data/app.db", +/// ["users", "products", "orders"] +/// )) +/// .build() +/// .await?; +/// ``` +pub fn read_only_tools_for_tables( + db_path: impl Into, + tables: I, +) -> Vec> +where + I: IntoIterator, + S: Into, +{ + let config = Arc::new( + SqliteConfig::builder() + .db_path(db_path) + .allow_read(tables) + .read_only() + .build() + .expect("read_only_tools_for_tables: config build should never fail"), + ); + + vec![box_tool(ConfiguredReadQueryTool::with_shared_config( + config, + ))] +} + +/// Create tools with custom configuration. +/// +/// This provides full control over database path and table permissions. +/// +/// # Example +/// +/// ```rust,ignore +/// use mixtape_tools::sqlite::{SqliteConfig, tools_with_config}; +/// +/// let config = SqliteConfig::builder() +/// .db_path("/data/app.db") +/// .allow_read(["users", "products", "orders", "analytics"]) +/// .allow_write(["analytics"]) +/// .build(); +/// +/// let agent = Agent::builder() +/// .add_tools(tools_with_config(config)) +/// .build() +/// .await?; +/// ``` +pub fn tools_with_config(config: SqliteConfig) -> Vec> { + let config = Arc::new(config); + + vec![ + box_tool(ConfiguredReadQueryTool::with_shared_config(config.clone())), + box_tool(ConfiguredWriteQueryTool::with_shared_config(config.clone())), + box_tool(ConfiguredSchemaQueryTool::with_shared_config( + config.clone(), + )), + box_tool(ConfiguredBulkInsertTool::with_shared_config(config)), + ] +} + #[cfg(test)] mod tests { use super::*; @@ -321,4 +494,38 @@ mod tests { "Duplicate tool names found" ); } + + #[test] + fn test_tools_for_database() { + let tools = tools_for_database("/test/path.db"); + assert_eq!(tools.len(), 4); + + let names: Vec<&str> = tools.iter().map(|t| t.name()).collect(); + assert!(names.contains(&"sqlite_read_query")); + assert!(names.contains(&"sqlite_write_query")); + assert!(names.contains(&"sqlite_schema_query")); + assert!(names.contains(&"sqlite_bulk_insert")); + } + + #[test] + fn test_read_only_tools_for_tables() { + let tools = read_only_tools_for_tables("/test/path.db", ["users", "orders"]); + assert_eq!(tools.len(), 1); + + let names: Vec<&str> = tools.iter().map(|t| t.name()).collect(); + assert!(names.contains(&"sqlite_read_query")); + } + + #[test] + fn test_tools_with_config() { + let config = SqliteConfig::builder() + .db_path("/test/path.db") + .allow_read(["users"]) + .allow_write(["orders"]) + .build() + .unwrap(); + + let tools = tools_with_config(config); + assert_eq!(tools.len(), 4); + } } diff --git a/mixtape-tools/src/sqlite/query/mod.rs b/mixtape-tools/src/sqlite/query/mod.rs index c4564ef..09d19d8 100644 --- a/mixtape-tools/src/sqlite/query/mod.rs +++ b/mixtape-tools/src/sqlite/query/mod.rs @@ -5,7 +5,7 @@ mod read; mod schema; mod write; -pub use bulk_insert::BulkInsertTool; -pub use read::ReadQueryTool; -pub use schema::SchemaQueryTool; -pub use write::WriteQueryTool; +pub use bulk_insert::{BulkInsertInput, BulkInsertTool}; +pub use read::{ReadQueryInput, ReadQueryTool}; +pub use schema::{SchemaQueryInput, SchemaQueryTool}; +pub use write::{WriteQueryInput, WriteQueryTool}; diff --git a/mixtape-tools/src/sqlite/sql_parser.rs b/mixtape-tools/src/sqlite/sql_parser.rs new file mode 100644 index 0000000..7a882b8 --- /dev/null +++ b/mixtape-tools/src/sqlite/sql_parser.rs @@ -0,0 +1,306 @@ +//! SQL parsing utilities for table extraction +//! +//! This module provides functions to extract table names from SQL queries +//! and categorize them as read or write operations. + +use sqlparser::ast::{ + Delete, Expr, FromTable, FunctionArg, FunctionArgExpr, Insert, Query, Select, SetExpr, + Statement, TableFactor, TableObject, TableWithJoins, Update, UpdateTableFromKind, +}; +use sqlparser::dialect::SQLiteDialect; +use sqlparser::parser::Parser; +use std::collections::HashSet; + +/// Categorize tables by operation type +#[derive(Debug, Default)] +pub struct TableOperations { + /// Tables being read from (SELECT, subqueries, JOINs) + pub read: HashSet, + /// Tables being written to (INSERT, UPDATE, DELETE target) + pub write: HashSet, +} + +/// Extract tables categorized by read/write operation +pub fn extract_table_operations(sql: &str) -> Result { + let dialect = SQLiteDialect {}; + let statements = + Parser::parse_sql(&dialect, sql).map_err(|e| format!("Failed to parse SQL: {}", e))?; + + let mut ops = TableOperations::default(); + for statement in statements { + categorize_tables(&statement, &mut ops); + } + Ok(ops) +} + +fn categorize_tables(stmt: &Statement, ops: &mut TableOperations) { + match stmt { + Statement::Query(query) => extract_tables_from_query(query, &mut ops.read), + Statement::Insert(Insert { table, source, .. }) => { + if let TableObject::TableName(name) = table { + ops.write.insert(name.to_string()); + } + if let Some(src) = source { + extract_tables_from_query(src, &mut ops.read); + } + } + Statement::Update(Update { + table, + from, + selection, + .. + }) => { + // The target table is being written + if let TableFactor::Table { name, .. } = &table.relation { + ops.write.insert(name.to_string()); + } + // JOINs in UPDATE are reads + for join in &table.joins { + if let TableFactor::Table { name, .. } = &join.relation { + ops.read.insert(name.to_string()); + } + } + // FROM clause tables are reads + if let Some(from_kind) = from { + let from_tables = match from_kind { + UpdateTableFromKind::BeforeSet(tables) + | UpdateTableFromKind::AfterSet(tables) => tables, + }; + for twj in from_tables { + extract_tables_from_table_with_joins(twj, &mut ops.read); + } + } + if let Some(expr) = selection { + extract_tables_from_expr(expr, &mut ops.read); + } + } + Statement::Delete(Delete { + from, selection, .. + }) => { + match from { + FromTable::WithFromKeyword(tables) | FromTable::WithoutKeyword(tables) => { + for twj in tables { + if let TableFactor::Table { name, .. } = &twj.relation { + ops.write.insert(name.to_string()); + } + } + } + } + if let Some(expr) = selection { + extract_tables_from_expr(expr, &mut ops.read); + } + } + _ => {} + } +} + +fn extract_tables_from_query(query: &Query, tables: &mut HashSet) { + // Handle CTEs (WITH clause) + if let Some(with) = &query.with { + for cte in &with.cte_tables { + extract_tables_from_query(&cte.query, tables); + } + } + extract_tables_from_set_expr(&query.body, tables); +} + +fn extract_tables_from_set_expr(body: &SetExpr, tables: &mut HashSet) { + match body { + SetExpr::Select(select) => extract_tables_from_select(select, tables), + SetExpr::Query(query) => extract_tables_from_query(query, tables), + SetExpr::SetOperation { left, right, .. } => { + extract_tables_from_set_expr(left, tables); + extract_tables_from_set_expr(right, tables); + } + _ => {} + } +} + +fn extract_tables_from_select(select: &Select, tables: &mut HashSet) { + for twj in &select.from { + extract_tables_from_table_with_joins(twj, tables); + } + if let Some(expr) = &select.selection { + extract_tables_from_expr(expr, tables); + } +} + +fn extract_tables_from_table_with_joins(twj: &TableWithJoins, tables: &mut HashSet) { + extract_tables_from_table_factor(&twj.relation, tables); + for join in &twj.joins { + extract_tables_from_table_factor(&join.relation, tables); + } +} + +fn extract_tables_from_table_factor(factor: &TableFactor, tables: &mut HashSet) { + match factor { + TableFactor::Table { name, .. } => { + tables.insert(name.to_string()); + } + TableFactor::Derived { subquery, .. } => { + extract_tables_from_query(subquery, tables); + } + TableFactor::NestedJoin { + table_with_joins, .. + } => { + extract_tables_from_table_with_joins(table_with_joins, tables); + } + _ => {} + } +} + +fn extract_tables_from_expr(expr: &Expr, tables: &mut HashSet) { + match expr { + Expr::Subquery(query) => extract_tables_from_query(query, tables), + Expr::InSubquery { subquery, .. } => extract_tables_from_query(subquery, tables), + Expr::Exists { subquery, .. } => extract_tables_from_query(subquery, tables), + Expr::BinaryOp { left, right, .. } => { + extract_tables_from_expr(left, tables); + extract_tables_from_expr(right, tables); + } + Expr::UnaryOp { expr, .. } => extract_tables_from_expr(expr, tables), + Expr::Between { + expr, low, high, .. + } => { + extract_tables_from_expr(expr, tables); + extract_tables_from_expr(low, tables); + extract_tables_from_expr(high, tables); + } + Expr::Case { + operand, + conditions, + else_result, + .. + } => { + if let Some(op) = operand { + extract_tables_from_expr(op, tables); + } + for case_when in conditions { + extract_tables_from_expr(&case_when.condition, tables); + extract_tables_from_expr(&case_when.result, tables); + } + if let Some(else_r) = else_result { + extract_tables_from_expr(else_r, tables); + } + } + Expr::Nested(inner) => extract_tables_from_expr(inner, tables), + Expr::InList { expr, list, .. } => { + extract_tables_from_expr(expr, tables); + for item in list { + extract_tables_from_expr(item, tables); + } + } + Expr::Function(func) => { + if let sqlparser::ast::FunctionArguments::List(arg_list) = &func.args { + for arg in &arg_list.args { + if let FunctionArg::Unnamed(FunctionArgExpr::Expr(e)) + | FunctionArg::Named { + arg: FunctionArgExpr::Expr(e), + .. + } = arg + { + extract_tables_from_expr(e, tables); + } + } + } + } + _ => {} + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_extract_simple_select() { + let ops = extract_table_operations("SELECT * FROM users").unwrap(); + assert!(ops.read.contains("users")); + assert_eq!(ops.read.len(), 1); + assert!(ops.write.is_empty()); + } + + #[test] + fn test_extract_join() { + let ops = + extract_table_operations("SELECT * FROM users u JOIN orders o ON u.id = o.user_id") + .unwrap(); + assert!(ops.read.contains("users")); + assert!(ops.read.contains("orders")); + assert_eq!(ops.read.len(), 2); + assert!(ops.write.is_empty()); + } + + #[test] + fn test_extract_subquery() { + let ops = extract_table_operations( + "SELECT * FROM users WHERE id IN (SELECT user_id FROM active_sessions)", + ) + .unwrap(); + assert!(ops.read.contains("users")); + assert!(ops.read.contains("active_sessions")); + assert_eq!(ops.read.len(), 2); + } + + #[test] + fn test_extract_insert() { + let ops = extract_table_operations("INSERT INTO users (name) VALUES ('test')").unwrap(); + assert!(ops.write.contains("users")); + assert!(ops.read.is_empty()); + } + + #[test] + fn test_categorize_insert_select() { + let ops = extract_table_operations( + "INSERT INTO archive SELECT * FROM logs WHERE created_at < '2024-01-01'", + ) + .unwrap(); + assert!(ops.write.contains("archive")); + assert!(ops.read.contains("logs")); + } + + #[test] + fn test_extract_update() { + let ops = + extract_table_operations("UPDATE users SET status = 'active' WHERE id = 1").unwrap(); + assert!(ops.write.contains("users")); + } + + #[test] + fn test_extract_delete() { + let ops = extract_table_operations("DELETE FROM users WHERE id = 1").unwrap(); + assert!(ops.write.contains("users")); + } + + #[test] + fn test_extract_cte() { + let ops = extract_table_operations( + "WITH active AS (SELECT * FROM users WHERE active = 1) SELECT * FROM active JOIN orders ON active.id = orders.user_id" + ).unwrap(); + assert!(ops.read.contains("users")); + assert!(ops.read.contains("orders")); + } + + #[test] + fn test_extract_union() { + let ops = + extract_table_operations("SELECT id FROM users UNION SELECT id FROM admins").unwrap(); + assert!(ops.read.contains("users")); + assert!(ops.read.contains("admins")); + } + + #[test] + fn test_invalid_sql() { + let result = extract_table_operations("NOT VALID SQL AT ALL"); + assert!(result.is_err()); + } + + #[test] + fn test_multiple_statements() { + let ops = + extract_table_operations("SELECT * FROM users; INSERT INTO logs (msg) VALUES ('test')") + .unwrap(); + assert!(ops.read.contains("users")); + assert!(ops.write.contains("logs")); + } +}