tools/tusk
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: harden security, reduce duplication, and improve robustness

Browse Source 
- Fix SQL injection in data.rs by wrapping get_table_data in READ ONLY transaction
- Fix SQL injection in docker.rs CREATE DATABASE via escape_ident
- Fix command injection in docker.rs by validating pg_version/container_name
  and escaping shell-interpolated values
- Fix UTF-8 panic on stderr truncation with char_indices
- Wrap delete_rows in a transaction for atomicity
- Replace .expect() with proper error propagation in lib.rs
- Cache AI settings in AppState to avoid repeated disk reads
- Cap JSONB column discovery at 50 to prevent unbounded queries
- Fix ERD colorMode to respect system theme via useTheme()
- Extract AppState::get_pool() replacing ~19 inline pool patterns
- Extract shared AiSettingsFields component (DRY popover + sheet)
- Make get_connections_path pub(crate) and reuse from docker.rs
- Deduplicate check_docker by delegating to check_docker_internal

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Alexey Shakhmatov
2026-02-21 11:41:14 +03:00
parent baa794b66a
commit d507162377
15 changed files with 1196 additions and 667 deletions
Download Patch File Download Diff File Expand all files Collapse all files

948
src-tauri/src/commands/ai.rs
View File

File diff suppressed because it is too large Load Diff

2
src-tauri/src/commands/connections.rs
View File

