aboutsummaryrefslogtreecommitdiff
diff options
context:
space:
mode:
authorEdwin Cheng <[email protected]>2019-04-22 08:33:55 +0100
committerEdwin Cheng <[email protected]>2019-04-22 08:33:55 +0100
commitb177813f3bef708636ec4be271e376b111c36a59 (patch)
treef9272f478f1b943b4e643ed9bd2b5fd1ee2e7e8b
parentbbc5c1d24e1a641b134f634516828301e8cfc320 (diff)
Add mbe expand limit and poision macro set
-rw-r--r--crates/ra_hir/src/ids.rs7
-rw-r--r--crates/ra_hir/src/nameres.rs16
-rw-r--r--crates/ra_hir/src/nameres/collector.rs186
-rw-r--r--crates/ra_mbe/src/subtree_parser.rs13
-rw-r--r--crates/ra_tt/src/lib.rs12
5 files changed, 216 insertions, 18 deletions
diff --git a/crates/ra_hir/src/ids.rs b/crates/ra_hir/src/ids.rs
index e771a311c..c7849c995 100644
--- a/crates/ra_hir/src/ids.rs
+++ b/crates/ra_hir/src/ids.rs
@@ -94,6 +94,13 @@ fn parse_macro(
94 94
95 let macro_rules = db.macro_def(loc.def).ok_or("Fail to find macro definition")?; 95 let macro_rules = db.macro_def(loc.def).ok_or("Fail to find macro definition")?;
96 let tt = macro_rules.expand(&macro_arg).map_err(|err| format!("{:?}", err))?; 96 let tt = macro_rules.expand(&macro_arg).map_err(|err| format!("{:?}", err))?;
97
98 // Set a hard limit for the expanded tt
99 let count = tt.count();
100 if count > 65536 {
101 return Err(format!("Total tokens count exceed limit : count = {}", count));
102 }
103
97 Ok(mbe::token_tree_to_ast_item_list(&tt)) 104 Ok(mbe::token_tree_to_ast_item_list(&tt))
98} 105}
99 106
diff --git a/crates/ra_hir/src/nameres.rs b/crates/ra_hir/src/nameres.rs
index 39152360c..fbfff4fd7 100644
--- a/crates/ra_hir/src/nameres.rs
+++ b/crates/ra_hir/src/nameres.rs
@@ -55,7 +55,7 @@ mod tests;
55 55
56use std::sync::Arc; 56use std::sync::Arc;
57 57
58use rustc_hash::FxHashMap; 58use rustc_hash::{FxHashMap, FxHashSet};
59use ra_arena::{Arena, RawId, impl_arena_id}; 59use ra_arena::{Arena, RawId, impl_arena_id};
60use ra_db::{FileId, Edition}; 60use ra_db::{FileId, Edition};
61use test_utils::tested_by; 61use test_utils::tested_by;
@@ -91,6 +91,19 @@ pub struct CrateDefMap {
91 root: CrateModuleId, 91 root: CrateModuleId,
92 modules: Arena<CrateModuleId, ModuleData>, 92 modules: Arena<CrateModuleId, ModuleData>,
93 public_macros: FxHashMap<Name, MacroDefId>, 93 public_macros: FxHashMap<Name, MacroDefId>,
94
95 /// Some macros are not well-behavior, which leads to infinite loop
96 /// e.g. macro_rules! foo { ($ty:ty) => { foo!($ty); } }
97 /// We mark it down and skip it in collector
98 ///
99 /// FIXME:
100 /// Right now it only handle a poison macro in a single crate,
101 /// such that if other crate try to call that macro,
102 /// the whole process will do again until it became poisoned in that crate.
103 /// We should handle this macro set globally
104 /// However, do we want to put it as a global variable?
105 poison_macros: FxHashSet<MacroDefId>,
106
94 diagnostics: Vec<DefDiagnostic>, 107 diagnostics: Vec<DefDiagnostic>,
95} 108}
96 109
@@ -195,6 +208,7 @@ impl CrateDefMap {
195 root, 208 root,
196 modules, 209 modules,
197 public_macros: FxHashMap::default(), 210 public_macros: FxHashMap::default(),
211 poison_macros: FxHashSet::default(),
198 diagnostics: Vec::new(), 212 diagnostics: Vec::new(),
199 } 213 }
200 }; 214 };
diff --git a/crates/ra_hir/src/nameres/collector.rs b/crates/ra_hir/src/nameres/collector.rs
index 6147b3219..82738cce3 100644
--- a/crates/ra_hir/src/nameres/collector.rs
+++ b/crates/ra_hir/src/nameres/collector.rs
@@ -42,14 +42,40 @@ pub(super) fn collect_defs(db: &impl DefDatabase, mut def_map: CrateDefMap) -> C
42 unresolved_imports: Vec::new(), 42 unresolved_imports: Vec::new(),
43 unexpanded_macros: Vec::new(), 43 unexpanded_macros: Vec::new(),
44 global_macro_scope: FxHashMap::default(), 44 global_macro_scope: FxHashMap::default(),
45 marco_stack_count: 0, 45 macro_stack_monitor: SimpleMacroStackMonitor::default(),
46 }; 46 };
47 collector.collect(); 47 collector.collect();
48 collector.finish() 48 collector.finish()
49} 49}
50 50
51trait MacroStackMonitor {
52 fn increase(&mut self, macro_def_id: MacroDefId);
53 fn decrease(&mut self, macro_def_id: MacroDefId);
54
55 fn is_poison(&self, macro_def_id: MacroDefId) -> bool;
56}
57
58#[derive(Default)]
59struct SimpleMacroStackMonitor {
60 counts: FxHashMap<MacroDefId, u32>,
61}
62
63impl MacroStackMonitor for SimpleMacroStackMonitor {
64 fn increase(&mut self, macro_def_id: MacroDefId) {
65 *self.counts.entry(macro_def_id).or_default() += 1;
66 }
67
68 fn decrease(&mut self, macro_def_id: MacroDefId) {
69 *self.counts.entry(macro_def_id).or_default() -= 1;
70 }
71
72 fn is_poison(&self, macro_def_id: MacroDefId) -> bool {
73 *self.counts.get(&macro_def_id).unwrap_or(&0) > 100
74 }
75}
76
51/// Walks the tree of module recursively 77/// Walks the tree of module recursively
52struct DefCollector<DB> { 78struct DefCollector<DB, M> {
53 db: DB, 79 db: DB,
54 def_map: CrateDefMap, 80 def_map: CrateDefMap,
55 glob_imports: FxHashMap<CrateModuleId, Vec<(CrateModuleId, raw::ImportId)>>, 81 glob_imports: FxHashMap<CrateModuleId, Vec<(CrateModuleId, raw::ImportId)>>,
@@ -59,12 +85,13 @@ struct DefCollector<DB> {
59 85
60 /// Some macro use `$tt:tt which mean we have to handle the macro perfectly 86 /// Some macro use `$tt:tt which mean we have to handle the macro perfectly
61 /// To prevent stackoverflow, we add a deep counter here for prevent that. 87 /// To prevent stackoverflow, we add a deep counter here for prevent that.
62 marco_stack_count: u32, 88 macro_stack_monitor: M,
63} 89}
64 90
65impl<'a, DB> DefCollector<&'a DB> 91impl<'a, DB, M> DefCollector<&'a DB, M>
66where 92where
67 DB: DefDatabase, 93 DB: DefDatabase,
94 M: MacroStackMonitor,
68{ 95{
69 fn collect(&mut self) { 96 fn collect(&mut self) {
70 let crate_graph = self.db.crate_graph(); 97 let crate_graph = self.db.crate_graph();
@@ -317,30 +344,40 @@ where
317 let def_map = self.db.crate_def_map(krate); 344 let def_map = self.db.crate_def_map(krate);
318 if let Some(macro_id) = def_map.public_macros.get(&path.segments[1].name).cloned() { 345 if let Some(macro_id) = def_map.public_macros.get(&path.segments[1].name).cloned() {
319 let call_id = MacroCallLoc { def: macro_id, ast_id: *ast_id }.id(self.db); 346 let call_id = MacroCallLoc { def: macro_id, ast_id: *ast_id }.id(self.db);
320 resolved.push((*module_id, call_id)); 347 resolved.push((*module_id, call_id, macro_id));
321 } 348 }
322 false 349 false
323 }); 350 });
324 351
325 for (module_id, macro_call_id) in resolved { 352 for (module_id, macro_call_id, macro_def_id) in resolved {
326 self.collect_macro_expansion(module_id, macro_call_id); 353 self.collect_macro_expansion(module_id, macro_call_id, macro_def_id);
327 } 354 }
328 res 355 res
329 } 356 }
330 357
331 fn collect_macro_expansion(&mut self, module_id: CrateModuleId, macro_call_id: MacroCallId) { 358 fn collect_macro_expansion(
332 self.marco_stack_count += 1; 359 &mut self,
360 module_id: CrateModuleId,
361 macro_call_id: MacroCallId,
362 macro_def_id: MacroDefId,
363 ) {
364 if self.def_map.poison_macros.contains(&macro_def_id) {
365 return;
366 }
367
368 self.macro_stack_monitor.increase(macro_def_id);
333 369
334 if self.marco_stack_count < 300 { 370 if !self.macro_stack_monitor.is_poison(macro_def_id) {
335 let file_id: HirFileId = macro_call_id.into(); 371 let file_id: HirFileId = macro_call_id.into();
336 let raw_items = self.db.raw_items(file_id); 372 let raw_items = self.db.raw_items(file_id);
337 ModCollector { def_collector: &mut *self, file_id, module_id, raw_items: &raw_items } 373 ModCollector { def_collector: &mut *self, file_id, module_id, raw_items: &raw_items }
338 .collect(raw_items.items()) 374 .collect(raw_items.items());
339 } else { 375 } else {
340 log::error!("Too deep macro expansion: {}", macro_call_id.debug_dump(self.db)); 376 log::error!("Too deep macro expansion: {}", macro_call_id.debug_dump(self.db));
377 self.def_map.poison_macros.insert(macro_def_id);
341 } 378 }
342 379
343 self.marco_stack_count -= 1; 380 self.macro_stack_monitor.decrease(macro_def_id);
344 } 381 }
345 382
346 fn finish(self) -> CrateDefMap { 383 fn finish(self) -> CrateDefMap {
@@ -356,9 +393,10 @@ struct ModCollector<'a, D> {
356 raw_items: &'a raw::RawItems, 393 raw_items: &'a raw::RawItems,
357} 394}
358 395
359impl<DB> ModCollector<'_, &'_ mut DefCollector<&'_ DB>> 396impl<DB, M> ModCollector<'_, &'_ mut DefCollector<&'_ DB, M>>
360where 397where
361 DB: DefDatabase, 398 DB: DefDatabase,
399 M: MacroStackMonitor,
362{ 400{
363 fn collect(&mut self, items: &[raw::RawItem]) { 401 fn collect(&mut self, items: &[raw::RawItem]) {
364 for item in items { 402 for item in items {
@@ -484,7 +522,7 @@ where
484 { 522 {
485 let macro_call_id = MacroCallLoc { def: macro_id, ast_id }.id(self.def_collector.db); 523 let macro_call_id = MacroCallLoc { def: macro_id, ast_id }.id(self.def_collector.db);
486 524
487 self.def_collector.collect_macro_expansion(self.module_id, macro_call_id); 525 self.def_collector.collect_macro_expansion(self.module_id, macro_call_id, macro_id);
488 return; 526 return;
489 } 527 }
490 528
@@ -530,3 +568,123 @@ fn resolve_submodule(
530 None => Err(if is_dir_owner { file_mod } else { file_dir_mod }), 568 None => Err(if is_dir_owner { file_mod } else { file_dir_mod }),
531 } 569 }
532} 570}
571
572#[cfg(test)]
573mod tests {
574 use ra_db::SourceDatabase;
575
576 use crate::{Crate, mock::MockDatabase, DefDatabase};
577 use ra_arena::{Arena};
578 use super::*;
579 use rustc_hash::FxHashSet;
580
581 struct LimitedMacroStackMonitor {
582 count: u32,
583 limit: u32,
584 poison_limit: u32,
585 }
586
587 impl MacroStackMonitor for LimitedMacroStackMonitor {
588 fn increase(&mut self, _: MacroDefId) {
589 self.count += 1;
590 assert!(self.count < self.limit);
591 }
592
593 fn decrease(&mut self, _: MacroDefId) {
594 self.count -= 1;
595 }
596
597 fn is_poison(&self, _: MacroDefId) -> bool {
598 self.count >= self.poison_limit
599 }
600 }
601
602 fn do_collect_defs(
603 db: &impl DefDatabase,
604 def_map: CrateDefMap,
605 monitor: impl MacroStackMonitor,
606 ) -> CrateDefMap {
607 let mut collector = DefCollector {
608 db,
609 def_map,
610 glob_imports: FxHashMap::default(),
611 unresolved_imports: Vec::new(),
612 unexpanded_macros: Vec::new(),
613 global_macro_scope: FxHashMap::default(),
614 macro_stack_monitor: monitor,
615 };
616 collector.collect();
617 collector.finish()
618 }
619
620 fn do_limited_resolve(code: &str, limit: u32, poison_limit: u32) -> CrateDefMap {
621 let (db, _source_root, _) = MockDatabase::with_single_file(&code);
622 let crate_id = db.crate_graph().iter().next().unwrap();
623 let krate = Crate { crate_id };
624
625 let def_map = {
626 let edition = krate.edition(&db);
627 let mut modules: Arena<CrateModuleId, ModuleData> = Arena::default();
628 let root = modules.alloc(ModuleData::default());
629 CrateDefMap {
630 krate,
631 edition,
632 extern_prelude: FxHashMap::default(),
633 prelude: None,
634 root,
635 modules,
636 public_macros: FxHashMap::default(),
637 poison_macros: FxHashSet::default(),
638 diagnostics: Vec::new(),
639 }
640 };
641
642 do_collect_defs(&db, def_map, LimitedMacroStackMonitor { count: 0, limit, poison_limit })
643 }
644
645 #[test]
646 fn test_macro_expand_limit_width() {
647 do_limited_resolve(
648 r#"
649 macro_rules! foo {
650 ($($ty:ty)*) => { foo!($($ty)*, $($ty)*); }
651 }
652foo!(KABOOM);
653 "#,
654 16,
655 1000,
656 );
657 }
658
659 #[test]
660 fn test_macro_expand_poisoned() {
661 let def = do_limited_resolve(
662 r#"
663 macro_rules! foo {
664 ($ty:ty) => { foo!($ty); }
665 }
666foo!(KABOOM);
667 "#,
668 100,
669 16,
670 );
671
672 assert_eq!(def.poison_macros.len(), 1);
673 }
674
675 #[test]
676 fn test_macro_expand_normal() {
677 let def = do_limited_resolve(
678 r#"
679 macro_rules! foo {
680 ($ident:ident) => { struct $ident {} }
681 }
682foo!(Bar);
683 "#,
684 16,
685 16,
686 );
687
688 assert_eq!(def.poison_macros.len(), 0);
689 }
690}
diff --git a/crates/ra_mbe/src/subtree_parser.rs b/crates/ra_mbe/src/subtree_parser.rs
index 528aa0f8a..f07107414 100644
--- a/crates/ra_mbe/src/subtree_parser.rs
+++ b/crates/ra_mbe/src/subtree_parser.rs
@@ -5,6 +5,7 @@ use ra_syntax::{SyntaxKind};
5 5
6struct OffsetTokenSink { 6struct OffsetTokenSink {
7 token_pos: usize, 7 token_pos: usize,
8 error: bool,
8} 9}
9 10
10impl TreeSink for OffsetTokenSink { 11impl TreeSink for OffsetTokenSink {
@@ -13,7 +14,9 @@ impl TreeSink for OffsetTokenSink {
13 } 14 }
14 fn start_node(&mut self, _kind: SyntaxKind) {} 15 fn start_node(&mut self, _kind: SyntaxKind) {}
15 fn finish_node(&mut self) {} 16 fn finish_node(&mut self) {}
16 fn error(&mut self, _error: ra_parser::ParseError) {} 17 fn error(&mut self, _error: ra_parser::ParseError) {
18 self.error = true;
19 }
17} 20}
18 21
19pub(crate) struct Parser<'a> { 22pub(crate) struct Parser<'a> {
@@ -67,11 +70,15 @@ impl<'a> Parser<'a> {
67 F: FnOnce(&dyn TokenSource, &mut dyn TreeSink), 70 F: FnOnce(&dyn TokenSource, &mut dyn TreeSink),
68 { 71 {
69 let mut src = SubtreeTokenSource::new(&self.subtree.token_trees[*self.cur_pos..]); 72 let mut src = SubtreeTokenSource::new(&self.subtree.token_trees[*self.cur_pos..]);
70 let mut sink = OffsetTokenSink { token_pos: 0 }; 73 let mut sink = OffsetTokenSink { token_pos: 0, error: false };
71 74
72 f(&src, &mut sink); 75 f(&src, &mut sink);
73 76
74 self.finish(sink.token_pos, &mut src) 77 let r = self.finish(sink.token_pos, &mut src);
78 if sink.error {
79 return None;
80 }
81 r
75 } 82 }
76 83
77 fn finish(self, parsed_token: usize, src: &mut SubtreeTokenSource) -> Option<tt::TokenTree> { 84 fn finish(self, parsed_token: usize, src: &mut SubtreeTokenSource) -> Option<tt::TokenTree> {
diff --git a/crates/ra_tt/src/lib.rs b/crates/ra_tt/src/lib.rs
index 0b0b9b4d2..9cc646140 100644
--- a/crates/ra_tt/src/lib.rs
+++ b/crates/ra_tt/src/lib.rs
@@ -149,3 +149,15 @@ impl fmt::Display for Punct {
149 fmt::Display::fmt(&self.char, f) 149 fmt::Display::fmt(&self.char, f)
150 } 150 }
151} 151}
152
153impl Subtree {
154 /// Count the number of tokens recursively
155 pub fn count(&self) -> usize {
156 self.token_trees.iter().fold(self.token_trees.len(), |acc, c| {
157 acc + match c {
158 TokenTree::Subtree(c) => c.count(),
159 _ => 0,
160 }
161 })
162 }
163}