github/DB-GPT
1
0
Fork 0
mirror of https://github.com/csunny/DB-GPT.git synced 2025-09-13 13:10:29 +00:00
Code Issues Packages Projects Releases Wiki Activity

[BUG] SQL Injection through CVE Bypass in DB-GPT 0.7.0 (CVE-2024-10835 & CVE-2024-10901) (#2650)

Browse Source 
Co-authored-by: nkoorty <amalyshau2002@gmail.com>
Co-authored-by: Fangyin Cheng <staneyffer@gmail.com>
This commit is contained in:
Gecko Security
2025-04-28 07:54:41 +02:00
committed by GitHub
parent b16c6793ec
commit bcb43266cf
2 changed files with 149 additions and 137 deletions
Download Patch File Download Diff File Expand all files Collapse all files

35
packages/dbgpt-core/src/dbgpt/datasource/rdbms/base.py
View File

@@ -477,7 +477,11 @@ class RDBMSConnector(BaseConnector):
return self._query(sql)
def query_ex(
self, query: str, fetch: str = "all", timeout: Optional[float] = None
self,
query: str,
params: Optional[Dict[str, Any]] = None,
fetch: str = "all",
timeout: Optional[float] = None,
) -> Tuple[List[str], Optional[List]]:
"""Execute a SQL command and return the results with optional timeout.
@@ -485,6 +489,7 @@ class RDBMSConnector(BaseConnector):
Args:
query (str): SQL query to run
params (Optional[dict]): Parameters for the query
fetch (str): fetch type, either 'all' or 'one'
timeout (Optional[float]): Query timeout in seconds. If None, no timeout is
applied.
@@ -501,13 +506,21 @@ class RDBMSConnector(BaseConnector):
return [], None
query = self._format_sql(query)
def _execute_query(session, sql_text):
cursor = session.execute(sql_text)
# Initialize params if None
if params is None:
params = {}
def _execute_query(session, sql_text, query_params):
cursor = session.execute(sql_text, query_params)
if cursor.returns_rows:
if fetch == "all":
result = cursor.fetchall()
elif fetch == "one":
result = cursor.fetchone()
if result:
result = [result]
else:
result = []
else:
raise ValueError("Fetch parameter must be either 'one' or 'all'")
field_names = list(cursor.keys())
@@ -526,14 +539,14 @@ class RDBMSConnector(BaseConnector):
session.execute(
text(f"SET SESSION MAX_EXECUTION_TIME = {mysql_timeout}")
)
return _execute_query(session, sql)
return _execute_query(session, sql, params)
elif self.dialect == "postgresql":
# PostgreSQL: Set statement_timeout in milliseconds
session.execute(
text(f"SET statement_timeout = {int(timeout * 1000)}")
)
return _execute_query(session, sql)
return _execute_query(session, sql, params)
elif self.dialect == "oceanbase":
# OceanBase: Set ob_query_timeout in microseconds
@@ -541,17 +554,19 @@ class RDBMSConnector(BaseConnector):
session.execute(
text(f"SET SESSION ob_query_timeout = {ob_timeout}")
)
return _execute_query(session, sql)
return _execute_query(session, sql, params)
elif self.dialect == "mssql":
# MSSQL: Use execution_options if supported by driver
sql_with_timeout = sql.execution_options(timeout=int(timeout))
return _execute_query(session, sql_with_timeout)
return _execute_query(session, sql_with_timeout, params)
elif self.dialect == "duckdb":
# DuckDB: Use ThreadPoolExecutor for timeout
with ThreadPoolExecutor(max_workers=1) as executor:
future = executor.submit(_execute_query, session, sql)
future = executor.submit(
_execute_query, session, sql, params
)
try:
return future.result(timeout=timeout)
except FutureTimeoutError:
@@ -564,10 +579,10 @@ class RDBMSConnector(BaseConnector):
f"Timeout not supported for dialect: {self.dialect}, "
"proceeding without timeout"
)
return _execute_query(session, sql)
return _execute_query(session, sql, params)
# No timeout specified, execute normally
return _execute_query(session, sql)
return _execute_query(session, sql, params)
except SQLAlchemyError as e:
if "timeout" in str(e).lower() or "timed out" in str(e).lower():