From 09e6a7e6d8bc8a9d3ef067f6aa95cb961407b3ac Mon Sep 17 00:00:00 2001 From: Illia Polosukhin Date: Fri, 13 Feb 2026 17:03:58 -0800 Subject: [PATCH] fix: review fixes for libSQL backend (shared connections, panics, indexes) - Replace .expect() with proper error propagation in 3 call sites - Share Arc between backend and stores instead of single Connection - Add connect-per-operation pattern to LibSqlSecretsStore and LibSqlWasmToolStore - Wrap store() INSERT + SELECT-back in a transaction - Add ~22 missing indexes for parity with PostgreSQL schema - Add 18 leak_detection_patterns seed rows matching PostgreSQL V2 migration - Fix super:: import to use crate:: style - Gate mask_password_in_url behind #[cfg(feature = "postgres")] - Rewrite secrets store init with or_else chain for runtime backend selection Co-Authored-By: Claude Opus 4.6 --- src/cli/mcp.rs | 12 ++--- src/cli/tool.rs | 12 ++--- src/db/libsql_backend.rs | 31 +++++++----- src/db/libsql_migrations.rs | 69 ++++++++++++++++++++++++++ src/main.rs | 93 +++++++++++++++++------------------ src/secrets/store.rs | 51 ++++++++++++-------- src/setup/wizard.rs | 7 ++- src/tools/wasm/storage.rs | 96 ++++++++++++++++++++++--------------- 8 files changed, 233 insertions(+), 138 deletions(-) diff --git a/src/cli/mcp.rs b/src/cli/mcp.rs index f73659fc..e61b65de 100644 --- a/src/cli/mcp.rs +++ b/src/cli/mcp.rs @@ -530,11 +530,9 @@ async fn get_secrets_store() -> anyhow::Result anyhow::Result, user_id: String) -> anyho .unwrap_or(&default_path); let backend = if let Some(ref url) = config.database.libsql_url { - let token = config - .database - .libsql_auth_token - .as_ref() - .expect("LIBSQL_AUTH_TOKEN required when LIBSQL_URL is set"); + let token = config.database.libsql_auth_token.as_ref().ok_or_else(|| { + anyhow::anyhow!("LIBSQL_AUTH_TOKEN is required when LIBSQL_URL is set") + })?; LibSqlBackend::new_remote_replica(db_path, url, token.expose_secret()) .await .map_err(|e| anyhow::anyhow!("{}", e))? @@ -766,10 +764,8 @@ async fn auth_tool(name: String, dir: Option, user_id: String) -> anyho .await .map_err(|e| anyhow::anyhow!("{}", e))?; - let conn = backend.connect().map_err(|e| anyhow::anyhow!("{}", e))?; - Arc::new(crate::secrets::LibSqlSecretsStore::new( - conn, + backend.shared_db(), Arc::new(crypto), )) } diff --git a/src/db/libsql_backend.rs b/src/db/libsql_backend.rs index ccae6a45..9fe183ee 100644 --- a/src/db/libsql_backend.rs +++ b/src/db/libsql_backend.rs @@ -8,6 +8,7 @@ use std::collections::HashMap; use std::path::Path; +use std::sync::Arc; use async_trait::async_trait; use chrono::{DateTime, NaiveDateTime, Utc}; @@ -31,7 +32,7 @@ use crate::workspace::{ reciprocal_rank_fusion, }; -use super::libsql_migrations; +use crate::db::libsql_migrations; /// Explicit column list for routines table (matches positional access in `row_to_routine_libsql`). const ROUTINE_COLUMNS: &str = "\ @@ -48,8 +49,12 @@ const ROUTINE_RUN_COLUMNS: &str = "\ status, completed_at, result_summary, tokens_used, job_id, created_at"; /// libSQL/Turso database backend. +/// +/// Stores the `Database` handle in an `Arc` so that the same underlying +/// database can be shared with stores (SecretsStore, WasmToolStore) that +/// create their own connections per-operation. pub struct LibSqlBackend { - db: LibSqlDatabase, + db: Arc, } impl LibSqlBackend { @@ -67,7 +72,7 @@ impl LibSqlBackend { .await .map_err(|e| DatabaseError::Pool(format!("Failed to open libSQL database: {}", e)))?; - Ok(Self { db }) + Ok(Self { db: Arc::new(db) }) } /// Create a new in-memory database (for testing). @@ -79,7 +84,7 @@ impl LibSqlBackend { DatabaseError::Pool(format!("Failed to create in-memory database: {}", e)) })?; - Ok(Self { db }) + Ok(Self { db: Arc::new(db) }) } /// Create with Turso cloud sync (embedded replica). @@ -99,18 +104,18 @@ impl LibSqlBackend { .await .map_err(|e| DatabaseError::Pool(format!("Failed to open remote replica: {}", e)))?; - Ok(Self { db }) + Ok(Self { db: Arc::new(db) }) } - /// Get the underlying database handle (for sync operations). - pub fn database(&self) -> &LibSqlDatabase { - &self.db - } - - /// Create a new connection to the same database. + /// Get a shared reference to the underlying database handle. /// - /// Used for creating separate connections for SecretsStore and WasmToolStore - /// which have their own connection needs. + /// Use this to pass the database to stores (SecretsStore, WasmToolStore) + /// that need to create their own connections per-operation. + pub fn shared_db(&self) -> Arc { + Arc::clone(&self.db) + } + + /// Create a new connection to the database. pub fn connect(&self) -> Result { self.db .connect() diff --git a/src/db/libsql_migrations.rs b/src/db/libsql_migrations.rs index 6de9ff5a..1480ed7d 100644 --- a/src/db/libsql_migrations.rs +++ b/src/db/libsql_migrations.rs @@ -477,4 +477,73 @@ CREATE TABLE IF NOT EXISTS settings ( CREATE INDEX IF NOT EXISTS idx_settings_user ON settings(user_id); +-- ==================== Missing indexes (parity with PostgreSQL) ==================== + +-- agent_jobs +CREATE INDEX IF NOT EXISTS idx_agent_jobs_stuck ON agent_jobs(stuck_since); + +-- secrets +CREATE INDEX IF NOT EXISTS idx_secrets_provider ON secrets(provider); +CREATE INDEX IF NOT EXISTS idx_secrets_expires ON secrets(expires_at); + +-- wasm_tools +CREATE INDEX IF NOT EXISTS idx_wasm_tools_trust ON wasm_tools(trust_level); + +-- tool_capabilities +CREATE INDEX IF NOT EXISTS idx_tool_capabilities_tool ON tool_capabilities(wasm_tool_id); + +-- leak_detection_patterns +CREATE INDEX IF NOT EXISTS idx_leak_patterns_enabled ON leak_detection_patterns(enabled); + +-- tool_rate_limit_state +CREATE INDEX IF NOT EXISTS idx_rate_limit_tool ON tool_rate_limit_state(wasm_tool_id); + +-- secret_usage_log +CREATE INDEX IF NOT EXISTS idx_secret_usage_secret ON secret_usage_log(secret_id); +CREATE INDEX IF NOT EXISTS idx_secret_usage_tool ON secret_usage_log(wasm_tool_id); +CREATE INDEX IF NOT EXISTS idx_secret_usage_created ON secret_usage_log(created_at DESC); + +-- leak_detection_events +CREATE INDEX IF NOT EXISTS idx_leak_events_pattern ON leak_detection_events(pattern_id); +CREATE INDEX IF NOT EXISTS idx_leak_events_tool ON leak_detection_events(wasm_tool_id); +CREATE INDEX IF NOT EXISTS idx_leak_events_user ON leak_detection_events(user_id); +CREATE INDEX IF NOT EXISTS idx_leak_events_created ON leak_detection_events(created_at DESC); + +-- tool_failures +CREATE INDEX IF NOT EXISTS idx_tool_failures_count ON tool_failures(error_count DESC); +CREATE INDEX IF NOT EXISTS idx_tool_failures_unrepaired ON tool_failures(tool_name); + +-- routines +CREATE INDEX IF NOT EXISTS idx_routines_next_fire ON routines(next_fire_at); +CREATE INDEX IF NOT EXISTS idx_routines_event_triggers ON routines(user_id); + +-- routine_runs +CREATE INDEX IF NOT EXISTS idx_routine_runs_status ON routine_runs(status); + +-- heartbeat_state +CREATE INDEX IF NOT EXISTS idx_heartbeat_next_run ON heartbeat_state(next_run); + +-- ==================== Seed data ==================== + +-- Pre-populate leak detection patterns (matches PostgreSQL V2 migration). +INSERT OR IGNORE INTO leak_detection_patterns (id, name, pattern, severity, action, enabled, created_at) VALUES + ('550e8400-e29b-41d4-a716-446655440001', 'openai_api_key', 'sk-(?:proj-)?[a-zA-Z0-9]{20,}(?:T3BlbkFJ[a-zA-Z0-9_-]*)?', 'critical', 'block', 1, datetime('now')), + ('550e8400-e29b-41d4-a716-446655440002', 'anthropic_api_key', 'sk-ant-api[a-zA-Z0-9_-]{90,}', 'critical', 'block', 1, datetime('now')), + ('550e8400-e29b-41d4-a716-446655440003', 'aws_access_key', 'AKIA[0-9A-Z]{16}', 'critical', 'block', 1, datetime('now')), + ('550e8400-e29b-41d4-a716-446655440004', 'aws_secret_key', '(? anyhow::Result<()> { // // NOTE: For simpler call sites (CLI commands, Memory handler) use the shared // helper `ironclaw::db::connect_from_config()`. This block is kept inline - // because it also captures backend-specific handles (`pg_pool`, `libsql_conn`) + // because it also captures backend-specific handles (`pg_pool`, `libsql_db`) // needed by the secrets store. #[cfg(feature = "postgres")] let mut pg_pool: Option = None; #[cfg(feature = "libsql")] - let mut libsql_conn: Option = None; + let mut libsql_db: Option> = None; let db: Option> = if cli.no_db { tracing::warn!("Running without database connection"); @@ -370,11 +372,9 @@ async fn main() -> anyhow::Result<()> { .unwrap_or(&default_path); let backend = if let Some(ref url) = config.database.libsql_url { - let token = config - .database - .libsql_auth_token - .as_ref() - .expect("LIBSQL_AUTH_TOKEN required when LIBSQL_URL is set"); + let token = config.database.libsql_auth_token.as_ref().ok_or_else(|| { + anyhow::anyhow!("LIBSQL_AUTH_TOKEN is required when LIBSQL_URL is set") + })?; LibSqlBackend::new_remote_replica(db_path, url, token.expose_secret()).await? } else { LibSqlBackend::new_local(db_path).await? @@ -382,8 +382,8 @@ async fn main() -> anyhow::Result<()> { backend.run_migrations().await?; tracing::info!("libSQL database connected and migrations applied"); - // Capture an extra connection for SecretsStore / WasmToolStore - libsql_conn = Some(backend.connect().map_err(|e| anyhow::anyhow!("{}", e))?); + // Capture the Database handle for SecretsStore (connection-per-op) + libsql_db = Some(backend.shared_db()); Some(Arc::new(backend) as Arc) } @@ -516,47 +516,48 @@ async fn main() -> anyhow::Result<()> { tracing::info!("Builder mode enabled"); } - // Create secrets store if master key is configured (needed for MCP auth and WASM channels) - let secrets_store: Option> = { - #[cfg(feature = "postgres")] - { - if let (Some(pool), Some(master_key)) = (&pg_pool, config.secrets.master_key()) { - match SecretsCrypto::new(master_key.clone()) { - Ok(crypto) => Some(Arc::new(PostgresSecretsStore::new( - pool.clone(), - Arc::new(crypto), - ))), - Err(e) => { - tracing::warn!("Failed to initialize secrets crypto: {}", e); - None - } + // Create secrets store if master key is configured (needed for MCP auth and WASM channels). + // + // When both `postgres` and `libsql` features are compiled, the runtime-selected + // backend determines which store is created: whichever DB init branch ran will + // have set its handle (pg_pool or libsql_db), and the or_else chain picks it up. + let secrets_store: Option> = + if let Some(master_key) = config.secrets.master_key() { + match SecretsCrypto::new(master_key.clone()) { + Ok(crypto) => { + let crypto = Arc::new(crypto); + let store: Option> = None; + + #[cfg(feature = "libsql")] + let store = store.or_else(|| { + libsql_db.take().map(|db| { + Arc::new(LibSqlSecretsStore::new(db, Arc::clone(&crypto))) + as Arc + }) + }); + + #[cfg(feature = "postgres")] + let store = store.or_else(|| { + pg_pool.as_ref().map(|pool| { + Arc::new(PostgresSecretsStore::new(pool.clone(), Arc::clone(&crypto))) + as Arc + }) + }); + + store } - } else { - None - } - } - #[cfg(all(feature = "libsql", not(feature = "postgres")))] - { - if let (Some(conn), Some(master_key)) = - (libsql_conn.take(), config.secrets.master_key()) - { - match SecretsCrypto::new(master_key.clone()) { - Ok(crypto) => Some(Arc::new(LibSqlSecretsStore::new(conn, Arc::new(crypto))) - as Arc), - Err(e) => { - tracing::warn!("Failed to initialize secrets crypto: {}", e); - None - } + Err(e) => { + tracing::warn!("Failed to initialize secrets crypto: {}", e); + #[cfg(feature = "libsql")] + let _ = libsql_db.take(); + None } - } else { - None } - } - #[cfg(not(any(feature = "postgres", feature = "libsql")))] - { + } else { + #[cfg(feature = "libsql")] + let _ = libsql_db.take(); None - } - }; + }; let mcp_session_manager = Arc::new(McpSessionManager::new()); diff --git a/src/secrets/store.rs b/src/secrets/store.rs index 04d896c6..1fee4086 100644 --- a/src/secrets/store.rs +++ b/src/secrets/store.rs @@ -307,17 +307,26 @@ fn row_to_secret(row: &tokio_postgres::Row) -> Secret { // ==================== libSQL implementation ==================== /// libSQL/Turso implementation of SecretsStore. +/// +/// Holds an `Arc` handle and creates a fresh connection per operation, +/// matching the connection-per-request pattern used by the main `LibSqlBackend`. #[cfg(feature = "libsql")] pub struct LibSqlSecretsStore { - conn: libsql::Connection, + db: Arc, crypto: Arc, } #[cfg(feature = "libsql")] impl LibSqlSecretsStore { - /// Create a new store with the given libsql connection and crypto instance. - pub fn new(conn: libsql::Connection, crypto: Arc) -> Self { - Self { conn, crypto } + /// Create a new store with the given shared libsql database handle and crypto instance. + pub fn new(db: Arc, crypto: Arc) -> Self { + Self { db, crypto } + } + + fn connect(&self) -> Result { + self.db + .connect() + .map_err(|e| SecretError::Database(format!("Connection failed: {}", e))) } } @@ -340,8 +349,8 @@ impl SecretsStore for LibSqlSecretsStore { .map(|dt| dt.to_rfc3339_opts(chrono::SecondsFormat::Millis, true)); // Start transaction for atomic upsert + read-back - let tx = self - .conn + let conn = self.connect()?; + let tx = conn .transaction() .await .map_err(|e| SecretError::Database(e.to_string()))?; @@ -401,8 +410,8 @@ impl SecretsStore for LibSqlSecretsStore { } async fn get(&self, user_id: &str, name: &str) -> Result { - let mut rows = self - .conn + let conn = self.connect()?; + let mut rows = conn .query( r#" SELECT id, user_id, name, encrypted_value, key_salt, provider, expires_at, @@ -446,8 +455,8 @@ impl SecretsStore for LibSqlSecretsStore { } async fn exists(&self, user_id: &str, name: &str) -> Result { - let mut rows = self - .conn + let conn = self.connect()?; + let mut rows = conn .query( "SELECT 1 FROM secrets WHERE user_id = ?1 AND name = ?2", libsql::params![user_id, name], @@ -463,8 +472,8 @@ impl SecretsStore for LibSqlSecretsStore { } async fn list(&self, user_id: &str) -> Result, SecretError> { - let mut rows = self - .conn + let conn = self.connect()?; + let mut rows = conn .query( "SELECT name, provider FROM secrets WHERE user_id = ?1 ORDER BY name", libsql::params![user_id], @@ -487,8 +496,8 @@ impl SecretsStore for LibSqlSecretsStore { } async fn delete(&self, user_id: &str, name: &str) -> Result { - let affected = self - .conn + let conn = self.connect()?; + let affected = conn .execute( "DELETE FROM secrets WHERE user_id = ?1 AND name = ?2", libsql::params![user_id, name], @@ -501,18 +510,18 @@ impl SecretsStore for LibSqlSecretsStore { async fn record_usage(&self, secret_id: Uuid) -> Result<(), SecretError> { let now = Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Millis, true); + let conn = self.connect()?; - self.conn - .execute( - r#" + conn.execute( + r#" UPDATE secrets SET last_used_at = ?1, usage_count = usage_count + 1 WHERE id = ?2 "#, - libsql::params![now.as_str(), secret_id.to_string()], - ) - .await - .map_err(|e| SecretError::Database(e.to_string()))?; + libsql::params![now.as_str(), secret_id.to_string()], + ) + .await + .map_err(|e| SecretError::Database(e.to_string()))?; Ok(()) } diff --git a/src/setup/wizard.rs b/src/setup/wizard.rs index 6fbb2c4d..0706296e 100644 --- a/src/setup/wizard.rs +++ b/src/setup/wizard.rs @@ -803,11 +803,8 @@ impl SetupWizard { crypto: &Arc, ) -> Result>, SetupError> { if let Some(ref backend) = self.db_backend { - let conn = backend - .connect() - .map_err(|e| SetupError::Database(format!("Failed to create connection: {}", e)))?; let store: Arc = Arc::new(crate::secrets::LibSqlSecretsStore::new( - conn, + backend.shared_db(), Arc::clone(crypto), )); Ok(Some(store)) @@ -1136,6 +1133,7 @@ impl Default for SetupWizard { } /// Mask password in a database URL for display. +#[cfg(feature = "postgres")] fn mask_password_in_url(url: &str) -> String { // URL format: scheme://user:password@host/database // Find "://" to locate start of credentials @@ -1326,6 +1324,7 @@ mod tests { } #[test] + #[cfg(feature = "postgres")] fn test_mask_password_in_url() { assert_eq!( mask_password_in_url("postgres://user:secret@localhost/db"), diff --git a/src/tools/wasm/storage.rs b/src/tools/wasm/storage.rs index 0b99bd9f..f3762aaa 100644 --- a/src/tools/wasm/storage.rs +++ b/src/tools/wasm/storage.rs @@ -567,15 +567,24 @@ fn row_to_tool(row: &tokio_postgres::Row) -> Result` handle and creates a fresh connection per operation, +/// matching the connection-per-request pattern used by the main `LibSqlBackend`. #[cfg(feature = "libsql")] pub struct LibSqlWasmToolStore { - conn: libsql::Connection, + db: std::sync::Arc, } #[cfg(feature = "libsql")] impl LibSqlWasmToolStore { - pub fn new(conn: libsql::Connection) -> Self { - Self { conn } + pub fn new(db: std::sync::Arc) -> Self { + Self { db } + } + + fn connect(&self) -> Result { + self.db + .connect() + .map_err(|e| WasmStorageError::Database(format!("Connection failed: {}", e))) } } @@ -589,9 +598,15 @@ impl WasmToolStore for LibSqlWasmToolStore { let schema_str = serde_json::to_string(¶ms.parameters_schema) .map_err(|e| WasmStorageError::InvalidData(e.to_string()))?; - self.conn - .execute( - r#" + // Wrap INSERT + read-back in a transaction to prevent TOCTOU races + let conn = self.connect()?; + let tx = conn + .transaction() + .await + .map_err(|e| WasmStorageError::Database(e.to_string()))?; + + tx.execute( + r#" INSERT INTO wasm_tools ( id, user_id, name, version, description, wasm_binary, binary_hash, parameters_schema, source_url, trust_level, status, created_at, updated_at @@ -605,26 +620,25 @@ impl WasmToolStore for LibSqlWasmToolStore { source_url = excluded.source_url, updated_at = ?11 "#, - libsql::params![ - id.to_string(), - params.user_id.as_str(), - params.name.as_str(), - params.version.as_str(), - params.description.as_str(), - libsql::Value::Blob(params.wasm_binary), - libsql::Value::Blob(binary_hash), - schema_str.as_str(), - libsql_wasm_opt_text(params.source_url.as_deref()), - params.trust_level.to_string(), - now.as_str(), - ], - ) - .await - .map_err(|e| WasmStorageError::Database(e.to_string()))?; + libsql::params![ + id.to_string(), + params.user_id.as_str(), + params.name.as_str(), + params.version.as_str(), + params.description.as_str(), + libsql::Value::Blob(params.wasm_binary), + libsql::Value::Blob(binary_hash), + schema_str.as_str(), + libsql_wasm_opt_text(params.source_url.as_deref()), + params.trust_level.to_string(), + now.as_str(), + ], + ) + .await + .map_err(|e| WasmStorageError::Database(e.to_string()))?; - // Read back the row - let mut rows = self - .conn + // Read back the row within the same transaction + let mut rows = tx .query( r#" SELECT id, user_id, name, version, description, parameters_schema, @@ -647,12 +661,18 @@ impl WasmToolStore for LibSqlWasmToolStore { WasmStorageError::Database("Insert succeeded but row not found".into()) })?; - libsql_row_to_tool(&row) + let tool = libsql_row_to_tool(&row)?; + + tx.commit() + .await + .map_err(|e| WasmStorageError::Database(e.to_string()))?; + + Ok(tool) } async fn get(&self, user_id: &str, name: &str) -> Result { - let mut rows = self - .conn + let conn = self.connect()?; + let mut rows = conn .query( r#" SELECT id, user_id, name, version, description, parameters_schema, @@ -689,8 +709,8 @@ impl WasmToolStore for LibSqlWasmToolStore { user_id: &str, name: &str, ) -> Result { - let mut rows = self - .conn + let conn = self.connect()?; + let mut rows = conn .query( r#" SELECT id, user_id, name, version, description, wasm_binary, binary_hash, @@ -748,8 +768,8 @@ impl WasmToolStore for LibSqlWasmToolStore { &self, tool_id: Uuid, ) -> Result, WasmStorageError> { - let mut rows = self - .conn + let conn = self.connect()?; + let mut rows = conn .query( r#" SELECT id, wasm_tool_id, http_allowlist, allowed_secrets, tool_aliases, @@ -818,8 +838,8 @@ impl WasmToolStore for LibSqlWasmToolStore { async fn list(&self, user_id: &str) -> Result, WasmStorageError> { // SQLite doesn't have DISTINCT ON, so we use a subquery to get latest version per name - let mut rows = self - .conn + let conn = self.connect()?; + let mut rows = conn .query( r#" SELECT id, user_id, name, version, description, parameters_schema, @@ -857,9 +877,9 @@ impl WasmToolStore for LibSqlWasmToolStore { status: ToolStatus, ) -> Result<(), WasmStorageError> { let now = Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Millis, true); + let conn = self.connect()?; - let result = self - .conn + let result = conn .execute( "UPDATE wasm_tools SET status = ?1, updated_at = ?2 WHERE user_id = ?3 AND name = ?4", libsql::params![status.to_string(), now.as_str(), user_id, name], @@ -875,8 +895,8 @@ impl WasmToolStore for LibSqlWasmToolStore { } async fn delete(&self, user_id: &str, name: &str) -> Result { - let result = self - .conn + let conn = self.connect()?; + let result = conn .execute( "DELETE FROM wasm_tools WHERE user_id = ?1 AND name = ?2", libsql::params![user_id, name],