proc_macro: Validate tokens coming from the compiler again

1 files changed, 56 insertions, 51 deletions

diff --git a/src/libsyntax_ext/proc_macro_server.rs b/src/libsyntax_ext/proc_macro_server.rs
index a04d6c92b78..ea4a8afc0aa 100644
--- a/src/libsyntax_ext/proc_macro_server.rs
+++ b/src/libsyntax_ext/proc_macro_server.rs
@@ -81,29 +81,23 @@ impl FromInternal<(TokenStream, &'_ ParseSess, &'_ mut Vec<Self>)>
                     $($field $(: $value)*,)*
                     span,
                 })
-            )
+            );
+            ($ty:ident::$method:ident($($value:expr),*)) => (
+                TokenTree::$ty(self::$ty::$method($($value,)* span))
+            );
         }
         macro_rules! op {
             ($a:expr) => {
-                tt!(Punct { ch: $a, joint })
+                tt!(Punct::new($a, joint))
             };
             ($a:expr, $b:expr) => {{
-                stack.push(tt!(Punct { ch: $b, joint }));
-                tt!(Punct {
-                    ch: $a,
-                    joint: true
-                })
+                stack.push(tt!(Punct::new($b, joint)));
+                tt!(Punct::new($a, true))
             }};
             ($a:expr, $b:expr, $c:expr) => {{
-                stack.push(tt!(Punct { ch: $c, joint }));
-                stack.push(tt!(Punct {
-                    ch: $b,
-                    joint: true
-                }));
-                tt!(Punct {
-                    ch: $a,
-                    joint: true
-                })
+                stack.push(tt!(Punct::new($c, joint)));
+                stack.push(tt!(Punct::new($b, true)));
+                tt!(Punct::new($a, true))
             }};
         }
 
@@ -156,20 +150,11 @@ impl FromInternal<(TokenStream, &'_ ParseSess, &'_ mut Vec<Self>)>
             Question => op!('?'),
             SingleQuote => op!('\''),
 
-            Ident(ident, is_raw) => tt!(Ident {
-                sym: ident.name,
-                is_raw
-            }),
+            Ident(ident, is_raw) => tt!(Ident::new(ident.name, is_raw)),
             Lifetime(ident) => {
                 let ident = ident.without_first_quote();
-                stack.push(tt!(Ident {
-                    sym: ident.name,
-                    is_raw: false
-                }));
-                tt!(Punct {
-                    ch: '\'',
-                    joint: true
-                })
+                stack.push(tt!(Ident::new(ident.name, false)));
+                tt!(Punct::new('\'', true))
             }
             Literal(lit, suffix) => tt!(Literal { lit, suffix }),
             DocComment(c) => {
@@ -193,15 +178,9 @@ impl FromInternal<(TokenStream, &'_ ParseSess, &'_ mut Vec<Self>)>
                     span: DelimSpan::from_single(span),
                 }));
                 if style == ast::AttrStyle::Inner {
-                    stack.push(tt!(Punct {
-                        ch: '!',
-                        joint: false
-                    }));
+                    stack.push(tt!(Punct::new('!', false)));
                 }
-                tt!(Punct {
-                    ch: '#',
-                    joint: false
-                })
+                tt!(Punct::new('#', false))
             }
 
             Interpolated(_) => {
@@ -237,7 +216,7 @@ impl ToInternal<TokenStream> for TokenTree<Group, Punct, Ident, Literal> {
                 )
                 .into();
             }
-            TokenTree::Ident(self::Ident { sym, span, is_raw }) => {
+            TokenTree::Ident(self::Ident { sym, is_raw, span }) => {
                 let token = Ident(ast::Ident::new(sym, span), is_raw);
                 return tokenstream::TokenTree::Token(span, token).into();
             }
@@ -338,11 +317,48 @@ pub struct Punct {
     span: Span,
 }
 
+impl Punct {
+    fn new(ch: char, joint: bool, span: Span) -> Punct {
+        const LEGAL_CHARS: &[char] = &['=', '<', '>', '!', '~', '+', '-', '*', '/', '%', '^',
+                                       '&', '|', '@', '.', ',', ';', ':', '#', '$', '?', '\''];
+        if !LEGAL_CHARS.contains(&ch) {
+            panic!("unsupported character `{:?}`", ch)
+        }
+        Punct { ch, joint, span }
+    }
+}
+
 #[derive(Copy, Clone, PartialEq, Eq, Hash)]
 pub struct Ident {
     sym: Symbol,
-    span: Span,
     is_raw: bool,
+    span: Span,
+}
+
+impl Ident {
+    fn is_valid(string: &str) -> bool {
+        let mut chars = string.chars();
+        if let Some(start) = chars.next() {
+            (start == '_' || start.is_xid_start())
+                && chars.all(|cont| cont == '_' || cont.is_xid_continue())
+        } else {
+            false
+        }
+    }
+    fn new(sym: Symbol, is_raw: bool, span: Span) -> Ident {
+        let string = sym.as_str().get();
+        if !Self::is_valid(string) {
+            panic!("`{:?}` is not a valid identifier", string)
+        }
+        if is_raw {
+            let normalized_sym = Symbol::intern(string);
+            if normalized_sym == keywords::Underscore.name() ||
+               ast::Ident::with_empty_ctxt(normalized_sym).is_path_segment_keyword() {
+                panic!("`{:?}` is not a valid raw identifier", string)
+            }
+        }
+        Ident { sym, is_raw, span }
+    }
 }
 
 // FIXME(eddyb) `Literal` should not expose internal `Debug` impls.
@@ -492,11 +508,7 @@ impl server::Group for Rustc<'_> {
 
 impl server::Punct for Rustc<'_> {
     fn new(&mut self, ch: char, spacing: Spacing) -> Self::Punct {
-        Punct {
-            ch,
-            joint: spacing == Spacing::Joint,
-            span: server::Span::call_site(self),
-        }
+        Punct::new(ch, spacing == Spacing::Joint, server::Span::call_site(self))
     }
     fn as_char(&mut self, punct: Self::Punct) -> char {
         punct.ch
@@ -518,14 +530,7 @@ impl server::Punct for Rustc<'_> {
 
 impl server::Ident for Rustc<'_> {
     fn new(&mut self, string: &str, span: Self::Span, is_raw: bool) -> Self::Ident {
-        let sym = Symbol::intern(string);
-        if is_raw
-            && (sym == keywords::Underscore.name()
-                || ast::Ident::with_empty_ctxt(sym).is_path_segment_keyword())
-        {
-            panic!("`{:?}` is not a valid raw identifier", string)
-        }
-        Ident { sym, span, is_raw }
+        Ident::new(Symbol::intern(string), is_raw, span)
     }
     fn span(&mut self, ident: Self::Ident) -> Self::Span {
         ident.span