xiaoxue_admin
/
dbx
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Merge pull request #91 from Mukesh-234234/main 

issue-88 fix
This commit is contained in:
skyler 2026-05-05 16:04:27 +08:00 committed by GitHub
parent 24537093ad f6857a52bd
commit 2fd76886d2
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
9 changed files with 487 additions and 134 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
crates/dbx-core/src/connection.rs
View File

@ -105,8 +105,8 @@ impl AppState {
PoolKind::Redis(tokio::sync::Mutex::new(con))
}
DatabaseType::DuckDb => {
let con = duckdb::Connection::open(&expand_tilde(&db_config.host)).map_err(|e| e.to_string())?;
PoolKind::DuckDb(Arc::new(std::sync::Mutex::new(con)))
let con = db::duckdb_driver::connect_path(&expand_tilde(&db_config.host))?;
PoolKind::DuckDb(con)
}
DatabaseType::MongoDb => {
let client = db::mongo_driver::connect(&url).await?;

25
crates/dbx-core/src/db/duckdb_driver.rs Normal file
View File

@ -0,0 +1,25 @@
use std::sync::Arc;
use std::sync::Mutex;
use super::file_validator::validate_file_path;
/// Connects to a DuckDb database file with file validation.
///
/// # Arguments
/// * `path` - The file path to the DuckDb database
///
/// # Returns
/// * `Ok(Arc<Mutex<duckdb::Connection>>)` on successful connection
/// * `Err(String)` with descriptive error message if connection fails
pub fn connect_path(path: &str) -> Result<Arc<Mutex<duckdb::Connection>>, String> {
// Validate file path using universal validator
validate_file_path(path, is_network_path)?;
let connection = duckdb::Connection::open(path)
.map_err(|e| format!("DuckDb connection failed: {e}"))?;
Ok(Arc::new(Mutex::new(connection)))
}
fn is_network_path(path: &str) -> bool {
path.starts_with("\\\\") || path.starts_with("//") || path.contains("wsl.localhost") || path.contains("wsl$")
}

98
crates/dbx-core/src/db/file_validator.rs Normal file
View File

@ -0,0 +1,98 @@
use std::path::Path;
/// Validates a file path for database connections.
///
/// Performs comprehensive checks including:
/// - Empty path validation
/// - Null character detection
/// - File existence (for local paths)
/// - File type validation (must be a file, not directory)
/// - Network path detection (skips validation for network paths)
///
/// # Arguments
/// * `path` - The file path to validate
/// * `is_network_path` - Closure to determine if path is a network path
///
/// # Returns
/// * `Ok(())` if validation passes
/// * `Err(String)` with descriptive error message if validation fails
pub fn validate_file_path<F>(path: &str, is_network_path: F) -> Result<(), String>
where
F: Fn(&str) -> bool,
{
// Check if path is empty
if path.is_empty() {
return Err("Database file path cannot be empty".to_string());
}
// Check if path contains invalid characters
if path.contains('\0') {
return Err("Database file path contains null characters".to_string());
}
let path_obj = Path::new(path);
// For non-network paths, perform file system checks
if !is_network_path(path) {
if !path_obj.exists() {
return Err(format!(
"Database file does not exist: {}",
path
));
}
// Check if path is actually a file, not a directory
if path_obj.is_dir() {
return Err(format!(
"Database file path is a directory, not a file: {}",
path
));
}
// Check if path is a valid file
if !path_obj.is_file() {
return Err(format!(
"Database file path is not a valid file: {}",
path
));
}
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
fn is_network_path_test(path: &str) -> bool {
path.starts_with("\\\\") || path.starts_with("//")
}
#[test]
fn test_empty_path() {
let result = validate_file_path("", is_network_path_test);
assert!(result.is_err());
assert!(result.unwrap_err().contains("empty"));
}
#[test]
fn test_null_character() {
let result = validate_file_path("path\0invalid", is_network_path_test);
assert!(result.is_err());
assert!(result.unwrap_err().contains("null"));
}
#[test]
fn test_network_path_skips_validation() {
let result = validate_file_path("//network/path/nonexistent.db", is_network_path_test);
assert!(result.is_ok());
}
#[test]
fn test_nonexistent_local_file() {
let result = validate_file_path("/nonexistent/path/to/file.db", is_network_path_test);
assert!(result.is_err());
assert!(result.unwrap_err().contains("does not exist"));
}
}

3
crates/dbx-core/src/db/mod.rs
View File

@ -8,12 +8,15 @@ pub mod redis_driver;
pub mod sqlite;
pub mod sqlserver;
pub mod ssh_tunnel;
pub mod file_validator;
pub mod duckdb_driver;
use std::future::Future;
use std::time::Duration;
// Re-export types so that `db::QueryResult` etc. work within dbx-core
pub use crate::types::*;
pub use file_validator::validate_file_path;
pub const CONNECTION_TIMEOUT_SECS: u64 = 5;
pub const TCP_PROBE_TIMEOUT_SECS: u64 = 3;

109
crates/dbx-core/src/db/mysql.rs
View File

@ -4,7 +4,6 @@ use sqlx::mysql::{MySqlPool, MySqlPoolOptions, MySqlRow};
use sqlx::{Column, Executor, Row, TypeInfo, ValueRef};
use std::time::{Duration, Instant};
use super::{connection_timeout, with_connection_timeout};
use crate::types::{ColumnInfo, DatabaseInfo, ForeignKeyInfo, IndexInfo, QueryResult, TableInfo, TriggerInfo};
fn quote_value(s: &str) -> String {
@ -24,9 +23,15 @@ fn get_str_by_name(row: &MySqlRow, name: &str) -> String {
}
fn get_opt_str(row: &MySqlRow, name: &str) -> Option<String> {
row.try_get::<Option<String>, _>(name).ok().flatten().or_else(|| {
row.try_get::<Option<Vec<u8>>, _>(name).ok().flatten().map(|b| String::from_utf8_lossy(&b).to_string())
})
row.try_get::<Option<String>, _>(name)
.ok()
.flatten()
.or_else(|| {
row.try_get::<Option<Vec<u8>>, _>(name)
.ok()
.flatten()
.map(|b| String::from_utf8_lossy(&b).to_string())
})
}
fn numeric_metadata_u64_to_i32(value: Option<u64>) -> Option<i32> {
@ -38,7 +43,8 @@ fn numeric_metadata_i64_to_i32(value: Option<i64>) -> Option<i32> {
}
fn numeric_metadata_str_to_i32(value: Option<String>) -> Option<i32> {
value.and_then(|v| v.parse::<i64>().ok()).and_then(|v| i32::try_from(v).ok())
value.and_then(|v| v.parse::<i64>().ok())
.and_then(|v| i32::try_from(v).ok())
}
fn get_opt_i32(row: &MySqlRow, name: &str) -> Option<i32> {
@ -95,14 +101,20 @@ fn mysql_value_to_json(row: &MySqlRow, idx: usize, type_name: &str) -> serde_jso
}
if upper_type == "BOOLEAN" {
return row.try_get::<bool, _>(idx).map(serde_json::Value::Bool).unwrap_or(serde_json::Value::Null);
return row
.try_get::<bool, _>(idx)
.map(serde_json::Value::Bool)
.unwrap_or(serde_json::Value::Null);
}
if upper_type.contains("BIGINT") {
return row
.try_get::<i64, _>(idx)
.map(|v| serde_json::Value::String(v.to_string()))
.or_else(|_| row.try_get::<u64, _>(idx).map(|v| serde_json::Value::String(v.to_string())))
.or_else(|_| {
row.try_get::<u64, _>(idx)
.map(|v| serde_json::Value::String(v.to_string()))
})
.unwrap_or(serde_json::Value::Null);
}
@ -128,24 +140,25 @@ fn mysql_value_to_json(row: &MySqlRow, idx: usize, type_name: &str) -> serde_jso
.map(serde_json::Value::String)
.or_else(|_| row.try_get::<i64, _>(idx).map(|v| serde_json::Value::Number(v.into())))
.or_else(|_| row.try_get::<u64, _>(idx).map(|v| serde_json::Value::Number(v.into())))
.or_else(|_| {
row.try_get::<f64, _>(idx).map(|v| {
serde_json::Number::from_f64(v).map(serde_json::Value::Number).unwrap_or(serde_json::Value::Null)
})
})
.or_else(|_| row.try_get::<f64, _>(idx).map(|v| {
serde_json::Number::from_f64(v)
.map(serde_json::Value::Number)
.unwrap_or(serde_json::Value::Null)
}))
.or_else(|_| row.try_get::<bool, _>(idx).map(serde_json::Value::Bool))
.or_else(|_| {
row.try_get::<Vec<u8>, _>(idx).map(|b| serde_json::Value::String(String::from_utf8_lossy(&b).to_string()))
row.try_get::<Vec<u8>, _>(idx)
.map(|b| serde_json::Value::String(String::from_utf8_lossy(&b).to_string()))
})
.or_else(|e| mysql_temporal_to_json_value(row, idx).ok_or(e))
.unwrap_or(serde_json::Value::Null)
}
pub async fn connect(url: &str) -> Result<MySqlPool, String> {
with_connection_timeout("MySQL", async {
super::with_connection_timeout("MySQL", async {
MySqlPoolOptions::new()
.max_connections(5)
.acquire_timeout(connection_timeout())
.acquire_timeout(super::connection_timeout())
.idle_timeout(Duration::from_secs(300))
.connect(url)
.await
@ -155,13 +168,17 @@ pub async fn connect(url: &str) -> Result<MySqlPool, String> {
}
pub async fn connect_bare(url: &str) -> Result<MySqlPool, String> {
let options: sqlx::mysql::MySqlConnectOptions =
url.parse().map_err(|e: sqlx::Error| format!("Invalid MySQL URL: {e}"))?;
let options = options.no_engine_substitution(false).set_names(false).pipes_as_concat(false).timezone(None);
with_connection_timeout("MySQL", async {
let options: sqlx::mysql::MySqlConnectOptions = url.parse()
.map_err(|e: sqlx::Error| format!("Invalid MySQL URL: {e}"))?;
let options = options
.no_engine_substitution(false)
.set_names(false)
.pipes_as_concat(false)
.timezone(None);
super::with_connection_timeout("MySQL", async {
MySqlPoolOptions::new()
.max_connections(5)
.acquire_timeout(connection_timeout())
.acquire_timeout(super::connection_timeout())
.idle_timeout(Duration::from_secs(300))
.connect_with(options)
.await
@ -184,7 +201,10 @@ pub async fn list_tables(pool: &MySqlPool, database: &str) -> Result<Vec<TableIn
"SELECT TABLE_NAME, TABLE_TYPE FROM information_schema.TABLES WHERE TABLE_SCHEMA = {} ORDER BY TABLE_NAME",
quote_value(database),
);
let rows: Vec<MySqlRow> = sqlx::raw_sql(&sql).fetch_all(pool).await.map_err(|e| e.to_string())?;
let rows: Vec<MySqlRow> = sqlx::raw_sql(&sql)
.fetch_all(pool)
.await
.map_err(|e| e.to_string())?;
Ok(rows
.iter()
@ -195,7 +215,11 @@ pub async fn list_tables(pool: &MySqlPool, database: &str) -> Result<Vec<TableIn
.collect())
}
pub async fn get_columns(pool: &MySqlPool, database: &str, table: &str) -> Result<Vec<ColumnInfo>, String> {
pub async fn get_columns(
pool: &MySqlPool,
database: &str,
table: &str,
) -> Result<Vec<ColumnInfo>, String> {
let sql = format!(
"SELECT c.COLUMN_NAME, c.COLUMN_TYPE, c.IS_NULLABLE, c.COLUMN_DEFAULT, c.EXTRA, c.COLUMN_COMMENT, \
CASE WHEN kcu.COLUMN_NAME IS NOT NULL THEN 1 ELSE 0 END AS IS_PK, \
@ -211,7 +235,10 @@ pub async fn get_columns(pool: &MySqlPool, database: &str, table: &str) -> Resul
quote_value(database),
quote_value(table),
);
let rows: Vec<MySqlRow> = sqlx::raw_sql(&sql).fetch_all(pool).await.map_err(|e| e.to_string())?;
let rows: Vec<MySqlRow> = sqlx::raw_sql(&sql)
.fetch_all(pool)
.await
.map_err(|e| e.to_string())?;
Ok(rows
.iter()
@ -234,13 +261,12 @@ pub async fn execute_query(pool: &MySqlPool, sql: &str, bare: bool) -> Result<Qu
let start = Instant::now();
let trimmed = sql.trim().to_uppercase();
if trimmed.starts_with("SELECT")
|| trimmed.starts_with("SHOW")
|| trimmed.starts_with("DESCRIBE")
|| trimmed.starts_with("EXPLAIN")
{
if trimmed.starts_with("SELECT") || trimmed.starts_with("SHOW") || trimmed.starts_with("DESCRIBE") || trimmed.starts_with("EXPLAIN") {
if bare {
let rows: Vec<MySqlRow> = sqlx::raw_sql(sql).fetch_all(pool).await.map_err(|e| e.to_string())?;
let rows: Vec<MySqlRow> = sqlx::raw_sql(sql)
.fetch_all(pool)
.await
.map_err(|e| e.to_string())?;
let (columns, column_types) = if let Some(first) = rows.first() {
let cols: Vec<String> = first.columns().iter().map(|c| c.name().to_string()).collect();
@ -271,7 +297,10 @@ pub async fn execute_query(pool: &MySqlPool, sql: &str, bare: bool) -> Result<Qu
let columns: Vec<String> = desc.columns().iter().map(|c| c.name().to_string()).collect();
let column_types: Vec<String> = desc.columns().iter().map(|c| c.type_info().name().to_string()).collect();
let rows: Vec<MySqlRow> = sqlx::query(sql).fetch_all(pool).await.map_err(|e| e.to_string())?;
let rows: Vec<MySqlRow> = sqlx::query(sql)
.fetch_all(pool)
.await
.map_err(|e| e.to_string())?;
let result_rows: Vec<Vec<serde_json::Value>> = rows
.iter()
@ -291,7 +320,10 @@ pub async fn execute_query(pool: &MySqlPool, sql: &str, bare: bool) -> Result<Qu
})
}
} else {
let result = sqlx::raw_sql(sql).execute(pool).await.map_err(|e| e.to_string())?;
let result = sqlx::raw_sql(sql)
.execute(pool)
.await
.map_err(|e| e.to_string())?;
Ok(QueryResult {
columns: vec![],
@ -315,7 +347,10 @@ pub async fn list_indexes(pool: &MySqlPool, database: &str, table: &str) -> Resu
quote_value(database),
quote_value(table),
);
let rows: Vec<MySqlRow> = sqlx::raw_sql(&sql).fetch_all(pool).await.map_err(|e| e.to_string())?;
let rows: Vec<MySqlRow> = sqlx::raw_sql(&sql)
.fetch_all(pool)
.await
.map_err(|e| e.to_string())?;
Ok(rows
.iter()
@ -346,7 +381,10 @@ pub async fn list_foreign_keys(pool: &MySqlPool, database: &str, table: &str) ->
quote_value(database),
quote_value(table),
);
let rows: Vec<MySqlRow> = sqlx::raw_sql(&sql).fetch_all(pool).await.map_err(|e| e.to_string())?;
let rows: Vec<MySqlRow> = sqlx::raw_sql(&sql)
.fetch_all(pool)
.await
.map_err(|e| e.to_string())?;
Ok(rows
.iter()
@ -368,7 +406,10 @@ pub async fn list_triggers(pool: &MySqlPool, database: &str, table: &str) -> Res
quote_value(database),
quote_value(table),
);
let rows: Vec<MySqlRow> = sqlx::raw_sql(&sql).fetch_all(pool).await.map_err(|e| e.to_string())?;
let rows: Vec<MySqlRow> = sqlx::raw_sql(&sql)
.fetch_all(pool)
.await
.map_err(|e| e.to_string())?;
Ok(rows
.iter()

109
crates/dbx-core/src/db/postgres.rs
View File

@ -3,9 +3,10 @@ use rust_decimal::Decimal;
use sqlx::postgres::{PgPool, PgPoolOptions, PgRow};
use sqlx::{Column, Executor, Row, TypeInfo, ValueRef};
use std::time::{Duration, Instant};
use percent_encoding::percent_decode_str;
use super::{connection_timeout, with_connection_timeout};
use crate::types::{ColumnInfo, DatabaseInfo, ForeignKeyInfo, IndexInfo, QueryResult, TableInfo, TriggerInfo};
use super::file_validator::validate_file_path;
fn pg_temporal_to_json_value(row: &PgRow, idx: usize) -> Option<serde_json::Value> {
if let Ok(v) = row.try_get::<DateTime<Utc>, _>(idx) {
@ -41,7 +42,10 @@ fn pg_value_to_json(row: &PgRow, idx: usize, type_name: &str) -> serde_json::Val
}
if upper == "BOOL" {
return row.try_get::<bool, _>(idx).map(serde_json::Value::Bool).unwrap_or(serde_json::Value::Null);
return row
.try_get::<bool, _>(idx)
.map(serde_json::Value::Bool)
.unwrap_or(serde_json::Value::Null);
}
if upper.contains("TIMESTAMP")
@ -64,11 +68,19 @@ fn pg_value_to_json(row: &PgRow, idx: usize, type_name: &str) -> serde_json::Val
row.try_get::<String, _>(idx)
.map(serde_json::Value::String)
.or_else(|_| row.try_get::<i64, _>(idx).map(|v| serde_json::Value::Number(v.into())))
.or_else(|_| row.try_get::<i32, _>(idx).map(|v| serde_json::Value::Number(v.into())))
.or_else(|_| {
row.try_get::<i64, _>(idx)
.map(|v| serde_json::Value::Number(v.into()))
})
.or_else(|_| {
row.try_get::<i32, _>(idx)
.map(|v| serde_json::Value::Number(v.into()))
})
.or_else(|_| {
row.try_get::<f64, _>(idx).map(|v| {
serde_json::Number::from_f64(v).map(serde_json::Value::Number).unwrap_or(serde_json::Value::Null)
serde_json::Number::from_f64(v)
.map(serde_json::Value::Number)
.unwrap_or(serde_json::Value::Null)
})
})
.or_else(|_| row.try_get::<bool, _>(idx).map(serde_json::Value::Bool))
@ -77,10 +89,13 @@ fn pg_value_to_json(row: &PgRow, idx: usize, type_name: &str) -> serde_json::Val
}
pub async fn connect(url: &str) -> Result<PgPool, String> {
with_connection_timeout("PostgreSQL", async {
// Validate SSL certificate paths if present in the URL
validate_postgres_ssl_paths(url)?;
super::with_connection_timeout("PostgreSQL", async {
PgPoolOptions::new()
.max_connections(5)
.acquire_timeout(connection_timeout())
.acquire_timeout(super::connection_timeout())
.idle_timeout(Duration::from_secs(300))
.connect(url)
.await
@ -89,13 +104,51 @@ pub async fn connect(url: &str) -> Result<PgPool, String> {
.await
}
pub async fn list_databases(pool: &PgPool) -> Result<Vec<DatabaseInfo>, String> {
let rows: Vec<PgRow> = sqlx::query("SELECT datname FROM pg_database WHERE datistemplate = false ORDER BY datname")
.fetch_all(pool)
.await
.map_err(|e| e.to_string())?;
/// Validates SSL certificate file paths in PostgreSQL connection URLs.
///
/// PostgreSQL connection strings can include SSL parameters like:
/// - sslcert=/path/to/cert.pem
/// - sslkey=/path/to/key.pem
/// - sslrootcert=/path/to/root.pem
fn validate_postgres_ssl_paths(url: &str) -> Result<(), String> {
// Extract query parameters from URL
if let Some(query_start) = url.find('?') {
let query_string = &url[query_start + 1..];
for param in query_string.split('&') {
if let Some((key, value)) = param.split_once('=') {
match key {
"sslcert" | "sslkey" | "sslrootcert" => {
// URL decode the value
let decoded = percent_decode_str(value)
.decode_utf8()
.map_err(|_| format!("Invalid URL encoding in {key}"))?;
// Validate the file path (skip network paths)
validate_file_path(&decoded, |_| false)?;
}
_ => {}
}
}
}
}
Ok(())
}
Ok(rows.iter().map(|row| DatabaseInfo { name: row.get::<String, _>("datname") }).collect())
pub async fn list_databases(pool: &PgPool) -> Result<Vec<DatabaseInfo>, String> {
let rows: Vec<PgRow> =
sqlx::query("SELECT datname FROM pg_database WHERE datistemplate = false ORDER BY datname")
.fetch_all(pool)
.await
.map_err(|e| e.to_string())?;
Ok(rows
.iter()
.map(|row| DatabaseInfo {
name: row.get::<String, _>("datname"),
})
.collect())
}
pub async fn list_tables(pool: &PgPool, schema: &str) -> Result<Vec<TableInfo>, String> {
@ -129,10 +182,17 @@ pub async fn list_schemas(pool: &PgPool) -> Result<Vec<String>, String> {
.await
.map_err(|e| e.to_string())?;
Ok(rows.iter().map(|row| row.get::<String, _>("schema_name")).collect())
Ok(rows
.iter()
.map(|row| row.get::<String, _>("schema_name"))
.collect())
}
pub async fn get_columns(pool: &PgPool, schema: &str, table: &str) -> Result<Vec<ColumnInfo>, String> {
pub async fn get_columns(
pool: &PgPool,
schema: &str,
table: &str,
) -> Result<Vec<ColumnInfo>, String> {
let rows: Vec<PgRow> = sqlx::query(
"SELECT a.attname AS column_name, \
format_type(a.atttypid, a.atttypmod) AS full_type, \
@ -194,7 +254,11 @@ pub async fn execute_query(pool: &PgPool, sql: &str) -> Result<QueryResult, Stri
|| trimmed.starts_with("WITH")
|| trimmed.starts_with("TABLE")
{
let rows: Vec<PgRow> = sqlx::query(sql).persistent(false).fetch_all(pool).await.map_err(|e| e.to_string())?;
let rows: Vec<PgRow> = sqlx::query(sql)
.persistent(false)
.fetch_all(pool)
.await
.map_err(|e| e.to_string())?;
let (columns, column_types): (Vec<String>, Vec<String>) = if let Some(first) = rows.first() {
let cols = first.columns();
@ -214,7 +278,13 @@ pub async fn execute_query(pool: &PgPool, sql: &str) -> Result<QueryResult, Stri
.iter()
.map(|row| {
(0..row.len())
.map(|i| pg_value_to_json(row, i, column_types.get(i).map(String::as_str).unwrap_or("")))
.map(|i| {
pg_value_to_json(
row,
i,
column_types.get(i).map(String::as_str).unwrap_or(""),
)
})
.collect()
})
.collect();
@ -227,7 +297,10 @@ pub async fn execute_query(pool: &PgPool, sql: &str) -> Result<QueryResult, Stri
truncated: false,
})
} else {
let result = sqlx::query(sql).execute(pool).await.map_err(|e| e.to_string())?;
let result = sqlx::query(sql)
.execute(pool)
.await
.map_err(|e| e.to_string())?;
Ok(QueryResult {
columns: vec![],

194
crates/dbx-core/src/db/redis_driver.rs
View File

@ -1,8 +1,6 @@
use redis::{AsyncCommands, FromRedisValue, Value as RedisRawValue};
use serde::{Deserialize, Serialize};
use super::{connection_timeout, CONNECTION_TIMEOUT_SECS};
const STREAM_ENTRY_LIMIT: usize = 100;
const DEFAULT_REDIS_DATABASES: u32 = 16;
@ -29,26 +27,44 @@ pub struct RedisValue {
pub async fn connect(url: &str) -> Result<redis::aio::MultiplexedConnection, String> {
let client = redis::Client::open(url).map_err(|e| format!("Redis connection failed: {e}"))?;
let mut con = tokio::time::timeout(connection_timeout(), client.get_multiplexed_async_connection())
.await
.map_err(|_| format!("Redis connection timed out ({CONNECTION_TIMEOUT_SECS}s)"))?
.map_err(|e| format!("Redis connection failed: {e}"))?;
let mut con = tokio::time::timeout(
super::connection_timeout(),
client.get_multiplexed_async_connection(),
)
.await
.map_err(|_| format!("Redis connection timed out ({}s)", super::CONNECTION_TIMEOUT_SECS))?
.map_err(|e| format!("Redis connection failed: {e}"))?;
tokio::time::timeout(connection_timeout(), redis::cmd("PING").query_async::<String>(&mut con))
.await
.map_err(|_| format!("Redis ping timed out ({CONNECTION_TIMEOUT_SECS}s)"))?
.map_err(|e| format!("Redis authentication failed or command rejected: {e}"))?;
tokio::time::timeout(
super::connection_timeout(),
redis::cmd("PING").query_async::<String>(&mut con),
)
.await
.map_err(|_| format!("Redis ping timed out ({}s)", super::CONNECTION_TIMEOUT_SECS))?
.map_err(|e| format!("Redis authentication failed or command rejected: {e}"))?;
Ok(con)
}
pub async fn list_databases(con: &mut redis::aio::MultiplexedConnection) -> Result<Vec<u32>, String> {
let configured_count =
redis::cmd("CONFIG").arg("GET").arg("databases").query_async(con).await.ok().and_then(parse_database_count);
pub async fn list_databases(
con: &mut redis::aio::MultiplexedConnection,
) -> Result<Vec<u32>, String> {
let configured_count = redis::cmd("CONFIG")
.arg("GET")
.arg("databases")
.query_async(con)
.await
.ok()
.and_then(parse_database_count);
let keyspace_dbs = list_keyspace_databases(con).await.unwrap_or_default();
let database_count = configured_count.unwrap_or(DEFAULT_REDIS_DATABASES);
let max_db = keyspace_dbs.iter().copied().max().map(|db| db + 1).unwrap_or(0);
let max_db = keyspace_dbs
.iter()
.copied()
.max()
.map(|db| db + 1)
.unwrap_or(0);
let visible_count = database_count.max(max_db).max(1);
Ok((0..visible_count).collect())
@ -70,8 +86,14 @@ fn parse_database_count(value: redis::Value) -> Option<u32> {
})
}
async fn list_keyspace_databases(con: &mut redis::aio::MultiplexedConnection) -> Result<Vec<u32>, String> {
let info: String = redis::cmd("INFO").arg("keyspace").query_async(con).await.map_err(|e| e.to_string())?;
async fn list_keyspace_databases(
con: &mut redis::aio::MultiplexedConnection,
) -> Result<Vec<u32>, String> {
let info: String = redis::cmd("INFO")
.arg("keyspace")
.query_async(con)
.await
.map_err(|e| e.to_string())?;
let mut dbs = Vec::new();
for line in info.lines() {
@ -87,7 +109,11 @@ async fn list_keyspace_databases(con: &mut redis::aio::MultiplexedConnection) ->
}
pub async fn select_db(con: &mut redis::aio::MultiplexedConnection, db: u32) -> Result<(), String> {
redis::cmd("SELECT").arg(db).query_async(con).await.map_err(|e| e.to_string())
redis::cmd("SELECT")
.arg(db)
.query_async(con)
.await
.map_err(|e| e.to_string())
}
pub async fn scan_keys_page(
@ -108,18 +134,35 @@ pub async fn scan_keys_page(
let mut result = Vec::new();
for key in &keys {
let key_type: String =
redis::cmd("TYPE").arg(key.as_str()).query_async(con).await.unwrap_or_else(|_| "unknown".to_string());
let key_type: String = redis::cmd("TYPE")
.arg(key.as_str())
.query_async(con)
.await
.unwrap_or_else(|_| "unknown".to_string());
let ttl: i64 = con.ttl(key.as_str()).await.unwrap_or(-1);
result.push(RedisKeyInfo { key: key.clone(), key_type, ttl });
result.push(RedisKeyInfo {
key: key.clone(),
key_type,
ttl,
});
}
Ok(RedisScanResult { cursor: next_cursor, keys: result })
Ok(RedisScanResult {
cursor: next_cursor,
keys: result,
})
}
pub async fn get_value(con: &mut redis::aio::MultiplexedConnection, key: &str) -> Result<RedisValue, String> {
let key_type: String = redis::cmd("TYPE").arg(key).query_async(con).await.map_err(|e| e.to_string())?;
pub async fn get_value(
con: &mut redis::aio::MultiplexedConnection,
key: &str,
) -> Result<RedisValue, String> {
let key_type: String = redis::cmd("TYPE")
.arg(key)
.query_async(con)
.await
.map_err(|e| e.to_string())?;
let ttl: i64 = con.ttl(key).await.unwrap_or(-1);
@ -137,20 +180,33 @@ pub async fn get_value(con: &mut redis::aio::MultiplexedConnection, key: &str) -
serde_json::json!(v)
}
"zset" => {
let v: Vec<(String, f64)> = con.zrange_withscores(key, 0, -1).await.map_err(|e| e.to_string())?;
serde_json::json!(v.iter().map(|(m, s)| serde_json::json!({"member": m, "score": s})).collect::<Vec<_>>())
let v: Vec<(String, f64)> = con
.zrange_withscores(key, 0, -1)
.await
.map_err(|e| e.to_string())?;
serde_json::json!(v
.iter()
.map(|(m, s)| serde_json::json!({"member": m, "score": s}))
.collect::<Vec<_>>())
}
"hash" => {
let v: Vec<(String, String)> = con.hgetall(key).await.map_err(|e| e.to_string())?;
let map: serde_json::Map<String, serde_json::Value> =
v.into_iter().map(|(k, v)| (k, serde_json::Value::String(v))).collect();
let map: serde_json::Map<String, serde_json::Value> = v
.into_iter()
.map(|(k, v)| (k, serde_json::Value::String(v)))
.collect();
serde_json::Value::Object(map)
}
"stream" => get_stream_entries(con, key).await?,
_ => serde_json::Value::Null,
};
Ok(RedisValue { key: key.to_string(), key_type, ttl, value })
Ok(RedisValue {
key: key.to_string(),
key_type,
ttl,
value,
})
}
async fn get_stream_entries(
@ -228,16 +284,23 @@ pub async fn set_string(
value: &str,
ttl: Option<i64>,
) -> Result<(), String> {
con.set::<_, _, ()>(key, value).await.map_err(|e| e.to_string())?;
con.set::<_, _, ()>(key, value)
.await
.map_err(|e| e.to_string())?;
if let Some(t) = ttl {
if t > 0 {
con.expire::<_, ()>(key, t).await.map_err(|e| e.to_string())?;
con.expire::<_, ()>(key, t)
.await
.map_err(|e| e.to_string())?;
}
}
Ok(())
}
pub async fn delete_key(con: &mut redis::aio::MultiplexedConnection, key: &str) -> Result<(), String> {
pub async fn delete_key(
con: &mut redis::aio::MultiplexedConnection,
key: &str,
) -> Result<(), String> {
con.del::<_, ()>(key).await.map_err(|e| e.to_string())
}
@ -247,29 +310,67 @@ pub async fn hash_set(
field: &str,
value: &str,
) -> Result<(), String> {
con.hset::<_, _, _, ()>(key, field, value).await.map_err(|e| e.to_string())
con.hset::<_, _, _, ()>(key, field, value)
.await
.map_err(|e| e.to_string())
}
pub async fn hash_del(con: &mut redis::aio::MultiplexedConnection, key: &str, field: &str) -> Result<(), String> {
con.hdel::<_, _, ()>(key, field).await.map_err(|e| e.to_string())
pub async fn hash_del(
con: &mut redis::aio::MultiplexedConnection,
key: &str,
field: &str,
) -> Result<(), String> {
con.hdel::<_, _, ()>(key, field)
.await
.map_err(|e| e.to_string())
}
pub async fn list_push(con: &mut redis::aio::MultiplexedConnection, key: &str, value: &str) -> Result<(), String> {
con.rpush::<_, _, ()>(key, value).await.map_err(|e| e.to_string())
pub async fn list_push(
con: &mut redis::aio::MultiplexedConnection,
key: &str,
value: &str,
) -> Result<(), String> {
con.rpush::<_, _, ()>(key, value)
.await
.map_err(|e| e.to_string())
}
pub async fn list_remove(con: &mut redis::aio::MultiplexedConnection, key: &str, index: i64) -> Result<(), String> {
pub async fn list_remove(
con: &mut redis::aio::MultiplexedConnection,
key: &str,
index: i64,
) -> Result<(), String> {
let placeholder = "__DELETED_PLACEHOLDER__";
redis::cmd("LSET").arg(key).arg(index).arg(placeholder).query_async::<()>(con).await.map_err(|e| e.to_string())?;
con.lrem::<_, _, ()>(key, 1, placeholder).await.map_err(|e| e.to_string())
redis::cmd("LSET")
.arg(key)
.arg(index)
.arg(placeholder)
.query_async::<()>(con)
.await
.map_err(|e| e.to_string())?;
con.lrem::<_, _, ()>(key, 1, placeholder)
.await
.map_err(|e| e.to_string())
}
pub async fn set_add(con: &mut redis::aio::MultiplexedConnection, key: &str, member: &str) -> Result<(), String> {
con.sadd::<_, _, ()>(key, member).await.map_err(|e| e.to_string())
pub async fn set_add(
con: &mut redis::aio::MultiplexedConnection,
key: &str,
member: &str,
) -> Result<(), String> {
con.sadd::<_, _, ()>(key, member)
.await
.map_err(|e| e.to_string())
}
pub async fn set_remove(con: &mut redis::aio::MultiplexedConnection, key: &str, member: &str) -> Result<(), String> {
con.srem::<_, _, ()>(key, member).await.map_err(|e| e.to_string())
pub async fn set_remove(
con: &mut redis::aio::MultiplexedConnection,
key: &str,
member: &str,
) -> Result<(), String> {
con.srem::<_, _, ()>(key, member)
.await
.map_err(|e| e.to_string())
}
#[cfg(test)]
@ -284,7 +385,12 @@ mod tests {
fn parses_stream_entries() {
let raw = RedisRawValue::Array(vec![RedisRawValue::Array(vec![
bulk("1714470000000-0"),
RedisRawValue::Array(vec![bulk("event"), bulk("login"), bulk("user_id"), bulk("42")]),
RedisRawValue::Array(vec![
bulk("event"),
bulk("login"),
bulk("user_id"),
bulk("42"),
]),
])]);
let parsed = parse_stream_entries(raw);

74
crates/dbx-core/src/db/sqlite.rs
View File

@ -3,9 +3,15 @@ use sqlx::{Column, Executor, Row};
use std::time::{Duration, Instant};
use crate::types::{ColumnInfo, DatabaseInfo, ForeignKeyInfo, IndexInfo, QueryResult, TableInfo, TriggerInfo};
use super::file_validator::validate_file_path;
pub async fn connect_path(path: &str) -> Result<SqlitePool, String> {
let mut options = SqliteConnectOptions::new().filename(path).create_if_missing(true);
// Validate file path using universal validator
validate_file_path(path, is_network_path)?;
let mut options = SqliteConnectOptions::new()
.filename(path)
.create_if_missing(false);
if is_network_path(path) {
options = options.vfs("unix-nolock");
@ -49,8 +55,10 @@ pub async fn list_tables(pool: &SqlitePool, _schema: &str) -> Result<Vec<TableIn
}
pub async fn get_columns(pool: &SqlitePool, _schema: &str, table: &str) -> Result<Vec<ColumnInfo>, String> {
let rows: Vec<SqliteRow> =
sqlx::query(&format!("PRAGMA table_info(\"{}\")", table)).fetch_all(pool).await.map_err(|e| e.to_string())?;
let rows: Vec<SqliteRow> = sqlx::query(&format!("PRAGMA table_info(\"{}\")", table))
.fetch_all(pool)
.await
.map_err(|e| e.to_string())?;
Ok(rows
.iter()
@ -60,8 +68,7 @@ pub async fn get_columns(pool: &SqlitePool, _schema: &str, table: &str) -> Resul
is_nullable: row.get::<i32, _>("notnull") == 0,
column_default: row.get::<Option<String>, _>("dflt_value"),
is_primary_key: row.get::<i32, _>("pk") > 0,
extra: None,
comment: None,
extra: None, comment: None,
numeric_precision: None,
numeric_scale: None,
character_maximum_length: None,
@ -123,33 +130,26 @@ pub async fn list_foreign_keys(pool: &SqlitePool, _schema: &str, table: &str) ->
}
pub async fn list_triggers(pool: &SqlitePool, _schema: &str, table: &str) -> Result<Vec<TriggerInfo>, String> {
let rows: Vec<SqliteRow> =
sqlx::query("SELECT name, sql FROM sqlite_master WHERE type = 'trigger' AND tbl_name = ? ORDER BY name")
.bind(table)
.fetch_all(pool)
.await
.map_err(|e| e.to_string())?;
let rows: Vec<SqliteRow> = sqlx::query(
"SELECT name, sql FROM sqlite_master WHERE type = 'trigger' AND tbl_name = ? ORDER BY name",
)
.bind(table)
.fetch_all(pool)
.await
.map_err(|e| e.to_string())?;
Ok(rows
.iter()
.map(|row| {
let sql_text: String = row.get::<Option<String>, _>("sql").unwrap_or_default();
let upper = sql_text.to_uppercase();
let timing = if upper.contains("BEFORE") {
"BEFORE"
} else if upper.contains("AFTER") {
"AFTER"
} else {
"INSTEAD OF"
};
let event = if upper.contains("INSERT") {
"INSERT"
} else if upper.contains("UPDATE") {
"UPDATE"
} else {
"DELETE"
};
TriggerInfo { name: row.get::<String, _>("name"), event: event.to_string(), timing: timing.to_string() }
let timing = if upper.contains("BEFORE") { "BEFORE" } else if upper.contains("AFTER") { "AFTER" } else { "INSTEAD OF" };
let event = if upper.contains("INSERT") { "INSERT" } else if upper.contains("UPDATE") { "UPDATE" } else { "DELETE" };
TriggerInfo {
name: row.get::<String, _>("name"),
event: event.to_string(),
timing: timing.to_string(),
}
})
.collect())
}
@ -166,7 +166,10 @@ pub async fn execute_query(pool: &SqlitePool, sql: &str) -> Result<QueryResult,
let desc = pool.describe(sql).await.map_err(|e| e.to_string())?;
let columns: Vec<String> = desc.columns().iter().map(|c| c.name().to_string()).collect();
let rows: Vec<SqliteRow> = sqlx::query(sql).fetch_all(pool).await.map_err(|e| e.to_string())?;
let rows: Vec<SqliteRow> = sqlx::query(sql)
.fetch_all(pool)
.await
.map_err(|e| e.to_string())?;
let result_rows: Vec<Vec<serde_json::Value>> = rows
.iter()
@ -176,13 +179,11 @@ pub async fn execute_query(pool: &SqlitePool, sql: &str) -> Result<QueryResult,
row.try_get::<String, _>(i)
.map(serde_json::Value::String)
.or_else(|_| row.try_get::<i64, _>(i).map(|v| serde_json::Value::Number(v.into())))
.or_else(|_| {
row.try_get::<f64, _>(i).map(|v| {
serde_json::Number::from_f64(v)
.map(serde_json::Value::Number)
.unwrap_or(serde_json::Value::Null)
})
})
.or_else(|_| row.try_get::<f64, _>(i).map(|v| {
serde_json::Number::from_f64(v)
.map(serde_json::Value::Number)
.unwrap_or(serde_json::Value::Null)
}))
.or_else(|_| row.try_get::<bool, _>(i).map(serde_json::Value::Bool))
.unwrap_or(serde_json::Value::Null)
})
@ -198,7 +199,10 @@ pub async fn execute_query(pool: &SqlitePool, sql: &str) -> Result<QueryResult,
truncated: false,
})
} else {
let result = sqlx::query(sql).execute(pool).await.map_err(|e| e.to_string())?;
let result = sqlx::query(sql)
.execute(pool)
.await
.map_err(|e| e.to_string())?;
Ok(QueryResult {
columns: vec![],

5
crates/dbx-core/src/db/ssh_tunnel.rs
View File

@ -9,7 +9,7 @@ use tokio::net::TcpListener;
use tokio::sync::Mutex;
use tokio::task::JoinHandle;
use super::{connection_timeout, CONNECTION_TIMEOUT_SECS};
use super::{connection_timeout, CONNECTION_TIMEOUT_SECS, file_validator::validate_file_path};
struct SshClient;
@ -41,6 +41,9 @@ async fn connect_and_authenticate(
.map_err(|e| format!("SSH connection failed: {e}"))?;
if !ssh_key_path.is_empty() {
// Validate SSH key file path
validate_file_path(ssh_key_path, |_| false)?;
let passphrase = if ssh_key_passphrase.is_empty() { None } else { Some(ssh_key_passphrase) };
let key_pair = load_secret_key(ssh_key_path, passphrase).map_err(|e| format!("Failed to load SSH key: {e}"))?;
let auth_res = tokio::time::timeout(