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

polyglot_sql/
guard.rs
