From 3e5b15571632467f203c3f93d0b56f41d0fb64f5 Mon Sep 17 00:00:00 2001
From: Aleksey Kladov <aleksey.kladov@gmail.com>
Date: Mon, 9 Aug 2021 16:06:49 +0300
Subject: [PATCH] fix: avoid pathological macro expansions

Today, rust-analyzer (and rustc, and bat, and IntelliJ) fail badly on
some kinds of maliciously constructed code, like a deep sequence of
nested parenthesis.

"Who writes 100k nested parenthesis" you'd ask?

Well, in a language with macros, a run-away macro expansion might do
that (see the added tests)! Such expansion can be broad, rather than
deep, so it bypasses recursion check at the macro-expansion layer, but
triggers deep recursion in parser.

In the ideal world, the parser would just handle deeply nested structs
gracefully. We'll get there some day, but at the moment, let's try to be
simple, and just avoid expanding macros with unbalanced parenthesis in
the first place.

closes #9358
---
 Cargo.lock                                    |  1 +
 crates/hir_def/src/nameres/tests/macros.rs    | 18 +++++++++++++++
 crates/hir_expand/Cargo.toml                  |  1 +
 crates/hir_expand/src/db.rs                   | 19 +++++++++++++++-
 crates/ide_completion/src/tests/expression.rs | 22 +------------------
 crates/ide_ssr/src/tests.rs                   |  4 ++--
 6 files changed, 41 insertions(+), 24 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 4d42c1c7be8..61180fcb9a8 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -508,6 +508,7 @@ version = "0.0.0"
 dependencies = [
  "base_db",
  "cfg",
+ "cov-mark",
  "either",
  "expect-test",
  "la-arena",
diff --git a/crates/hir_def/src/nameres/tests/macros.rs b/crates/hir_def/src/nameres/tests/macros.rs
index 45448ba81b3..763dff5a214 100644
--- a/crates/hir_def/src/nameres/tests/macros.rs
+++ b/crates/hir_def/src/nameres/tests/macros.rs
@@ -1,4 +1,5 @@
 use super::*;
+
 use crate::nameres::proc_macro::{ProcMacroDef, ProcMacroKind};
 
 #[test]
@@ -1021,3 +1022,20 @@ pub mod prelude {
         "#]],
     )
 }
+
+#[test]
+fn issue9358_bad_macro_stack_overflow() {
+    cov_mark::check!(issue9358_bad_macro_stack_overflow);
+    check(
+        r#"
+macro_rules! m {
+  ($cond:expr) => { m!($cond, stringify!($cond)) };
+  ($cond:expr, $($arg:tt)*) => { $cond };
+}
+m!(
+"#,
+        expect![[r#"
+            crate
+        "#]],
+    )
+}
diff --git a/crates/hir_expand/Cargo.toml b/crates/hir_expand/Cargo.toml
index 92d6c3e96a1..4645970266c 100644
--- a/crates/hir_expand/Cargo.toml
+++ b/crates/hir_expand/Cargo.toml
@@ -9,6 +9,7 @@ edition = "2018"
 doctest = false
 
 [dependencies]
+cov-mark = "2.0.0-pre.1"
 log = "0.4.8"
 either = "1.5.3"
 rustc-hash = "1.0.0"
diff --git a/crates/hir_expand/src/db.rs b/crates/hir_expand/src/db.rs
index fa693e74ee1..a9bf6db476f 100644
--- a/crates/hir_expand/src/db.rs
+++ b/crates/hir_expand/src/db.rs
@@ -5,7 +5,7 @@ use std::sync::Arc;
 use base_db::{salsa, SourceDatabase};
 use limit::Limit;
 use mbe::{ExpandError, ExpandResult};
-use parser::FragmentKind;
+use parser::{FragmentKind, T};
 use syntax::{
     algo::diff,
     ast::{self, NameOwner},
@@ -273,6 +273,23 @@ fn macro_arg_text(db: &dyn AstDatabase, id: MacroCallId) -> Option<GreenNode> {
     let loc = db.lookup_intern_macro(id);
     let arg = loc.kind.arg(db)?;
     let arg = process_macro_input(db, arg, id);
+    if matches!(loc.kind, MacroCallKind::FnLike { .. }) {
+        let first = arg.first_child_or_token().map_or(T![.], |it| it.kind());
+        let last = arg.last_child_or_token().map_or(T![.], |it| it.kind());
+        let well_formed_tt =
+            matches!((first, last), (T!['('], T![')']) | (T!['['], T![']']) | (T!['{'], T!['}']));
+        if !well_formed_tt {
+            // Don't expand malformed (unbalanced) macro invocations. This is
+            // less than ideal, but trying to expand unbalanced  macro calls
+            // sometimes produces pathological, deeply nested code which breaks
+            // all kinds of things.
+            //
+            // Some day, we'll have explicit recursion counters for all
+            // recursive things, at which point this code might be removed.
+            cov_mark::hit!(issue9358_bad_macro_stack_overflow);
+            return None;
+        }
+    }
     Some(arg.green().into())
 }
 
diff --git a/crates/ide_completion/src/tests/expression.rs b/crates/ide_completion/src/tests/expression.rs
index 95aaff01ac3..9a42dd7b275 100644
--- a/crates/ide_completion/src/tests/expression.rs
+++ b/crates/ide_completion/src/tests/expression.rs
@@ -308,27 +308,7 @@ fn quux(x: i32) {
     m!(x$0
 }
 "#,
-        expect![[r#"
-            kw unsafe
-            kw match
-            kw while
-            kw while let
-            kw loop
-            kw if
-            kw if let
-            kw for
-            kw true
-            kw false
-            kw return
-            kw self
-            kw super
-            kw crate
-            lc y         i32
-            bt u32
-            lc x         i32
-            fn quux(…)   fn(i32)
-            ma m!(…)     macro_rules! m
-        "#]],
+        expect![[r#""#]],
     );
 }
 
diff --git a/crates/ide_ssr/src/tests.rs b/crates/ide_ssr/src/tests.rs
index 444c6b0afd0..0b0c1111c46 100644
--- a/crates/ide_ssr/src/tests.rs
+++ b/crates/ide_ssr/src/tests.rs
@@ -921,13 +921,13 @@ fn preserves_whitespace_within_macro_expansion() {
             macro_rules! macro1 {
                 ($a:expr) => {$a}
             }
-            fn f() {macro1!(1   *   2 + 3 + 4}
+            fn f() {macro1!(1   *   2 + 3 + 4)}
             "#,
         expect![[r#"
             macro_rules! macro1 {
                 ($a:expr) => {$a}
             }
-            fn f() {macro1!(4 - (3 - 1   *   2)}
+            fn f() {macro1!(4 - (3 - 1   *   2))}
             "#]],
     )
 }