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// WASM has no stacker support, and grammar frames differ substantially in size.
15// Keep its default conservative; an explicit option still replaces this value.
16#[cfg(target_arch = "wasm32")]
17const DEFAULT_MAX_PARSER_DEPTH: usize = 32;
18#[cfg(not(target_arch = "wasm32"))]
19const DEFAULT_MAX_PARSER_DEPTH: usize = 1_024;
20
21fn default_max_parser_depth() -> Option<usize> {
22    Some(DEFAULT_MAX_PARSER_DEPTH)
23}
24
25/// Resource exhaustion must not be treated as a speculative syntax mismatch.
26pub(crate) fn is_guard_error(error: &Error) -> bool {
27    matches!(error, Error::Parse { message, .. } if message.starts_with("E_GUARD_"))
28}
29
30fn default_max_input_bytes() -> Option<usize> {
31    Some(DEFAULT_MAX_INPUT_BYTES)
32}
33
34fn default_max_tokens() -> Option<usize> {
35    Some(DEFAULT_MAX_TOKENS)
36}
37
38fn default_max_ast_nodes() -> Option<usize> {
39    Some(DEFAULT_MAX_AST_NODES)
40}
41
42fn default_max_ast_depth() -> Option<usize> {
43    Some(DEFAULT_MAX_AST_DEPTH)
44}
45
46fn default_max_parenthesis_depth() -> Option<usize> {
47    Some(DEFAULT_MAX_PARENTHESES_DEPTH)
48}
49
50fn default_max_function_call_depth() -> Option<usize> {
51    Some(DEFAULT_MAX_FUNCTION_CALL_DEPTH)
52}
53
54/// Guard options for parse/transpile/generate complexity.
55///
56/// These limits turn excessively deep or large inputs into regular errors
57/// instead of relying on process stack exhaustion as the failure mode.
58#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
59#[serde(rename_all = "camelCase")]
60pub struct ComplexityGuardOptions {
61    /// Maximum logical nesting during parsing, before an AST exists.
62    /// Defaults to 32 on WASM and 1024 on native targets.
63    /// Raising this limit or disabling it with `None` can permit stack exhaustion.
64    #[serde(default = "default_max_parser_depth")]
65    pub max_parser_depth: Option<usize>,
66    /// Maximum allowed SQL input size in bytes.
67    /// `None` disables this check.
68    #[serde(default = "default_max_input_bytes")]
69    pub max_input_bytes: Option<usize>,
70    /// Maximum allowed number of tokens after tokenization.
71    /// `None` disables this check.
72    #[serde(default = "default_max_tokens")]
73    pub max_tokens: Option<usize>,
74    /// Maximum allowed AST node count after parsing.
75    /// `None` disables this check.
76    #[serde(default = "default_max_ast_nodes")]
77    pub max_ast_nodes: Option<usize>,
78    /// Maximum allowed AST depth after parsing.
79    /// `None` disables this check.
80    #[serde(default = "default_max_ast_depth")]
81    pub max_ast_depth: Option<usize>,
82    /// Maximum allowed nested parenthesis depth before parsing.
83    /// `None` disables this check.
84    #[serde(default = "default_max_parenthesis_depth")]
85    pub max_parenthesis_depth: Option<usize>,
86    /// Maximum allowed nested function-call depth before parsing.
87    /// `None` disables this check.
88    #[serde(default = "default_max_function_call_depth")]
89    pub max_function_call_depth: Option<usize>,
90}
91
92#[derive(Debug, Default)]
93pub(crate) struct TokenGuardStats {
94    pub token_count: usize,
95    parenthesis_depth_spans: Vec<Span>,
96    function_depth_spans: Vec<Span>,
97    parenthesis_stack: Vec<bool>,
98    parenthesis_depth: usize,
99    function_depth: usize,
100    previous_significant: Option<TokenType>,
101}
102
103impl TokenGuardStats {
104    pub(crate) fn observe(&mut self, token_type: TokenType, span: Span) {
105        self.token_count += 1;
106        match token_type {
107            TokenType::LParen => {
108                self.parenthesis_depth += 1;
109                if self.parenthesis_depth_spans.len() < self.parenthesis_depth {
110                    self.parenthesis_depth_spans.push(span);
111                }
112
113                let is_function_call = self
114                    .previous_significant
115                    .map(is_function_call_name_token)
116                    .unwrap_or(false);
117                self.parenthesis_stack.push(is_function_call);
118                if is_function_call {
119                    self.function_depth += 1;
120                    if self.function_depth_spans.len() < self.function_depth {
121                        self.function_depth_spans.push(span);
122                    }
123                }
124            }
125            TokenType::RParen => {
126                self.parenthesis_depth = self.parenthesis_depth.saturating_sub(1);
127                if self.parenthesis_stack.pop().unwrap_or(false) {
128                    self.function_depth = self.function_depth.saturating_sub(1);
129                }
130            }
131            _ => {}
132        }
133
134        if !is_trivia_token(token_type) {
135            self.previous_significant = Some(token_type);
136        }
137    }
138}
139
140impl Default for ComplexityGuardOptions {
141    fn default() -> Self {
142        Self {
143            max_parser_depth: default_max_parser_depth(),
144            max_input_bytes: default_max_input_bytes(),
145            max_tokens: default_max_tokens(),
146            max_ast_nodes: default_max_ast_nodes(),
147            max_ast_depth: default_max_ast_depth(),
148            max_parenthesis_depth: default_max_parenthesis_depth(),
149            max_function_call_depth: default_max_function_call_depth(),
150        }
151    }
152}
153
154fn parse_guard_error(code: &str, actual: usize, limit: usize, span: Option<Span>) -> Error {
155    let message = format!("{code}: value {actual} exceeds configured limit {limit}");
156    if let Some(span) = span {
157        Error::parse(message, span.line, span.column, span.start, span.end)
158    } else {
159        Error::parse(message, 0, 0, 0, 0)
160    }
161}
162
163fn generate_guard_error(code: &str, actual: usize, limit: usize) -> Error {
164    Error::generate(format!(
165        "{code}: value {actual} exceeds configured limit {limit}"
166    ))
167}
168
169/// Enforce raw SQL input limits before tokenization.
170pub fn enforce_input(sql: &str, options: &ComplexityGuardOptions) -> Result<()> {
171    if let Some(max) = options.max_input_bytes {
172        let input_bytes = sql.len();
173        if input_bytes > max {
174            return Err(parse_guard_error(
175                "E_GUARD_INPUT_TOO_LARGE",
176                input_bytes,
177                max,
178                None,
179            ));
180        }
181    }
182
183    Ok(())
184}
185
186/// Enforce token and pre-parse nesting limits.
187trait GuardToken {
188    fn token_type(&self) -> TokenType;
189    fn span(&self) -> Span;
190}
191
192impl GuardToken for Token {
193    fn token_type(&self) -> TokenType {
194        self.token_type
195    }
196
197    fn span(&self) -> Span {
198        self.span
199    }
200}
201
202impl GuardToken for ParserToken {
203    fn token_type(&self) -> TokenType {
204        self.token_type
205    }
206
207    fn span(&self) -> Span {
208        self.span
209    }
210}
211
212fn enforce_token_slice<T: GuardToken>(
213    tokens: &[T],
214    options: &ComplexityGuardOptions,
215) -> Result<()> {
216    if let Some(max) = options.max_tokens {
217        let token_count = tokens.len();
218        if token_count > max {
219            let span = tokens
220                .get(max)
221                .or_else(|| tokens.last())
222                .map(GuardToken::span);
223            return Err(parse_guard_error(
224                "E_GUARD_TOKEN_BUDGET_EXCEEDED",
225                token_count,
226                max,
227                span,
228            ));
229        }
230    }
231
232    if options.max_parenthesis_depth.is_some() || options.max_function_call_depth.is_some() {
233        let mut paren_depth = 0usize;
234        let mut function_depth = 0usize;
235        let mut paren_stack = Vec::new();
236        let mut previous_significant: Option<TokenType> = None;
237
238        for token in tokens {
239            match token.token_type() {
240                TokenType::LParen => {
241                    paren_depth += 1;
242                    if let Some(max) = options.max_parenthesis_depth {
243                        if paren_depth > max {
244                            return Err(parse_guard_error(
245                                "E_GUARD_NESTING_DEPTH_EXCEEDED",
246                                paren_depth,
247                                max,
248                                Some(token.span()),
249                            ));
250                        }
251                    }
252
253                    let is_function_call = previous_significant
254                        .map(is_function_call_name_token)
255                        .unwrap_or(false);
256                    paren_stack.push(is_function_call);
257                    if is_function_call {
258                        function_depth += 1;
259                        if let Some(max) = options.max_function_call_depth {
260                            if function_depth > max {
261                                return Err(parse_guard_error(
262                                    "E_GUARD_FUNCTION_NESTING_DEPTH_EXCEEDED",
263                                    function_depth,
264                                    max,
265                                    Some(token.span()),
266                                ));
267                            }
268                        }
269                    }
270                }
271                TokenType::RParen => {
272                    paren_depth = paren_depth.saturating_sub(1);
273                    if paren_stack.pop().unwrap_or(false) {
274                        function_depth = function_depth.saturating_sub(1);
275                    }
276                }
277                _ => {}
278            }
279
280            if !is_trivia_token(token.token_type()) {
281                previous_significant = Some(token.token_type());
282            }
283        }
284    }
285
286    Ok(())
287}
288
289pub fn enforce_tokens(tokens: &[Token], options: &ComplexityGuardOptions) -> Result<()> {
290    enforce_token_slice(tokens, options)
291}
292
293pub(crate) fn enforce_parser_tokens(
294    tokens: &[ParserToken],
295    options: &ComplexityGuardOptions,
296) -> Result<()> {
297    enforce_token_slice(tokens, options)
298}
299
300pub(crate) fn enforce_parser_token_stats(
301    tokens: &[ParserToken],
302    stats: &TokenGuardStats,
303    options: &ComplexityGuardOptions,
304) -> Result<()> {
305    if let Some(max) = options.max_tokens {
306        if stats.token_count > max {
307            let span = tokens
308                .get(max)
309                .or_else(|| tokens.last())
310                .map(|token| token.span);
311            return Err(parse_guard_error(
312                "E_GUARD_TOKEN_BUDGET_EXCEEDED",
313                stats.token_count,
314                max,
315                span,
316            ));
317        }
318    }
319
320    let parenthesis_error = options.max_parenthesis_depth.and_then(|max| {
321        stats
322            .parenthesis_depth_spans
323            .get(max)
324            .copied()
325            .map(|span| (span, max))
326    });
327    let function_error = options.max_function_call_depth.and_then(|max| {
328        stats
329            .function_depth_spans
330            .get(max)
331            .copied()
332            .map(|span| (span, max))
333    });
334
335    match (parenthesis_error, function_error) {
336        (Some((paren_span, _paren_max)), Some((function_span, function_max)))
337            if function_span.start < paren_span.start =>
338        {
339            Err(parse_guard_error(
340                "E_GUARD_FUNCTION_NESTING_DEPTH_EXCEEDED",
341                function_max + 1,
342                function_max,
343                Some(function_span),
344            ))
345        }
346        (Some((span, max)), _) => Err(parse_guard_error(
347            "E_GUARD_NESTING_DEPTH_EXCEEDED",
348            max + 1,
349            max,
350            Some(span),
351        )),
352        (None, Some((span, max))) => Err(parse_guard_error(
353            "E_GUARD_FUNCTION_NESTING_DEPTH_EXCEEDED",
354            max + 1,
355            max,
356            Some(span),
357        )),
358        (None, None) => Ok(()),
359    }
360}
361
362fn is_trivia_token(token_type: TokenType) -> bool {
363    matches!(
364        token_type,
365        TokenType::Space | TokenType::Break | TokenType::LineComment | TokenType::BlockComment
366    )
367}
368
369fn is_function_call_name_token(token_type: TokenType) -> bool {
370    matches!(
371        token_type,
372        TokenType::Identifier
373            | TokenType::Var
374            | TokenType::QuotedIdentifier
375            | TokenType::CurrentDate
376            | TokenType::CurrentDateTime
377            | TokenType::CurrentTime
378            | TokenType::CurrentTimestamp
379            | TokenType::CurrentUser
380            | TokenType::If
381            | TokenType::Index
382            | TokenType::Insert
383            | TokenType::Left
384            | TokenType::Replace
385            | TokenType::Right
386            | TokenType::Row
387    )
388}
389
390#[cfg(test)]
391mod token_guard_tests {
392    use super::*;
393    use crate::tokens::Tokenizer;
394    use std::sync::Arc;
395
396    #[test]
397    fn parser_depth_json_contract() {
398        let omitted: ComplexityGuardOptions = serde_json::from_str("{}").unwrap();
399        assert_eq!(omitted.max_parser_depth, Some(DEFAULT_MAX_PARSER_DEPTH));
400        assert_eq!(
401            omitted.max_parser_depth,
402            ComplexityGuardOptions::default().max_parser_depth
403        );
404        for limit in [Some(0), Some(1), Some(2048), None] {
405            let guard = ComplexityGuardOptions {
406                max_parser_depth: limit,
407                ..Default::default()
408            };
409            let json = serde_json::to_string(&guard).unwrap();
410            assert_eq!(
411                serde_json::from_str::<ComplexityGuardOptions>(&json).unwrap(),
412                guard
413            );
414        }
415        for value in ["-1", "1.5", "true", "\"10\""] {
416            assert!(serde_json::from_str::<ComplexityGuardOptions>(&format!(
417                "{{\"maxParserDepth\":{value}}}"
418            ))
419            .is_err());
420        }
421    }
422
423    #[test]
424    fn collected_parser_stats_match_full_token_guard_pass() {
425        let sql: Arc<str> = Arc::from("SELECT outer(inner((value))), other(1, 2) FROM t");
426        let tokenizer = Tokenizer::default();
427        let public_tokens = tokenizer.tokenize(&sql).unwrap();
428        let (parser_tokens, stats) = tokenizer.tokenize_for_parser(&sql).unwrap();
429        let options = [
430            ComplexityGuardOptions::default(),
431            ComplexityGuardOptions {
432                max_tokens: Some(5),
433                ..Default::default()
434            },
435            ComplexityGuardOptions {
436                max_parenthesis_depth: Some(1),
437                ..Default::default()
438            },
439            ComplexityGuardOptions {
440                max_function_call_depth: Some(1),
441                ..Default::default()
442            },
443            ComplexityGuardOptions {
444                max_parenthesis_depth: Some(2),
445                max_function_call_depth: Some(1),
446                ..Default::default()
447            },
448        ];
449
450        for option in options {
451            let full_pass =
452                enforce_tokens(&public_tokens, &option).map_err(|error| error.to_string());
453            let collected = enforce_parser_token_stats(&parser_tokens, &stats, &option)
454                .map_err(|error| error.to_string());
455            assert_eq!(collected, full_pass);
456        }
457    }
458}
459
460/// Enforce AST size/depth limits and report parse-oriented errors.
461pub fn enforce_ast(expressions: &[Expression], options: &ComplexityGuardOptions) -> Result<()> {
462    let Some(max_nodes) = options.max_ast_nodes else {
463        return enforce_ast_depth_only(expressions, options);
464    };
465
466    let stats = ast_stats(expressions, Some(max_nodes), options.max_ast_depth)?;
467    if stats.node_count > max_nodes {
468        return Err(parse_guard_error(
469            "E_GUARD_AST_BUDGET_EXCEEDED",
470            stats.node_count,
471            max_nodes,
472            None,
473        ));
474    }
475
476    if let Some(max_depth) = options.max_ast_depth {
477        if stats.max_depth > max_depth {
478            return Err(parse_guard_error(
479                "E_GUARD_AST_DEPTH_EXCEEDED",
480                stats.max_depth,
481                max_depth,
482                None,
483            ));
484        }
485    }
486
487    Ok(())
488}
489
490/// Enforce AST size/depth limits and report generation-oriented errors.
491pub fn enforce_generate_ast(
492    expression: &Expression,
493    options: &ComplexityGuardOptions,
494) -> Result<()> {
495    let stats = ast_stats(
496        std::slice::from_ref(expression),
497        options.max_ast_nodes,
498        options.max_ast_depth,
499    )?;
500
501    if let Some(max_nodes) = options.max_ast_nodes {
502        if stats.node_count > max_nodes {
503            return Err(generate_guard_error(
504                "E_GUARD_AST_BUDGET_EXCEEDED",
505                stats.node_count,
506                max_nodes,
507            ));
508        }
509    }
510
511    if let Some(max_depth) = options.max_ast_depth {
512        if stats.max_depth > max_depth {
513            return Err(generate_guard_error(
514                "E_GUARD_AST_DEPTH_EXCEEDED",
515                stats.max_depth,
516                max_depth,
517            ));
518        }
519    }
520
521    Ok(())
522}
523
524fn enforce_ast_depth_only(
525    expressions: &[Expression],
526    options: &ComplexityGuardOptions,
527) -> Result<()> {
528    let Some(max_depth) = options.max_ast_depth else {
529        return Ok(());
530    };
531
532    let stats = ast_stats(expressions, None, Some(max_depth))?;
533    if stats.max_depth > max_depth {
534        return Err(parse_guard_error(
535            "E_GUARD_AST_DEPTH_EXCEEDED",
536            stats.max_depth,
537            max_depth,
538            None,
539        ));
540    }
541
542    Ok(())
543}
544
545#[derive(Debug, Clone, Copy, Default)]
546struct AstStats {
547    node_count: usize,
548    max_depth: usize,
549}
550
551fn ast_stats(
552    expressions: &[Expression],
553    max_nodes: Option<usize>,
554    max_depth: Option<usize>,
555) -> Result<AstStats> {
556    let mut stats = AstStats::default();
557    let mut stack: Vec<(&Expression, usize)> = expressions.iter().rev().map(|e| (e, 0)).collect();
558
559    while let Some((expr, depth)) = stack.pop() {
560        stats.node_count += 1;
561        stats.max_depth = stats.max_depth.max(depth);
562
563        if let Some(max) = max_nodes {
564            if stats.node_count > max {
565                return Ok(stats);
566            }
567        }
568
569        if let Some(max) = max_depth {
570            if stats.max_depth > max {
571                return Ok(stats);
572            }
573        }
574
575        push_ast_children(&mut stack, expr, depth);
576    }
577
578    Ok(stats)
579}
580
581fn push_ast_children<'a>(
582    stack: &mut Vec<(&'a Expression, usize)>,
583    expr: &'a Expression,
584    depth: usize,
585) {
586    match expr {
587        Expression::And(op) if is_commentless_binary_op(op) => {
588            push_connector_child(stack, &op.right, depth, ConnectorKind::And);
589            push_connector_child(stack, &op.left, depth, ConnectorKind::And);
590        }
591        Expression::Or(op) if is_commentless_binary_op(op) => {
592            push_connector_child(stack, &op.right, depth, ConnectorKind::Or);
593            push_connector_child(stack, &op.left, depth, ConnectorKind::Or);
594        }
595        _ => {
596            let child_start = stack.len();
597            crate::ast_children::for_each_child(expr, |_, child| {
598                stack.push((child, depth + 1));
599            });
600            stack[child_start..].reverse();
601        }
602    }
603}
604
605fn push_connector_child<'a>(
606    stack: &mut Vec<(&'a Expression, usize)>,
607    child: &'a Expression,
608    depth: usize,
609    kind: ConnectorKind,
610) {
611    let child_depth = if is_commentless_connector(child, kind) {
612        depth
613    } else {
614        depth + 1
615    };
616    stack.push((child, child_depth));
617}
618
619fn is_commentless_connector(expr: &Expression, kind: ConnectorKind) -> bool {
620    match (kind, expr) {
621        (ConnectorKind::And, Expression::And(op)) | (ConnectorKind::Or, Expression::Or(op)) => {
622            is_commentless_binary_op(op)
623        }
624        _ => false,
625    }
626}
627
628fn is_commentless_binary_op(op: &crate::expressions::BinaryOp) -> bool {
629    op.left_comments.is_empty()
630        && op.operator_comments.is_empty()
631        && op.trailing_comments.is_empty()
632}
633
634#[derive(Debug, Clone, Copy)]
635enum ConnectorKind {
636    And,
637    Or,
638}
polyglot_sql/guard.rs

polyglot_sql/
guard.rs