@@ -14,7 +14,7 @@ pub struct ConnectResult {
pub flavor: DbFlavor,
}
fn get_connections_path(app: &AppHandle) -> TuskResult<std::path::PathBuf> {
pub(crate) fn get_connections_path(app: &AppHandle) -> TuskResult<std::path::PathBuf> {
let dir = app
.path()
.app_data_dir()

58
src-tauri/src/commands/data.rs
View File

@@ -21,10 +21,7 @@ pub async fn get_table_data(
sort_direction: Option<String>,
filter: Option<String>,
) -> TuskResult<PaginatedQueryResult> {
let pools = state.pools.read().await;
let pool = pools
.get(&connection_id)
.ok_or(TuskError::NotConnected(connection_id))?;
let pool = state.get_pool(&connection_id).await?;
let qualified = format!("{}.{}", escape_ident(&schema), escape_ident(&table));
@@ -56,11 +53,24 @@ pub async fn get_table_data(
let start = Instant::now();
let (rows, count_row) = tokio::try_join!(
sqlx::query(&data_sql).fetch_all(pool),
sqlx::query(&count_sql).fetch_one(pool),
)
.map_err(TuskError::Database)?;
// Always run table data queries in a read-only transaction to prevent
// writable CTEs or other mutation via the raw filter parameter.
let mut tx = (&pool).begin().await.map_err(TuskError::Database)?;
sqlx::query("SET TRANSACTION READ ONLY")
.execute(&mut *tx)
.await
.map_err(TuskError::Database)?;
let rows = sqlx::query(&data_sql)
.fetch_all(&mut *tx)
.await
.map_err(TuskError::Database)?;
let count_row = sqlx::query(&count_sql)
.fetch_one(&mut *tx)
.await
.map_err(TuskError::Database)?;
tx.rollback().await.map_err(TuskError::Database)?;
let execution_time_ms = start.elapsed().as_millis();
let total_rows: i64 = count_row.get(0);
@@ -134,10 +144,7 @@ pub async fn update_row(
return Err(TuskError::ReadOnly);
}
let pools = state.pools.read().await;
let pool = pools
.get(&connection_id)
.ok_or(TuskError::NotConnected(connection_id))?;
let pool = state.get_pool(&connection_id).await?;
let qualified = format!("{}.{}", escape_ident(&schema), escape_ident(&table));
@@ -155,7 +162,7 @@ pub async fn update_row(
let mut query = sqlx::query(&sql);
query = bind_json_value(query, &value);
query = query.bind(ctid_val);
query.execute(pool).await.map_err(TuskError::Database)?;
query.execute(&pool).await.map_err(TuskError::Database)?;
} else {
let where_parts: Vec<String> = pk_columns
.iter()
@@ -174,7 +181,7 @@ pub async fn update_row(
for pk_val in &pk_values {
query = bind_json_value(query, pk_val);
}
query.execute(pool).await.map_err(TuskError::Database)?;
query.execute(&pool).await.map_err(TuskError::Database)?;
}
Ok(())
@@ -193,10 +200,7 @@ pub async fn insert_row(
return Err(TuskError::ReadOnly);
}
let pools = state.pools.read().await;
let pool = pools
.get(&connection_id)
.ok_or(TuskError::NotConnected(connection_id))?;
let pool = state.get_pool(&connection_id).await?;
let qualified = format!("{}.{}", escape_ident(&schema), escape_ident(&table));
@@ -215,7 +219,7 @@ pub async fn insert_row(
query = bind_json_value(query, val);
}
query.execute(pool).await.map_err(TuskError::Database)?;
query.execute(&pool).await.map_err(TuskError::Database)?;
Ok(())
}
@@ -234,14 +238,14 @@ pub async fn delete_rows(
return Err(TuskError::ReadOnly);
}
let pools = state.pools.read().await;
let pool = pools
.get(&connection_id)
.ok_or(TuskError::NotConnected(connection_id))?;
let pool = state.get_pool(&connection_id).await?;
let qualified = format!("{}.{}", escape_ident(&schema), escape_ident(&table));
let mut total_affected: u64 = 0;
// Wrap all deletes in a transaction for atomicity
let mut tx = (&pool).begin().await.map_err(TuskError::Database)?;
if pk_columns.is_empty() {
// Fallback: use ctids for row identification
let ctid_list = ctids.ok_or_else(|| {
@@ -250,7 +254,7 @@ pub async fn delete_rows(
for ctid_val in &ctid_list {
let sql = format!("DELETE FROM {} WHERE ctid = $1::tid", qualified);
let query = sqlx::query(&sql).bind(ctid_val);
let result = query.execute(pool).await.map_err(TuskError::Database)?;
let result = query.execute(&mut *tx).await.map_err(TuskError::Database)?;
total_affected += result.rows_affected();
}
} else {
@@ -269,11 +273,13 @@ pub async fn delete_rows(
query = bind_json_value(query, val);
}
let result = query.execute(pool).await.map_err(TuskError::Database)?;
let result = query.execute(&mut *tx).await.map_err(TuskError::Database)?;
total_affected += result.rows_affected();
}
}
tx.commit().await.map_err(TuskError::Database)?;
Ok(total_affected)
}

127
src-tauri/src/commands/docker.rs
View File

@@ -4,9 +4,10 @@ use crate::models::docker::{
CloneMode, CloneProgress, CloneResult, CloneToDockerParams, DockerStatus, TuskContainer,
};
use crate::state::AppState;
use crate::utils::escape_ident;
use std::fs;
use std::sync::Arc;
use tauri::{AppHandle, Emitter, Manager, State};
use tauri::{AppHandle, Emitter, State};
use tokio::process::Command;
async fn docker_cmd(state: &AppState) -> Command {
@@ -42,17 +43,8 @@ fn emit_progress(
);
}
fn get_connections_path(app: &AppHandle) -> TuskResult<std::path::PathBuf> {
let dir = app
.path()
.app_data_dir()
.map_err(|e| TuskError::Custom(e.to_string()))?;
fs::create_dir_all(&dir)?;
Ok(dir.join("connections.json"))
}
fn load_connection_config(app: &AppHandle, connection_id: &str) -> TuskResult<ConnectionConfig> {
let path = get_connections_path(app)?;
let path = super::connections::get_connections_path(app)?;
if !path.exists() {
return Err(TuskError::ConnectionNotFound(connection_id.to_string()));
}
@@ -69,43 +61,58 @@ fn shell_escape(s: &str) -> String {
s.replace('\'', "'\\''")
}
/// Validate pg_version matches a safe pattern (e.g. "16", "16.2", "17.1")
fn validate_pg_version(version: &str) -> TuskResult<()> {
let is_valid = !version.is_empty()
&& version
.chars()
.all(|c| c.is_ascii_digit() || c == '.');
if !is_valid {
return Err(docker_err(format!(
"Invalid pg_version '{}': must contain only digits and dots (e.g. '16', '16.2')",
version
)));
}
Ok(())
}
/// Validate container name matches Docker naming rules: [a-zA-Z0-9][a-zA-Z0-9_.-]*
fn validate_container_name(name: &str) -> TuskResult<()> {
if name.is_empty() {
return Err(docker_err("Container name cannot be empty"));
}
let first = name.chars().next().unwrap();
if !first.is_ascii_alphanumeric() {
return Err(docker_err(format!(
"Invalid container name '{}': must start with an alphanumeric character",
name
)));
}
let is_valid = name
.chars()
.all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '.' || c == '-');
if !is_valid {
return Err(docker_err(format!(
"Invalid container name '{}': only [a-zA-Z0-9_.-] characters are allowed",
name
)));
}
Ok(())
}
/// Shell-escape a string for use inside double-quoted shell contexts
fn shell_escape_double(s: &str) -> String {
s.replace('\\', "\\\\")
.replace('"', "\\\"")
.replace('$', "\\$")
.replace('`', "\\`")
.replace('!', "\\!")
}
#[tauri::command]
pub async fn check_docker(state: State<'_, Arc<AppState>>) -> TuskResult<DockerStatus> {
let output = docker_cmd(&state)
.await
.args(["version", "--format", "{{.Server.Version}}"])
.output()
.await;
match output {
Ok(out) => {
if out.status.success() {
let version = String::from_utf8_lossy(&out.stdout).trim().to_string();
Ok(DockerStatus {
installed: true,
daemon_running: true,
version: Some(version),
error: None,
})
} else {
let stderr = String::from_utf8_lossy(&out.stderr).trim().to_string();
let daemon_running = !stderr.contains("Cannot connect")
&& !stderr.contains("connection refused");
Ok(DockerStatus {
installed: true,
daemon_running,
version: None,
error: Some(stderr),
})
}
}
Err(_) => Ok(DockerStatus {
installed: false,
daemon_running: false,
version: None,
error: Some("Docker CLI not found. Please install Docker.".to_string()),
}),
}
let docker_host = state.docker_host.read().await.clone();
check_docker_internal(&docker_host).await
}
#[tauri::command]
@@ -252,6 +259,10 @@ async fn do_clone(
params: &CloneToDockerParams,
clone_id: &str,
) -> TuskResult<CloneResult> {
// Validate user inputs before any operations
validate_pg_version(&params.pg_version)?;
validate_container_name(&params.container_name)?;
let docker_host = state.docker_host.read().await.clone();
// Step 1: Check Docker
@@ -313,7 +324,7 @@ async fn do_clone(
.args([
"exec", &params.container_name,
"psql", "-U", "postgres", "-c",
&format!("CREATE DATABASE \"{}\"", params.source_database),
&format!("CREATE DATABASE {}", escape_ident(&params.source_database)),
])
.output()
.await
@@ -492,7 +503,11 @@ async fn run_pipe_cmd(
if !stderr.is_empty() {
// Truncate for progress display (full log can be long)
let short = if stderr.len() > 500 {
format!("{}...", &stderr[..500])
let truncated = stderr.char_indices()
.nth(500)
.map(|(i, _)| &stderr[..i])
.unwrap_or(&stderr);
format!("{}...", truncated)
} else {
stderr.clone()
};
@@ -633,13 +648,16 @@ async fn transfer_sample_data(
let table = parts[1];
// Use COPY (SELECT ... LIMIT N) TO STDOUT piped to COPY ... FROM STDIN
// Escape schema/table for use inside double-quoted shell strings
let escaped_schema = shell_escape_double(schema);
let escaped_table = shell_escape_double(table);
let copy_out_sql = format!(
"\\copy (SELECT * FROM \\\"{}\\\".\\\"{}\\\" LIMIT {}) TO STDOUT",
schema, table, sample_rows
escaped_schema, escaped_table, sample_rows
);
let copy_in_sql = format!(
"\\copy \\\"{}\\\".\\\"{}\\\" FROM STDIN",
schema, table
escaped_schema, escaped_table
);
let escaped_url = shell_escape(source_url);
@@ -693,7 +711,7 @@ async fn transfer_sample_data(
}
fn save_connection_config(app: &AppHandle, config: &ConnectionConfig) -> TuskResult<()> {
let path = get_connections_path(app)?;
let path = super::connections::get_connections_path(app)?;
let mut connections = if path.exists() {
let data = fs::read_to_string(&path)?;
serde_json::from_str::<Vec<ConnectionConfig>>(&data)?
@@ -701,7 +719,12 @@ fn save_connection_config(app: &AppHandle, config: &ConnectionConfig) -> TuskRes
vec![]
};
connections.push(config.clone());
// Upsert by ID to avoid duplicate entries on retry
if let Some(pos) = connections.iter().position(|c| c.id == config.id) {
connections[pos] = config.clone();
} else {
connections.push(config.clone());
}
let data = serde_json::to_string_pretty(&connections)?;
fs::write(&path, data)?;

100
src-tauri/src/commands/schema.rs
View File

@@ -14,17 +14,14 @@ pub async fn list_databases(
state: State<'_, Arc<AppState>>,
connection_id: String,
) -> TuskResult<Vec<String>> {
let pools = state.pools.read().await;
let pool = pools
.get(&connection_id)
.ok_or(TuskError::NotConnected(connection_id))?;
let pool = state.get_pool(&connection_id).await?;
let rows = sqlx::query(
"SELECT datname FROM pg_database \
WHERE datistemplate = false \
ORDER BY datname",
)
.fetch_all(pool)
.fetch_all(&pool)
.await
.map_err(TuskError::Database)?;
@@ -35,10 +32,7 @@ pub async fn list_schemas_core(
state: &AppState,
connection_id: &str,
) -> TuskResult<Vec<String>> {
let pools = state.pools.read().await;
let pool = pools
.get(connection_id)
.ok_or_else(|| TuskError::NotConnected(connection_id.to_string()))?;
let pool = state.get_pool(connection_id).await?;
let flavor = state.get_flavor(connection_id).await;
let sql = if flavor == DbFlavor::Greenplum {
@@ -52,7 +46,7 @@ pub async fn list_schemas_core(
};
let rows = sqlx::query(sql)
.fetch_all(pool)
.fetch_all(&pool)
.await
.map_err(TuskError::Database)?;
@@ -72,10 +66,7 @@ pub async fn list_tables_core(
connection_id: &str,
schema: &str,
) -> TuskResult<Vec<SchemaObject>> {
let pools = state.pools.read().await;
let pool = pools
.get(connection_id)
.ok_or_else(|| TuskError::NotConnected(connection_id.to_string()))?;
let pool = state.get_pool(connection_id).await?;
let rows = sqlx::query(
"SELECT t.table_name, \
@@ -88,7 +79,7 @@ pub async fn list_tables_core(
ORDER BY t.table_name",
)
.bind(schema)
.fetch_all(pool)
.fetch_all(&pool)
.await
.map_err(TuskError::Database)?;
@@ -119,10 +110,7 @@ pub async fn list_views(
connection_id: String,
schema: String,
) -> TuskResult<Vec<SchemaObject>> {
let pools = state.pools.read().await;
let pool = pools
.get(&connection_id)
.ok_or(TuskError::NotConnected(connection_id))?;
let pool = state.get_pool(&connection_id).await?;
let rows = sqlx::query(
"SELECT table_name FROM information_schema.views \
@@ -130,7 +118,7 @@ pub async fn list_views(
ORDER BY table_name",
)
.bind(&schema)
.fetch_all(pool)
.fetch_all(&pool)
.await
.map_err(TuskError::Database)?;
@@ -152,10 +140,7 @@ pub async fn list_functions(
connection_id: String,
schema: String,
) -> TuskResult<Vec<SchemaObject>> {
let pools = state.pools.read().await;
let pool = pools
.get(&connection_id)
.ok_or(TuskError::NotConnected(connection_id))?;
let pool = state.get_pool(&connection_id).await?;
let rows = sqlx::query(
"SELECT routine_name FROM information_schema.routines \
@@ -163,7 +148,7 @@ pub async fn list_functions(
ORDER BY routine_name",
)
.bind(&schema)
.fetch_all(pool)
.fetch_all(&pool)
.await
.map_err(TuskError::Database)?;
@@ -185,10 +170,7 @@ pub async fn list_indexes(
connection_id: String,
schema: String,
) -> TuskResult<Vec<SchemaObject>> {
let pools = state.pools.read().await;
let pool = pools
.get(&connection_id)
.ok_or(TuskError::NotConnected(connection_id))?;
let pool = state.get_pool(&connection_id).await?;
let rows = sqlx::query(
"SELECT indexname FROM pg_indexes \
@@ -196,7 +178,7 @@ pub async fn list_indexes(
ORDER BY indexname",
)
.bind(&schema)
.fetch_all(pool)
.fetch_all(&pool)
.await
.map_err(TuskError::Database)?;
@@ -218,10 +200,7 @@ pub async fn list_sequences(
connection_id: String,
schema: String,
) -> TuskResult<Vec<SchemaObject>> {
let pools = state.pools.read().await;
let pool = pools
.get(&connection_id)
.ok_or(TuskError::NotConnected(connection_id))?;
let pool = state.get_pool(&connection_id).await?;
let rows = sqlx::query(
"SELECT sequence_name FROM information_schema.sequences \
@@ -229,7 +208,7 @@ pub async fn list_sequences(
ORDER BY sequence_name",
)
.bind(&schema)
.fetch_all(pool)
.fetch_all(&pool)
.await
.map_err(TuskError::Database)?;
@@ -251,10 +230,7 @@ pub async fn get_table_columns_core(
schema: &str,
table: &str,
) -> TuskResult<Vec<ColumnInfo>> {
let pools = state.pools.read().await;
let pool = pools
.get(connection_id)
.ok_or_else(|| TuskError::NotConnected(connection_id.to_string()))?;
let pool = state.get_pool(connection_id).await?;
let rows = sqlx::query(
"SELECT \
@@ -287,7 +263,7 @@ pub async fn get_table_columns_core(
)
.bind(schema)
.bind(table)
.fetch_all(pool)
.fetch_all(&pool)
.await
.map_err(TuskError::Database)?;
@@ -323,10 +299,7 @@ pub async fn get_table_constraints(
schema: String,
table: String,
) -> TuskResult<Vec<ConstraintInfo>> {
let pools = state.pools.read().await;
let pool = pools
.get(&connection_id)
.ok_or(TuskError::NotConnected(connection_id))?;
let pool = state.get_pool(&connection_id).await?;
let rows = sqlx::query(
"SELECT \
@@ -376,7 +349,7 @@ pub async fn get_table_constraints(
)
.bind(&schema)
.bind(&table)
.fetch_all(pool)
.fetch_all(&pool)
.await
.map_err(TuskError::Database)?;
@@ -402,10 +375,7 @@ pub async fn get_table_indexes(
schema: String,
table: String,
) -> TuskResult<Vec<IndexInfo>> {
let pools = state.pools.read().await;
let pool = pools
.get(&connection_id)
.ok_or(TuskError::NotConnected(connection_id))?;
let pool = state.get_pool(&connection_id).await?;
let rows = sqlx::query(
"SELECT \
@@ -422,7 +392,7 @@ pub async fn get_table_indexes(
)
.bind(&schema)
.bind(&table)
.fetch_all(pool)
.fetch_all(&pool)
.await
.map_err(TuskError::Database)?;
@@ -443,10 +413,7 @@ pub async fn get_completion_schema(
connection_id: String,
) -> TuskResult<HashMap<String, HashMap<String, Vec<String>>>> {
let flavor = state.get_flavor(&connection_id).await;
let pools = state.pools.read().await;
let pool = pools
.get(&connection_id)
.ok_or(TuskError::NotConnected(connection_id))?;
let pool = state.get_pool(&connection_id).await?;
let sql = if flavor == DbFlavor::Greenplum {
"SELECT table_schema, table_name, column_name \
@@ -461,7 +428,7 @@ pub async fn get_completion_schema(
};
let rows = sqlx::query(sql)
.fetch_all(pool)
.fetch_all(&pool)
.await
.map_err(TuskError::Database)?;
@@ -490,10 +457,7 @@ pub async fn get_column_details(
table: String,
) -> TuskResult<Vec<ColumnDetail>> {
let flavor = state.get_flavor(&connection_id).await;
let pools = state.pools.read().await;
let pool = pools
.get(&connection_id)
.ok_or(TuskError::NotConnected(connection_id))?;
let pool = state.get_pool(&connection_id).await?;
let sql = if flavor == DbFlavor::Greenplum {
"SELECT c.column_name, c.data_type, \
@@ -516,7 +480,7 @@ pub async fn get_column_details(
let rows = sqlx::query(sql)
.bind(&schema)
.bind(&table)
.fetch_all(pool)
.fetch_all(&pool)
.await
.map_err(TuskError::Database)?;
@@ -539,10 +503,7 @@ pub async fn get_table_triggers(
schema: String,
table: String,
) -> TuskResult<Vec<TriggerInfo>> {
let pools = state.pools.read().await;
let pool = pools
.get(&connection_id)
.ok_or(TuskError::NotConnected(connection_id))?;
let pool = state.get_pool(&connection_id).await?;
let rows = sqlx::query(
"SELECT \
@@ -571,7 +532,7 @@ pub async fn get_table_triggers(
)
.bind(&schema)
.bind(&table)
.fetch_all(pool)
.fetch_all(&pool)
.await
.map_err(TuskError::Database)?;
@@ -595,10 +556,7 @@ pub async fn get_schema_erd(
connection_id: String,
schema: String,
) -> TuskResult<ErdData> {
let pools = state.pools.read().await;
let pool = pools
.get(&connection_id)
.ok_or(TuskError::NotConnected(connection_id))?;
let pool = state.get_pool(&connection_id).await?;
// Get all tables with columns
let col_rows = sqlx::query(
@@ -627,7 +585,7 @@ pub async fn get_schema_erd(
ORDER BY c.table_name, c.ordinal_position",
)
.bind(&schema)
.fetch_all(pool)
.fetch_all(&pool)
.await
.map_err(TuskError::Database)?;
@@ -690,7 +648,7 @@ pub async fn get_schema_erd(
ORDER BY c.conname",
)
.bind(&schema)
.fetch_all(pool)
.fetch_all(&pool)
.await
.map_err(TuskError::Database)?;