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::{ParserToken, Span, 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
72#[derive(Debug, Default)]
73pub(crate) struct TokenGuardStats {
74    pub token_count: usize,
75    parenthesis_depth_spans: Vec<Span>,
76    function_depth_spans: Vec<Span>,
77    parenthesis_stack: Vec<bool>,
78    parenthesis_depth: usize,
79    function_depth: usize,
80    previous_significant: Option<TokenType>,
81}
82
83impl TokenGuardStats {
84    pub(crate) fn observe(&mut self, token_type: TokenType, span: Span) {
85        self.token_count += 1;
86        match token_type {
87            TokenType::LParen => {
88                self.parenthesis_depth += 1;
89                if self.parenthesis_depth_spans.len() < self.parenthesis_depth {
90                    self.parenthesis_depth_spans.push(span);
91                }
92
93                let is_function_call = self
94                    .previous_significant
95                    .map(is_function_call_name_token)
96                    .unwrap_or(false);
97                self.parenthesis_stack.push(is_function_call);
98                if is_function_call {
99                    self.function_depth += 1;
100                    if self.function_depth_spans.len() < self.function_depth {
101                        self.function_depth_spans.push(span);
102                    }
103                }
104            }
105            TokenType::RParen => {
106                self.parenthesis_depth = self.parenthesis_depth.saturating_sub(1);
107                if self.parenthesis_stack.pop().unwrap_or(false) {
108                    self.function_depth = self.function_depth.saturating_sub(1);
109                }
110            }
111            _ => {}
112        }
113
114        if !is_trivia_token(token_type) {
115            self.previous_significant = Some(token_type);
116        }
117    }
118}
119
120impl Default for ComplexityGuardOptions {
121    fn default() -> Self {
122        Self {
123            max_input_bytes: default_max_input_bytes(),
124            max_tokens: default_max_tokens(),
125            max_ast_nodes: default_max_ast_nodes(),
126            max_ast_depth: default_max_ast_depth(),
127            max_parenthesis_depth: default_max_parenthesis_depth(),
128            max_function_call_depth: default_max_function_call_depth(),
129        }
130    }
131}
132
133fn parse_guard_error(code: &str, actual: usize, limit: usize, span: Option<Span>) -> Error {
134    let message = format!("{code}: value {actual} exceeds configured limit {limit}");
135    if let Some(span) = span {
136        Error::parse(message, span.line, span.column, span.start, span.end)
137    } else {
138        Error::parse(message, 0, 0, 0, 0)
139    }
140}
141
142fn generate_guard_error(code: &str, actual: usize, limit: usize) -> Error {
143    Error::generate(format!(
144        "{code}: value {actual} exceeds configured limit {limit}"
145    ))
146}
147
148/// Enforce raw SQL input limits before tokenization.
149pub fn enforce_input(sql: &str, options: &ComplexityGuardOptions) -> Result<()> {
150    if let Some(max) = options.max_input_bytes {
151        let input_bytes = sql.len();
152        if input_bytes > max {
153            return Err(parse_guard_error(
154                "E_GUARD_INPUT_TOO_LARGE",
155                input_bytes,
156                max,
157                None,
158            ));
159        }
160    }
161
162    Ok(())
163}
164
165/// Enforce token and pre-parse nesting limits.
166trait GuardToken {
167    fn token_type(&self) -> TokenType;
168    fn span(&self) -> Span;
169}
170
171impl GuardToken for Token {
172    fn token_type(&self) -> TokenType {
173        self.token_type
174    }
175
176    fn span(&self) -> Span {
177        self.span
178    }
179}
180
181impl GuardToken for ParserToken {
182    fn token_type(&self) -> TokenType {
183        self.token_type
184    }
185
186    fn span(&self) -> Span {
187        self.span
188    }
189}
190
191fn enforce_token_slice<T: GuardToken>(
192    tokens: &[T],
193    options: &ComplexityGuardOptions,
194) -> Result<()> {
195    if let Some(max) = options.max_tokens {
196        let token_count = tokens.len();
197        if token_count > max {
198            let span = tokens
199                .get(max)
200                .or_else(|| tokens.last())
201                .map(GuardToken::span);
202            return Err(parse_guard_error(
203                "E_GUARD_TOKEN_BUDGET_EXCEEDED",
204                token_count,
205                max,
206                span,
207            ));
208        }
209    }
210
211    if options.max_parenthesis_depth.is_some() || options.max_function_call_depth.is_some() {
212        let mut paren_depth = 0usize;
213        let mut function_depth = 0usize;
214        let mut paren_stack = Vec::new();
215        let mut previous_significant: Option<TokenType> = None;
216
217        for token in tokens {
218            match token.token_type() {
219                TokenType::LParen => {
220                    paren_depth += 1;
221                    if let Some(max) = options.max_parenthesis_depth {
222                        if paren_depth > max {
223                            return Err(parse_guard_error(
224                                "E_GUARD_NESTING_DEPTH_EXCEEDED",
225                                paren_depth,
226                                max,
227                                Some(token.span()),
228                            ));
229                        }
230                    }
231
232                    let is_function_call = previous_significant
233                        .map(is_function_call_name_token)
234                        .unwrap_or(false);
235                    paren_stack.push(is_function_call);
236                    if is_function_call {
237                        function_depth += 1;
238                        if let Some(max) = options.max_function_call_depth {
239                            if function_depth > max {
240                                return Err(parse_guard_error(
241                                    "E_GUARD_FUNCTION_NESTING_DEPTH_EXCEEDED",
242                                    function_depth,
243                                    max,
244                                    Some(token.span()),
245                                ));
246                            }
247                        }
248                    }
249                }
250                TokenType::RParen => {
251                    paren_depth = paren_depth.saturating_sub(1);
252                    if paren_stack.pop().unwrap_or(false) {
253                        function_depth = function_depth.saturating_sub(1);
254                    }
255                }
256                _ => {}
257            }
258
259            if !is_trivia_token(token.token_type()) {
260                previous_significant = Some(token.token_type());
261            }
262        }
263    }
264
265    Ok(())
266}
267
268pub fn enforce_tokens(tokens: &[Token], options: &ComplexityGuardOptions) -> Result<()> {
269    enforce_token_slice(tokens, options)
270}
271
272pub(crate) fn enforce_parser_tokens(
273    tokens: &[ParserToken],
274    options: &ComplexityGuardOptions,
275) -> Result<()> {
276    enforce_token_slice(tokens, options)
277}
278
279pub(crate) fn enforce_parser_token_stats(
280    tokens: &[ParserToken],
281    stats: &TokenGuardStats,
282    options: &ComplexityGuardOptions,
283) -> Result<()> {
284    if let Some(max) = options.max_tokens {
285        if stats.token_count > max {
286            let span = tokens
287                .get(max)
288                .or_else(|| tokens.last())
289                .map(|token| token.span);
290            return Err(parse_guard_error(
291                "E_GUARD_TOKEN_BUDGET_EXCEEDED",
292                stats.token_count,
293                max,
294                span,
295            ));
296        }
297    }
298
299    let parenthesis_error = options.max_parenthesis_depth.and_then(|max| {
300        stats
301            .parenthesis_depth_spans
302            .get(max)
303            .copied()
304            .map(|span| (span, max))
305    });
306    let function_error = options.max_function_call_depth.and_then(|max| {
307        stats
308            .function_depth_spans
309            .get(max)
310            .copied()
311            .map(|span| (span, max))
312    });
313
314    match (parenthesis_error, function_error) {
315        (Some((paren_span, _paren_max)), Some((function_span, function_max)))
316            if function_span.start < paren_span.start =>
317        {
318            Err(parse_guard_error(
319                "E_GUARD_FUNCTION_NESTING_DEPTH_EXCEEDED",
320                function_max + 1,
321                function_max,
322                Some(function_span),
323            ))
324        }
325        (Some((span, max)), _) => Err(parse_guard_error(
326            "E_GUARD_NESTING_DEPTH_EXCEEDED",
327            max + 1,
328            max,
329            Some(span),
330        )),
331        (None, Some((span, max))) => Err(parse_guard_error(
332            "E_GUARD_FUNCTION_NESTING_DEPTH_EXCEEDED",
333            max + 1,
334            max,
335            Some(span),
336        )),
337        (None, None) => Ok(()),
338    }
339}
340
341fn is_trivia_token(token_type: TokenType) -> bool {
342    matches!(
343        token_type,
344        TokenType::Space | TokenType::Break | TokenType::LineComment | TokenType::BlockComment
345    )
346}
347
348fn is_function_call_name_token(token_type: TokenType) -> bool {
349    matches!(
350        token_type,
351        TokenType::Identifier
352            | TokenType::Var
353            | TokenType::QuotedIdentifier
354            | TokenType::CurrentDate
355            | TokenType::CurrentDateTime
356            | TokenType::CurrentTime
357            | TokenType::CurrentTimestamp
358            | TokenType::CurrentUser
359            | TokenType::If
360            | TokenType::Index
361            | TokenType::Insert
362            | TokenType::Left
363            | TokenType::Replace
364            | TokenType::Right
365            | TokenType::Row
366    )
367}
368
369#[cfg(test)]
370mod token_guard_tests {
371    use super::*;
372    use crate::tokens::Tokenizer;
373    use std::sync::Arc;
374
375    #[test]
376    fn collected_parser_stats_match_full_token_guard_pass() {
377        let sql: Arc<str> = Arc::from("SELECT outer(inner((value))), other(1, 2) FROM t");
378        let tokenizer = Tokenizer::default();
379        let public_tokens = tokenizer.tokenize(&sql).unwrap();
380        let (parser_tokens, stats) = tokenizer.tokenize_for_parser(&sql).unwrap();
381        let options = [
382            ComplexityGuardOptions::default(),
383            ComplexityGuardOptions {
384                max_tokens: Some(5),
385                ..Default::default()
386            },
387            ComplexityGuardOptions {
388                max_parenthesis_depth: Some(1),
389                ..Default::default()
390            },
391            ComplexityGuardOptions {
392                max_function_call_depth: Some(1),
393                ..Default::default()
394            },
395            ComplexityGuardOptions {
396                max_parenthesis_depth: Some(2),
397                max_function_call_depth: Some(1),
398                ..Default::default()
399            },
400        ];
401
402        for option in options {
403            let full_pass =
404                enforce_tokens(&public_tokens, &option).map_err(|error| error.to_string());
405            let collected = enforce_parser_token_stats(&parser_tokens, &stats, &option)
406                .map_err(|error| error.to_string());
407            assert_eq!(collected, full_pass);
408        }
409    }
410}
411
412/// Enforce AST size/depth limits and report parse-oriented errors.
413pub fn enforce_ast(expressions: &[Expression], options: &ComplexityGuardOptions) -> Result<()> {
414    let Some(max_nodes) = options.max_ast_nodes else {
415        return enforce_ast_depth_only(expressions, options);
416    };
417
418    let stats = ast_stats(expressions, Some(max_nodes), options.max_ast_depth)?;
419    if stats.node_count > max_nodes {
420        return Err(parse_guard_error(
421            "E_GUARD_AST_BUDGET_EXCEEDED",
422            stats.node_count,
423            max_nodes,
424            None,
425        ));
426    }
427
428    if let Some(max_depth) = options.max_ast_depth {
429        if stats.max_depth > max_depth {
430            return Err(parse_guard_error(
431                "E_GUARD_AST_DEPTH_EXCEEDED",
432                stats.max_depth,
433                max_depth,
434                None,
435            ));
436        }
437    }
438
439    Ok(())
440}
441
442/// Enforce AST size/depth limits and report generation-oriented errors.
443pub fn enforce_generate_ast(
444    expression: &Expression,
445    options: &ComplexityGuardOptions,
446) -> Result<()> {
447    let stats = ast_stats(
448        std::slice::from_ref(expression),
449        options.max_ast_nodes,
450        options.max_ast_depth,
451    )?;
452
453    if let Some(max_nodes) = options.max_ast_nodes {
454        if stats.node_count > max_nodes {
455            return Err(generate_guard_error(
456                "E_GUARD_AST_BUDGET_EXCEEDED",
457                stats.node_count,
458                max_nodes,
459            ));
460        }
461    }
462
463    if let Some(max_depth) = options.max_ast_depth {
464        if stats.max_depth > max_depth {
465            return Err(generate_guard_error(
466                "E_GUARD_AST_DEPTH_EXCEEDED",
467                stats.max_depth,
468                max_depth,
469            ));
470        }
471    }
472
473    Ok(())
474}
475
476fn enforce_ast_depth_only(
477    expressions: &[Expression],
478    options: &ComplexityGuardOptions,
479) -> Result<()> {
480    let Some(max_depth) = options.max_ast_depth else {
481        return Ok(());
482    };
483
484    let stats = ast_stats(expressions, None, Some(max_depth))?;
485    if stats.max_depth > max_depth {
486        return Err(parse_guard_error(
487            "E_GUARD_AST_DEPTH_EXCEEDED",
488            stats.max_depth,
489            max_depth,
490            None,
491        ));
492    }
493
494    Ok(())
495}
496
497#[derive(Debug, Clone, Copy, Default)]
498struct AstStats {
499    node_count: usize,
500    max_depth: usize,
501}
502
503fn ast_stats(
504    expressions: &[Expression],
505    max_nodes: Option<usize>,
506    max_depth: Option<usize>,
507) -> Result<AstStats> {
508    let mut stats = AstStats::default();
509    let mut stack: Vec<(&Expression, usize)> = expressions.iter().rev().map(|e| (e, 0)).collect();
510
511    while let Some((expr, depth)) = stack.pop() {
512        stats.node_count += 1;
513        stats.max_depth = stats.max_depth.max(depth);
514
515        if let Some(max) = max_nodes {
516            if stats.node_count > max {
517                return Ok(stats);
518            }
519        }
520
521        if let Some(max) = max_depth {
522            if stats.max_depth > max {
523                return Ok(stats);
524            }
525        }
526
527        push_ast_children(&mut stack, expr, depth);
528    }
529
530    Ok(stats)
531}
532
533fn push_ast_children<'a>(
534    stack: &mut Vec<(&'a Expression, usize)>,
535    expr: &'a Expression,
536    depth: usize,
537) {
538    match expr {
539        Expression::And(op) if is_commentless_binary_op(op) => {
540            push_connector_child(stack, &op.right, depth, ConnectorKind::And);
541            push_connector_child(stack, &op.left, depth, ConnectorKind::And);
542        }
543        Expression::Or(op) if is_commentless_binary_op(op) => {
544            push_connector_child(stack, &op.right, depth, ConnectorKind::Or);
545            push_connector_child(stack, &op.left, depth, ConnectorKind::Or);
546        }
547        _ => {
548            let child_start = stack.len();
549            crate::ast_children::for_each_child(expr, |_, child| {
550                stack.push((child, depth + 1));
551            });
552            stack[child_start..].reverse();
553        }
554    }
555}
556
557fn push_connector_child<'a>(
558    stack: &mut Vec<(&'a Expression, usize)>,
559    child: &'a Expression,
560    depth: usize,
561    kind: ConnectorKind,
562) {
563    let child_depth = if is_commentless_connector(child, kind) {
564        depth
565    } else {
566        depth + 1
567    };
568    stack.push((child, child_depth));
569}
570
571fn is_commentless_connector(expr: &Expression, kind: ConnectorKind) -> bool {
572    match (kind, expr) {
573        (ConnectorKind::And, Expression::And(op)) | (ConnectorKind::Or, Expression::Or(op)) => {
574            is_commentless_binary_op(op)
575        }
576        _ => false,
577    }
578}
579
580fn is_commentless_binary_op(op: &crate::expressions::BinaryOp) -> bool {
581    op.left_comments.is_empty()
582        && op.operator_comments.is_empty()
583        && op.trailing_comments.is_empty()
584}
585
586#[derive(Debug, Clone, Copy)]
587enum ConnectorKind {
588    And,
589    Or,
590}
polyglot_sql/guard.rs

polyglot_sql/
guard.rs
