mirror of
https://github.com/rustfs/rustfs.git
synced 2026-08-07 13:53:12 +00:00
fix(s3select): enforce query and resource limits (#5028)
* fix(s3select): enforce query and resource limits * fix(s3select): close query resource limit gaps * fix(s3select): preserve timeout and stream invariants * fix(s3select): enforce staged query limits * fix(s3select): preserve policy error compatibility * fix(s3select): bound error source traversal
This commit is contained in:
GatewayJ
@@ -12,11 +12,18 @@
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
use std::ops::ControlFlow;
|
||||
|
||||
use async_recursion::async_recursion;
|
||||
use async_trait::async_trait;
|
||||
use datafusion::sql::{planner::SqlToRel, sqlparser::ast::Statement};
|
||||
use datafusion::sql::{
|
||||
planner::SqlToRel,
|
||||
sqlparser::ast::{
|
||||
GroupByExpr, ObjectNamePart, OrderByKind, Query, Select, SelectFlavor, SetExpr, Statement, TableFactor, Visit, Visitor,
|
||||
},
|
||||
};
|
||||
use rustfs_s3select_api::{
|
||||
QueryError, QueryResult,
|
||||
QueryError, QueryResult, S3SelectPolicyError,
|
||||
query::{
|
||||
ast::ExtStatement,
|
||||
logical_planner::{LogicalPlanner, Plan, QueryPlan},
|
||||
@@ -60,6 +67,7 @@ impl<'a, S: ContextProviderExtension + Send + Sync + 'a> SqlPlanner<'a, S> {
|
||||
async fn df_sql_to_plan(&self, stmt: Statement, _session: &SessionCtx) -> QueryResult<Plan> {
|
||||
match stmt {
|
||||
Statement::Query(_) => {
|
||||
validate_s3_select_statement(&stmt)?;
|
||||
let df_plan = self.df_planner.sql_statement_to_plan(stmt)?;
|
||||
let plan = Plan::Query(QueryPlan {
|
||||
df_plan,
|
||||
@@ -72,3 +80,256 @@ impl<'a, S: ContextProviderExtension + Send + Sync + 'a> SqlPlanner<'a, S> {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn validate_s3_select_statement(statement: &Statement) -> QueryResult<()> {
|
||||
let Statement::Query(query) = statement else {
|
||||
return Err(unsupported_structure("only SELECT queries are supported"));
|
||||
};
|
||||
|
||||
if query.with.is_some()
|
||||
|| query.order_by.as_ref().is_some_and(|order_by| {
|
||||
order_by.interpolate.is_some()
|
||||
|| match &order_by.kind {
|
||||
OrderByKind::Expressions(expressions) => expressions.iter().any(|expression| expression.with_fill.is_some()),
|
||||
OrderByKind::All(_) => true,
|
||||
}
|
||||
})
|
||||
|| query.fetch.is_some()
|
||||
|| !query.locks.is_empty()
|
||||
|| query.for_clause.is_some()
|
||||
|| query.settings.is_some()
|
||||
|| query.format_clause.is_some()
|
||||
|| !query.pipe_operators.is_empty()
|
||||
{
|
||||
return Err(unsupported_structure("the query contains an unsupported clause"));
|
||||
}
|
||||
if let Some(limit_clause) = query.limit_clause.as_ref()
|
||||
&& !matches!(
|
||||
limit_clause,
|
||||
datafusion::sql::sqlparser::ast::LimitClause::LimitOffset {
|
||||
limit: Some(_),
|
||||
offset: None,
|
||||
limit_by,
|
||||
} if limit_by.is_empty()
|
||||
)
|
||||
{
|
||||
return Err(unsupported_structure("only LIMIT without OFFSET is supported"));
|
||||
}
|
||||
if let Some(datafusion::sql::sqlparser::ast::LimitClause::LimitOffset { limit: Some(limit), .. }) =
|
||||
query.limit_clause.as_ref()
|
||||
&& limit.to_string().parse::<u64>().is_err()
|
||||
{
|
||||
return Err(unsupported_structure("LIMIT must be a non-negative integer"));
|
||||
}
|
||||
|
||||
let mut detector = SubqueryDetector { visited_root: false };
|
||||
if query.visit(&mut detector).is_break() {
|
||||
return Err(unsupported_structure("subqueries are not supported"));
|
||||
}
|
||||
|
||||
let SetExpr::Select(select) = query.body.as_ref() else {
|
||||
return Err(unsupported_structure("set operations and nested queries are not supported"));
|
||||
};
|
||||
validate_select(select)
|
||||
}
|
||||
|
||||
fn validate_select(select: &Select) -> QueryResult<()> {
|
||||
if !select.optimizer_hints.is_empty()
|
||||
|| select.distinct.is_some()
|
||||
|| select.select_modifiers.is_some()
|
||||
|| select.top.is_some()
|
||||
|| select.exclude.is_some()
|
||||
|| select.into.is_some()
|
||||
|| !select.lateral_views.is_empty()
|
||||
|| select.prewhere.is_some()
|
||||
|| !select.connect_by.is_empty()
|
||||
|| !select.cluster_by.is_empty()
|
||||
|| !select.distribute_by.is_empty()
|
||||
|| !select.sort_by.is_empty()
|
||||
|| select.having.is_some()
|
||||
|| !select.named_window.is_empty()
|
||||
|| select.qualify.is_some()
|
||||
|| select.value_table_mode.is_some()
|
||||
|| select.flavor != SelectFlavor::Standard
|
||||
|| !matches!(&select.group_by, GroupByExpr::Expressions(_, modifiers) if modifiers.is_empty())
|
||||
{
|
||||
return Err(unsupported_structure("the SELECT contains an unsupported clause"));
|
||||
}
|
||||
|
||||
let [table] = select.from.as_slice() else {
|
||||
return Err(unsupported_structure("exactly one S3Object source is required"));
|
||||
};
|
||||
if !table.joins.is_empty() {
|
||||
return Err(unsupported_structure("JOIN is not supported"));
|
||||
}
|
||||
let TableFactor::Table {
|
||||
name,
|
||||
alias,
|
||||
args,
|
||||
with_hints,
|
||||
version,
|
||||
with_ordinality,
|
||||
partitions,
|
||||
sample,
|
||||
index_hints,
|
||||
..
|
||||
} = &table.relation
|
||||
else {
|
||||
return Err(unsupported_structure("subqueries and table functions are not supported"));
|
||||
};
|
||||
if args.is_some()
|
||||
|| !with_hints.is_empty()
|
||||
|| version.is_some()
|
||||
|| *with_ordinality
|
||||
|| !partitions.is_empty()
|
||||
|| sample.is_some()
|
||||
|| !index_hints.is_empty()
|
||||
|| alias.as_ref().is_some_and(|alias| !alias.columns.is_empty())
|
||||
{
|
||||
return Err(unsupported_structure("the S3Object source contains unsupported modifiers"));
|
||||
}
|
||||
let ([ObjectNamePart::Identifier(table_name)] | [ObjectNamePart::Identifier(table_name), ObjectNamePart::Identifier(_)]) =
|
||||
name.0.as_slice()
|
||||
else {
|
||||
return Err(unsupported_structure("the source must be S3Object"));
|
||||
};
|
||||
let is_s3_object = if table_name.quote_style.is_some() {
|
||||
table_name.value == "S3Object"
|
||||
} else {
|
||||
table_name.value.eq_ignore_ascii_case("S3Object")
|
||||
};
|
||||
if !is_s3_object {
|
||||
return Err(unsupported_structure("the source must be S3Object"));
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn unsupported_structure(message: &str) -> QueryError {
|
||||
S3SelectPolicyError::UnsupportedSqlStructure {
|
||||
message: message.to_string(),
|
||||
}
|
||||
.into()
|
||||
}
|
||||
|
||||
struct SubqueryDetector {
|
||||
visited_root: bool,
|
||||
}
|
||||
|
||||
impl Visitor for SubqueryDetector {
|
||||
type Break = ();
|
||||
|
||||
fn pre_visit_query(&mut self, _query: &Query) -> ControlFlow<Self::Break> {
|
||||
if self.visited_root {
|
||||
ControlFlow::Break(())
|
||||
} else {
|
||||
self.visited_root = true;
|
||||
ControlFlow::Continue(())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::validate_s3_select_statement;
|
||||
use crate::sql::parser::ExtParser;
|
||||
use datafusion::sql::sqlparser::ast::Statement;
|
||||
use rustfs_s3select_api::{S3SelectPolicyError, query::ast::ExtStatement};
|
||||
|
||||
fn parse_statement(sql: &str) -> Statement {
|
||||
let mut statements = ExtParser::parse_sql(sql).expect("SQL should parse");
|
||||
let ExtStatement::SqlStatement(statement) = statements.pop_front().expect("one SQL statement");
|
||||
*statement
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn accepts_s3_select_query_shape() {
|
||||
let statement = parse_statement("SELECT s.id FROM S3Object AS s WHERE s.id = '1' LIMIT 10");
|
||||
|
||||
assert!(validate_s3_select_statement(&statement).is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn accepts_json_sub_path_source() {
|
||||
let statement = parse_statement("SELECT e.name FROM S3Object.employees AS e");
|
||||
|
||||
assert!(validate_s3_select_statement(&statement).is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn accepts_group_by_and_order_by() {
|
||||
let statement = parse_statement("SELECT department, COUNT(*) FROM S3Object GROUP BY department ORDER BY department");
|
||||
|
||||
assert!(validate_s3_select_statement(&statement).is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_join() {
|
||||
let statement = parse_statement("SELECT * FROM S3Object a JOIN S3Object b ON a.id = b.id");
|
||||
|
||||
assert!(matches!(
|
||||
validate_s3_select_statement(&statement),
|
||||
Err(ref err) if matches!(
|
||||
err.s3_select_policy_error(),
|
||||
Some(S3SelectPolicyError::UnsupportedSqlStructure { message }) if message == "JOIN is not supported"
|
||||
)
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_subquery() {
|
||||
let statement = parse_statement("SELECT * FROM S3Object WHERE id IN (SELECT id FROM S3Object)");
|
||||
|
||||
assert!(matches!(
|
||||
validate_s3_select_statement(&statement),
|
||||
Err(ref err) if matches!(
|
||||
err.s3_select_policy_error(),
|
||||
Some(S3SelectPolicyError::UnsupportedSqlStructure { message }) if message == "subqueries are not supported"
|
||||
)
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_subquery_in_order_by() {
|
||||
let statement = parse_statement("SELECT id FROM S3Object ORDER BY (SELECT id FROM S3Object)");
|
||||
|
||||
assert!(matches!(
|
||||
validate_s3_select_statement(&statement),
|
||||
Err(ref err) if matches!(
|
||||
err.s3_select_policy_error(),
|
||||
Some(S3SelectPolicyError::UnsupportedSqlStructure { message }) if message == "subqueries are not supported"
|
||||
)
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_non_s3_object_source() {
|
||||
let statement = parse_statement("SELECT * FROM other_table");
|
||||
|
||||
assert!(matches!(
|
||||
validate_s3_select_statement(&statement),
|
||||
Err(ref err) if matches!(
|
||||
err.s3_select_policy_error(),
|
||||
Some(S3SelectPolicyError::UnsupportedSqlStructure { message }) if message == "the source must be S3Object"
|
||||
)
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_unsupported_select_clauses() {
|
||||
for sql in [
|
||||
"SELECT DISTINCT id FROM S3Object",
|
||||
"SELECT * FROM S3Object OFFSET 1",
|
||||
"SELECT * FROM S3Object UNION SELECT * FROM S3Object",
|
||||
] {
|
||||
let statement = parse_statement(sql);
|
||||
assert!(
|
||||
matches!(
|
||||
validate_s3_select_statement(&statement),
|
||||
Err(ref err) if matches!(err.s3_select_policy_error(), Some(S3SelectPolicyError::UnsupportedSqlStructure { .. }))
|
||||
),
|
||||
"query should be rejected: {sql}"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user