polyglot_sql/
guard.rs

1//! Shared complexity guards for recursion-heavy SQL operations.
2
3use crate::error::{Error, Result};
4use crate::expressions::Expression;
5use crate::tokens::{Token, TokenType};
6use serde::{Deserialize, Serialize};
7
8const DEFAULT_MAX_INPUT_BYTES: usize = 16 * 1024 * 1024;
9const DEFAULT_MAX_TOKENS: usize = 1_000_000;
10const DEFAULT_MAX_AST_NODES: usize = 1_000_000;
11const DEFAULT_MAX_AST_DEPTH: usize = 512;
12const DEFAULT_MAX_PARENTHESES_DEPTH: usize = 512;
13const DEFAULT_MAX_FUNCTION_CALL_DEPTH: usize = 64;
14
15fn default_max_input_bytes() -> Option<usize> {
16    Some(DEFAULT_MAX_INPUT_BYTES)
17}
18
19fn default_max_tokens() -> Option<usize> {
20    Some(DEFAULT_MAX_TOKENS)
21}
22
23fn default_max_ast_nodes() -> Option<usize> {
24    Some(DEFAULT_MAX_AST_NODES)
25}
26
27fn default_max_ast_depth() -> Option<usize> {
28    Some(DEFAULT_MAX_AST_DEPTH)
29}
30
31fn default_max_parenthesis_depth() -> Option<usize> {
32    Some(DEFAULT_MAX_PARENTHESES_DEPTH)
33}
34
35fn default_max_function_call_depth() -> Option<usize> {
36    Some(DEFAULT_MAX_FUNCTION_CALL_DEPTH)
37}
38
39/// Guard options for parse/transpile/generate complexity.
40///
41/// These limits turn excessively deep or large inputs into regular errors
42/// instead of relying on process stack exhaustion as the failure mode.
43#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
44#[serde(rename_all = "camelCase")]
45pub struct ComplexityGuardOptions {
46    /// Maximum allowed SQL input size in bytes.
47    /// `None` disables this check.
48    #[serde(default = "default_max_input_bytes")]
49    pub max_input_bytes: Option<usize>,
50    /// Maximum allowed number of tokens after tokenization.
51    /// `None` disables this check.
52    #[serde(default = "default_max_tokens")]
53    pub max_tokens: Option<usize>,
54    /// Maximum allowed AST node count after parsing.
55    /// `None` disables this check.
56    #[serde(default = "default_max_ast_nodes")]
57    pub max_ast_nodes: Option<usize>,
58    /// Maximum allowed AST depth after parsing.
59    /// `None` disables this check.
60    #[serde(default = "default_max_ast_depth")]
61    pub max_ast_depth: Option<usize>,
62    /// Maximum allowed nested parenthesis depth before parsing.
63    /// `None` disables this check.
64    #[serde(default = "default_max_parenthesis_depth")]
65    pub max_parenthesis_depth: Option<usize>,
66    /// Maximum allowed nested function-call depth before parsing.
67    /// `None` disables this check.
68    #[serde(default = "default_max_function_call_depth")]
69    pub max_function_call_depth: Option<usize>,
70}
71
72impl Default for ComplexityGuardOptions {
73    fn default() -> Self {
74        Self {
75            max_input_bytes: default_max_input_bytes(),
76            max_tokens: default_max_tokens(),
77            max_ast_nodes: default_max_ast_nodes(),
78            max_ast_depth: default_max_ast_depth(),
79            max_parenthesis_depth: default_max_parenthesis_depth(),
80            max_function_call_depth: default_max_function_call_depth(),
81        }
82    }
83}
84
85fn parse_guard_error(code: &str, actual: usize, limit: usize, token: Option<&Token>) -> Error {
86    let message = format!("{code}: value {actual} exceeds configured limit {limit}");
87    if let Some(token) = token {
88        Error::parse(
89            message,
90            token.span.line,
91            token.span.column,
92            token.span.start,
93            token.span.end,
94        )
95    } else {
96        Error::parse(message, 0, 0, 0, 0)
97    }
98}
99
100fn generate_guard_error(code: &str, actual: usize, limit: usize) -> Error {
101    Error::generate(format!(
102        "{code}: value {actual} exceeds configured limit {limit}"
103    ))
104}
105
106/// Enforce raw SQL input limits before tokenization.
107pub fn enforce_input(sql: &str, options: &ComplexityGuardOptions) -> Result<()> {
108    if let Some(max) = options.max_input_bytes {
109        let input_bytes = sql.len();
110        if input_bytes > max {
111            return Err(parse_guard_error(
112                "E_GUARD_INPUT_TOO_LARGE",
113                input_bytes,
114                max,
115                None,
116            ));
117        }
118    }
119
120    Ok(())
121}
122
123/// Enforce token and pre-parse nesting limits.
124pub fn enforce_tokens(tokens: &[Token], options: &ComplexityGuardOptions) -> Result<()> {
125    if let Some(max) = options.max_tokens {
126        let token_count = tokens.len();
127        if token_count > max {
128            let token = tokens.get(max).or_else(|| tokens.last());
129            return Err(parse_guard_error(
130                "E_GUARD_TOKEN_BUDGET_EXCEEDED",
131                token_count,
132                max,
133                token,
134            ));
135        }
136    }
137
138    if options.max_parenthesis_depth.is_some() || options.max_function_call_depth.is_some() {
139        let mut paren_depth = 0usize;
140        let mut function_depth = 0usize;
141        let mut paren_stack = Vec::new();
142        let mut previous_significant: Option<TokenType> = None;
143
144        for token in tokens {
145            match token.token_type {
146                TokenType::LParen => {
147                    paren_depth += 1;
148                    if let Some(max) = options.max_parenthesis_depth {
149                        if paren_depth > max {
150                            return Err(parse_guard_error(
151                                "E_GUARD_NESTING_DEPTH_EXCEEDED",
152                                paren_depth,
153                                max,
154                                Some(token),
155                            ));
156                        }
157                    }
158
159                    let is_function_call = previous_significant
160                        .map(is_function_call_name_token)
161                        .unwrap_or(false);
162                    paren_stack.push(is_function_call);
163                    if is_function_call {
164                        function_depth += 1;
165                        if let Some(max) = options.max_function_call_depth {
166                            if function_depth > max {
167                                return Err(parse_guard_error(
168                                    "E_GUARD_FUNCTION_NESTING_DEPTH_EXCEEDED",
169                                    function_depth,
170                                    max,
171                                    Some(token),
172                                ));
173                            }
174                        }
175                    }
176                }
177                TokenType::RParen => {
178                    paren_depth = paren_depth.saturating_sub(1);
179                    if paren_stack.pop().unwrap_or(false) {
180                        function_depth = function_depth.saturating_sub(1);
181                    }
182                }
183                _ => {}
184            }
185
186            if !is_trivia_token(token.token_type) {
187                previous_significant = Some(token.token_type);
188            }
189        }
190    }
191
192    Ok(())
193}
194
195fn is_trivia_token(token_type: TokenType) -> bool {
196    matches!(
197        token_type,
198        TokenType::Space | TokenType::Break | TokenType::LineComment | TokenType::BlockComment
199    )
200}
201
202fn is_function_call_name_token(token_type: TokenType) -> bool {
203    matches!(
204        token_type,
205        TokenType::Identifier
206            | TokenType::Var
207            | TokenType::QuotedIdentifier
208            | TokenType::CurrentDate
209            | TokenType::CurrentDateTime
210            | TokenType::CurrentTime
211            | TokenType::CurrentTimestamp
212            | TokenType::CurrentUser
213            | TokenType::If
214            | TokenType::Index
215            | TokenType::Insert
216            | TokenType::Left
217            | TokenType::Replace
218            | TokenType::Right
219            | TokenType::Row
220    )
221}
222
223/// Enforce AST size/depth limits and report parse-oriented errors.
224pub fn enforce_ast(expressions: &[Expression], options: &ComplexityGuardOptions) -> Result<()> {
225    let Some(max_nodes) = options.max_ast_nodes else {
226        return enforce_ast_depth_only(expressions, options);
227    };
228
229    let stats = ast_stats(expressions, Some(max_nodes), options.max_ast_depth)?;
230    if stats.node_count > max_nodes {
231        return Err(parse_guard_error(
232            "E_GUARD_AST_BUDGET_EXCEEDED",
233            stats.node_count,
234            max_nodes,
235            None,
236        ));
237    }
238
239    if let Some(max_depth) = options.max_ast_depth {
240        if stats.max_depth > max_depth {
241            return Err(parse_guard_error(
242                "E_GUARD_AST_DEPTH_EXCEEDED",
243                stats.max_depth,
244                max_depth,
245                None,
246            ));
247        }
248    }
249
250    Ok(())
251}
252
253/// Enforce AST size/depth limits and report generation-oriented errors.
254pub fn enforce_generate_ast(
255    expression: &Expression,
256    options: &ComplexityGuardOptions,
257) -> Result<()> {
258    let stats = ast_stats(
259        std::slice::from_ref(expression),
260        options.max_ast_nodes,
261        options.max_ast_depth,
262    )?;
263
264    if let Some(max_nodes) = options.max_ast_nodes {
265        if stats.node_count > max_nodes {
266            return Err(generate_guard_error(
267                "E_GUARD_AST_BUDGET_EXCEEDED",
268                stats.node_count,
269                max_nodes,
270            ));
271        }
272    }
273
274    if let Some(max_depth) = options.max_ast_depth {
275        if stats.max_depth > max_depth {
276            return Err(generate_guard_error(
277                "E_GUARD_AST_DEPTH_EXCEEDED",
278                stats.max_depth,
279                max_depth,
280            ));
281        }
282    }
283
284    Ok(())
285}
286
287fn enforce_ast_depth_only(
288    expressions: &[Expression],
289    options: &ComplexityGuardOptions,
290) -> Result<()> {
291    let Some(max_depth) = options.max_ast_depth else {
292        return Ok(());
293    };
294
295    let stats = ast_stats(expressions, None, Some(max_depth))?;
296    if stats.max_depth > max_depth {
297        return Err(parse_guard_error(
298            "E_GUARD_AST_DEPTH_EXCEEDED",
299            stats.max_depth,
300            max_depth,
301            None,
302        ));
303    }
304
305    Ok(())
306}
307
308#[derive(Debug, Clone, Copy, Default)]
309struct AstStats {
310    node_count: usize,
311    max_depth: usize,
312}
313
314fn ast_stats(
315    expressions: &[Expression],
316    max_nodes: Option<usize>,
317    max_depth: Option<usize>,
318) -> Result<AstStats> {
319    let mut stats = AstStats::default();
320    let mut stack: Vec<(&Expression, usize)> = expressions.iter().rev().map(|e| (e, 0)).collect();
321
322    while let Some((expr, depth)) = stack.pop() {
323        stats.node_count += 1;
324        stats.max_depth = stats.max_depth.max(depth);
325
326        if let Some(max) = max_nodes {
327            if stats.node_count > max {
328                return Ok(stats);
329            }
330        }
331
332        if let Some(max) = max_depth {
333            if stats.max_depth > max {
334                return Ok(stats);
335            }
336        }
337
338        push_ast_children(&mut stack, expr, depth);
339    }
340
341    Ok(stats)
342}
343
344fn push_ast_children<'a>(
345    stack: &mut Vec<(&'a Expression, usize)>,
346    expr: &'a Expression,
347    depth: usize,
348) {
349    match expr {
350        Expression::And(op) if is_commentless_binary_op(op) => {
351            push_connector_child(stack, &op.right, depth, ConnectorKind::And);
352            push_connector_child(stack, &op.left, depth, ConnectorKind::And);
353        }
354        Expression::Or(op) if is_commentless_binary_op(op) => {
355            push_connector_child(stack, &op.right, depth, ConnectorKind::Or);
356            push_connector_child(stack, &op.left, depth, ConnectorKind::Or);
357        }
358        _ => {
359            use crate::traversal::ExpressionWalk;
360
361            for child in expr.children().into_iter().rev() {
362                stack.push((child, depth + 1));
363            }
364        }
365    }
366}
367
368fn push_connector_child<'a>(
369    stack: &mut Vec<(&'a Expression, usize)>,
370    child: &'a Expression,
371    depth: usize,
372    kind: ConnectorKind,
373) {
374    let child_depth = if is_commentless_connector(child, kind) {
375        depth
376    } else {
377        depth + 1
378    };
379    stack.push((child, child_depth));
380}
381
382fn is_commentless_connector(expr: &Expression, kind: ConnectorKind) -> bool {
383    match (kind, expr) {
384        (ConnectorKind::And, Expression::And(op)) | (ConnectorKind::Or, Expression::Or(op)) => {
385            is_commentless_binary_op(op)
386        }
387        _ => false,
388    }
389}
390
391fn is_commentless_binary_op(op: &crate::expressions::BinaryOp) -> bool {
392    op.left_comments.is_empty()
393        && op.operator_comments.is_empty()
394        && op.trailing_comments.is_empty()
395}
396
397#[derive(Debug, Clone, Copy)]
398enum ConnectorKind {
399    And,
400    Or,
401}
polyglot_sql/guard.rs

polyglot_sql/
guard.rs
