fix(wgsl-in)!: limit brace recursion

ErichDonGubler · ErichDonGubler · commit b21a3265dee9 · 2024-04-03T15:54:43.000-04:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -155,6 +155,7 @@ Bottom level categories:
 - In spv-in, remove unnecessary "gl_PerVertex" name check so unused builtins will always be skipped. By @Imberflur in [#5227](https://github.com/gfx-rs/wgpu/pull/5227).
 - GLSL 410 does not support layout(binding = ...), enable only for GLSL 420. By @bes in [#5357](https://github.com/gfx-rs/wgpu/pull/5357)
 - In spv-out, check for acceleration and ray-query types when enabling ray-query extension to prevent validation error. By @Vecvec in [#5463](https://github.com/gfx-rs/wgpu/pull/5463)
+- Add a limit for curly brace nesting in WGSL parsing. By @ErichDonGubler in [#5447](https://github.com/gfx-rs/wgpu/pull/5447).
 
 #### Tests
 
diff --git a/naga/src/front/wgsl/error.rs b/naga/src/front/wgsl/error.rs
@@ -269,6 +269,10 @@ pub enum Error<'a> {
         scalar: String,
         inner: ConstantEvaluatorError,
     },
+    ExceededLimitForNestedBraces {
+        span: Span,
+        limit: u8,
+    },
 }
 
 impl<'a> Error<'a> {
@@ -770,6 +774,13 @@ impl<'a> Error<'a> {
                     format!("the expression should have been converted to have {} scalar type", scalar),
                 ]
             },
+            Error::ExceededLimitForNestedBraces { span, limit } => ParseError {
+                message: "brace nesting limit reached".into(),
+                labels: vec![(span, "limit reached at this brace".into())],
+                notes: vec![
+                    format!("nesting limit is currently set to {limit}"),
+                ],
+            },
         }
     }
 }
