diff options
author | Edwin Cheng <[email protected]> | 2019-04-22 08:33:55 +0100 |
---|---|---|
committer | Edwin Cheng <[email protected]> | 2019-04-22 08:33:55 +0100 |
commit | b177813f3bef708636ec4be271e376b111c36a59 (patch) | |
tree | f9272f478f1b943b4e643ed9bd2b5fd1ee2e7e8b | |
parent | bbc5c1d24e1a641b134f634516828301e8cfc320 (diff) |
Add mbe expand limit and poision macro set
-rw-r--r-- | crates/ra_hir/src/ids.rs | 7 | ||||
-rw-r--r-- | crates/ra_hir/src/nameres.rs | 16 | ||||
-rw-r--r-- | crates/ra_hir/src/nameres/collector.rs | 186 | ||||
-rw-r--r-- | crates/ra_mbe/src/subtree_parser.rs | 13 | ||||
-rw-r--r-- | crates/ra_tt/src/lib.rs | 12 |
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(¯o_arg).map_err(|err| format!("{:?}", err))?; | 96 | let tt = macro_rules.expand(¯o_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 | ||
56 | use std::sync::Arc; | 56 | use std::sync::Arc; |
57 | 57 | ||
58 | use rustc_hash::FxHashMap; | 58 | use rustc_hash::{FxHashMap, FxHashSet}; |
59 | use ra_arena::{Arena, RawId, impl_arena_id}; | 59 | use ra_arena::{Arena, RawId, impl_arena_id}; |
60 | use ra_db::{FileId, Edition}; | 60 | use ra_db::{FileId, Edition}; |
61 | use test_utils::tested_by; | 61 | use 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 | ||
51 | trait 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)] | ||
59 | struct SimpleMacroStackMonitor { | ||
60 | counts: FxHashMap<MacroDefId, u32>, | ||
61 | } | ||
62 | |||
63 | impl 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(¯o_def_id).unwrap_or(&0) > 100 | ||
74 | } | ||
75 | } | ||
76 | |||
51 | /// Walks the tree of module recursively | 77 | /// Walks the tree of module recursively |
52 | struct DefCollector<DB> { | 78 | struct 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 | ||
65 | impl<'a, DB> DefCollector<&'a DB> | 91 | impl<'a, DB, M> DefCollector<&'a DB, M> |
66 | where | 92 | where |
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(¯o_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 | ||
359 | impl<DB> ModCollector<'_, &'_ mut DefCollector<&'_ DB>> | 396 | impl<DB, M> ModCollector<'_, &'_ mut DefCollector<&'_ DB, M>> |
360 | where | 397 | where |
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)] | ||
573 | mod 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 | } | ||
652 | foo!(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 | } | ||
666 | foo!(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 | } | ||
682 | foo!(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 | ||
6 | struct OffsetTokenSink { | 6 | struct OffsetTokenSink { |
7 | token_pos: usize, | 7 | token_pos: usize, |
8 | error: bool, | ||
8 | } | 9 | } |
9 | 10 | ||
10 | impl TreeSink for OffsetTokenSink { | 11 | impl 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 | ||
19 | pub(crate) struct Parser<'a> { | 22 | pub(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 | |||
153 | impl 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 | } | ||