diff --git a/naga/src/front/wgsl/parse/mod.rs b/naga/src/front/wgsl/parse/mod.rs
@@ -1619,6 +1619,7 @@ impl Parser {
         lexer: &mut Lexer<'a>,
         ctx: &mut ExpressionContext<'a, '_, '_>,
         block: &mut ast::Block<'a>,
+        brace_nesting_level: u8,
     ) -> Result<(), Error<'a>> {
         self.push_rule_span(Rule::Statement, lexer);
         match lexer.peek() {
@@ -1627,7 +1628,7 @@ impl Parser {
                 self.pop_rule_span(lexer);
             }
             (Token::Paren('{'), _) => {
-                let (inner, span) = self.block(lexer, ctx)?;
+                let (inner, span) = self.block(lexer, ctx, brace_nesting_level)?;
                 block.stmts.push(ast::Statement {
                     kind: ast::StatementKind::Block(inner),
                     span,
@@ -1709,7 +1710,7 @@ impl Parser {
                         let _ = lexer.next();
                         let condition = self.general_expression(lexer, ctx)?;
 
-                        let accept = self.block(lexer, ctx)?.0;
+                        let accept = self.block(lexer, ctx, brace_nesting_level)?.0;
 
                         let mut elsif_stack = Vec::new();
                         let mut elseif_span_start = lexer.start_byte_offset();
@@ -1720,12 +1721,12 @@ impl Parser {
 
                             if !lexer.skip(Token::Word("if")) {
                                 // ... else { ... }
-                                break self.block(lexer, ctx)?.0;
+                                break self.block(lexer, ctx, brace_nesting_level)?.0;
                             }
 
                             // ... else if (...) { ... }
                             let other_condition = self.general_expression(lexer, ctx)?;
-                            let other_block = self.block(lexer, ctx)?;
+                            let other_block = self.block(lexer, ctx, brace_nesting_level)?;
                             elsif_stack.push((elseif_span_start, other_condition, other_block));
                             elseif_span_start = lexer.start_byte_offset();
                         };
@@ -1757,7 +1758,9 @@ impl Parser {
                     "switch" => {
                         let _ = lexer.next();
                         let selector = self.general_expression(lexer, ctx)?;
-                        lexer.expect(Token::Paren('{'))?;
+                        let brace_span = lexer.expect_span(Token::Paren('{'))?;
+                        let brace_nesting_level =
+                            Self::increase_brace_nesting(brace_nesting_level, brace_span)?;
                         let mut cases = Vec::new();
 
                         loop {
@@ -1782,7 +1785,7 @@ impl Parser {
                                         });
                                     };
 
-                                    let body = self.block(lexer, ctx)?.0;
+                                    let body = self.block(lexer, ctx, brace_nesting_level)?.0;
 
                                     cases.push(ast::SwitchCase {
                                         value,
@@ -1792,7 +1795,7 @@ impl Parser {
                                 }
                                 (Token::Word("default"), _) => {
                                     lexer.skip(Token::Separator(':'));
-                                    let body = self.block(lexer, ctx)?.0;
+                                    let body = self.block(lexer, ctx, brace_nesting_level)?.0;
                                     cases.push(ast::SwitchCase {
                                         value: ast::SwitchValue::Default,
                                         body,
@@ -1808,7 +1811,7 @@ impl Parser {
 
                         ast::StatementKind::Switch { selector, cases }
                     }
-                    "loop" => self.r#loop(lexer, ctx)?,
+                    "loop" => self.r#loop(lexer, ctx, brace_nesting_level)?,
                     "while" => {
                         let _ = lexer.next();
                         let mut body = ast::Block::default();
@@ -1832,7 +1835,7 @@ impl Parser {
                             span,
                         });
 
-                        let (block, span) = self.block(lexer, ctx)?;
+                        let (block, span) = self.block(lexer, ctx, brace_nesting_level)?;
                         body.stmts.push(ast::Statement {
                             kind: ast::StatementKind::Block(block),
                             span,
@@ -1855,7 +1858,9 @@ impl Parser {
                             let (_, span) = {
                                 let ctx = &mut *ctx;
                                 let block = &mut *block;
-                                lexer.capture_span(|lexer| self.statement(lexer, ctx, block))?
+                                lexer.capture_span(|lexer| {
+                                    self.statement(lexer, ctx, block, brace_nesting_level)
+                                })?
                             };
 
                             if block.stmts.len() != num_statements {
@@ -1900,7 +1905,7 @@ impl Parser {
                             lexer.expect(Token::Paren(')'))?;
                         }
 
-                        let (block, span) = self.block(lexer, ctx)?;
+                        let (block, span) = self.block(lexer, ctx, brace_nesting_level)?;
                         body.stmts.push(ast::Statement {
                             kind: ast::StatementKind::Block(block),
                             span,
@@ -1962,13 +1967,15 @@ impl Parser {
         &mut self,
         lexer: &mut Lexer<'a>,
         ctx: &mut ExpressionContext<'a, '_, '_>,
+        brace_nesting_level: u8,
     ) -> Result<ast::StatementKind<'a>, Error<'a>> {
         let _ = lexer.next();
         let mut body = ast::Block::default();
         let mut continuing = ast::Block::default();
         let mut break_if = None;
 
-        lexer.expect(Token::Paren('{'))?;
+        let brace_span = lexer.expect_span(Token::Paren('{'))?;
+        let brace_nesting_level = Self::increase_brace_nesting(brace_nesting_level, brace_span)?;
 
         ctx.local_table.push_scope();
 
@@ -1978,7 +1985,9 @@ impl Parser {
                 // the last thing in the loop body
 
                 // Expect a opening brace to start the continuing block
-                lexer.expect(Token::Paren('{'))?;
+                let brace_span = lexer.expect_span(Token::Paren('{'))?;
+                let brace_nesting_level =
+                    Self::increase_brace_nesting(brace_nesting_level, brace_span)?;
                 loop {
                     if lexer.skip(Token::Word("break")) {
                         // Branch for the `break if` statement, this statement
@@ -2007,7 +2016,7 @@ impl Parser {
                         break;
                     } else {
                         // Otherwise try to parse a statement
-                        self.statement(lexer, ctx, &mut continuing)?;
+                        self.statement(lexer, ctx, &mut continuing, brace_nesting_level)?;
                     }
                 }
                 // Since the continuing block must be the last part of the loop body,
@@ -2021,7 +2030,7 @@ impl Parser {
                 break;
             }
             // Otherwise try to parse a statement
-            self.statement(lexer, ctx, &mut body)?;
+            self.statement(lexer, ctx, &mut body, brace_nesting_level)?;
         }
 
         ctx.local_table.pop_scope();
@@ -2038,15 +2047,17 @@ impl Parser {
         &mut self,
         lexer: &mut Lexer<'a>,
         ctx: &mut ExpressionContext<'a, '_, '_>,
+        brace_nesting_level: u8,
     ) -> Result<(ast::Block<'a>, Span), Error<'a>> {
         self.push_rule_span(Rule::Block, lexer);
 
         ctx.local_table.push_scope();
 
-        lexer.expect(Token::Paren('{'))?;
+        let brace_span = lexer.expect_span(Token::Paren('{'))?;
+        let brace_nesting_level = Self::increase_brace_nesting(brace_nesting_level, brace_span)?;
         let mut block = ast::Block::default();
         while !lexer.skip(Token::Paren('}')) {
-            self.statement(lexer, ctx, &mut block)?;
+            self.statement(lexer, ctx, &mut block, brace_nesting_level)?;
         }
 
         ctx.local_table.pop_scope();
@@ -2133,9 +2144,10 @@ impl Parser {
 
         // do not use `self.block` here, since we must not push a new scope
         lexer.expect(Token::Paren('{'))?;
+        let brace_nesting_level = 1;
         let mut body = ast::Block::default();
         while !lexer.skip(Token::Paren('}')) {
-            self.statement(lexer, &mut ctx, &mut body)?;
+            self.statement(lexer, &mut ctx, &mut body, brace_nesting_level)?;
         }
 
         ctx.local_table.pop_scope();
@@ -2345,4 +2357,30 @@ impl Parser {
 
         Ok(tu)
     }
+
+    const fn increase_brace_nesting(
+        brace_nesting_level: u8,
+        brace_span: Span,
+    ) -> Result<u8, Error<'static>> {
+        // From [spec.](https://gpuweb.github.io/gpuweb/wgsl/#limits):
+        //
+        // > § 2.4. Limits
+        // >
+        // > …
+        // >
+        // > Maximum nesting depth of brace-enclosed statements in a function[:] 127
+        //
+        // _However_, we choose 64 instead because (a) it avoids stack overflows in CI and
+        // (b) we expect the limit to be decreased to 63 based on this conversation in
+        // WebGPU CTS upstream:
+        // <https://github.com/gpuweb/cts/pull/3389#discussion_r1543742701>
+        const BRACE_NESTING_MAXIMUM: u8 = 64;
+        if brace_nesting_level + 1 > BRACE_NESTING_MAXIMUM {
+            return Err(Error::ExceededLimitForNestedBraces {
+                span: brace_span,
+                limit: BRACE_NESTING_MAXIMUM,
+            });
+        }
+        Ok(brace_nesting_level + 1)
+    }
 }
diff --git a/naga/tests/wgsl_errors.rs b/naga/tests/wgsl_errors.rs
@@ -2151,3 +2151,129 @@ fn compaction_preserves_spans() {
         panic!("Error message has wrong span:\n\n{err:#?}");
     }
 }
+
+#[test]
+fn limit_braced_statement_nesting() {
+    let too_many_braces = "fn f() {{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{";
+
+    let expected_diagnostic = r###"error: brace nesting limit reached
+  ┌─ wgsl:1:72
+  │
+1 │ fn f() {{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{
+  │                                                                        ^ limit reached at this brace
+  │
+  = note: nesting limit is currently set to 64
+
+"###;
+
+    // In debug builds, we might actually overflow the stack before exercising this error case,
+    // depending on the platform and the `RUST_MIN_STACK` env. var. Use a thread with a custom
+    // stack size that works on all platforms.
+    std::thread::Builder::new()
+        .stack_size(1024 * 1024 * 2 /* MB */)
+        .spawn(|| check(too_many_braces, expected_diagnostic))
+        .unwrap()
+        .join()
+        .unwrap()
+}
+
+#[test]
+fn too_many_unclosed_loops() {
+    let too_many_braces = "fn f() {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+       loop {
+           ";
+
+    let expected_diagnostic = r###"error: brace nesting limit reached
+   ┌─ wgsl:65:13
+   │
+65 │        loop {
+   │             ^ limit reached at this brace
+   │
+   = note: nesting limit is currently set to 64
+
+"###;
+
+    // In debug builds, we might actually overflow the stack before exercising this error case,
+    // depending on the platform and the `RUST_MIN_STACK` env. var. Use a thread with a custom
+    // stack size that works on all platforms.
+    std::thread::Builder::new()
+        .stack_size(1024 * 1024 * 2 /* MB */)
+        .spawn(|| check(too_many_braces, expected_diagnostic))
+        .unwrap()
+        .join()
+        .unwrap()
+}

Original file line number	Diff line number	Diff line change
`@@ -269,6 +269,10 @@ pub enum Error<'a> {`
`269`	`269`	`scalar: String,`
`270`	`270`	`inner: ConstantEvaluatorError,`
`271`	`271`	`},`
	`272`	`+ ExceededLimitForNestedBraces {`
	`273`	`+ span: Span,`
	`274`	`+ limit: u8,`
	`275`	`+ },`
`272`	`276`	`}`
`273`	`277`
`274`	`278`	`impl<'a> Error<'a> {`
`@@ -770,6 +774,13 @@ impl<'a> Error<'a> {`
`770`	`774`	`format!("the expression should have been converted to have {} scalar type", scalar),`
`771`	`775`	`]`
`772`	`776`	`},`
	`777`	`+ Error::ExceededLimitForNestedBraces { span, limit } => ParseError {`
	`778`	`+ message: "brace nesting limit reached".into(),`
	`779`	`+ labels: vec![(span, "limit reached at this brace".into())],`
	`780`	`+ notes: vec![`
	`781`	`+ format!("nesting limit is currently set to {limit}"),`
	`782`	`+ ],`
	`783`	`+ },`
`773`	`784`	`}`
`774`	`785`	`}`
`775`	`786`	`}`