preprocessor: Fix use after free when #undef the macro being invoked

src/compiler/preprocessor/DiagnosticsBase.cpp

@@ -74,6 +74,8 @@ std::string Diagnostics::message(ID id)
         return "predefined macro undefined";
       case PP_MACRO_UNTERMINATED_INVOCATION:
         return "unterminated macro invocation";
+      case PP_MACRO_UNDEFINED_WHILE_INVOKED:
+          return "macro undefined while being invoked";
       case PP_MACRO_TOO_FEW_ARGS:
         return "Not enough arguments for macro";
       case PP_MACRO_TOO_MANY_ARGS:

src/compiler/preprocessor/DiagnosticsBase.h

@@ -44,6 +44,7 @@ class Diagnostics
         PP_MACRO_PREDEFINED_REDEFINED,
         PP_MACRO_PREDEFINED_UNDEFINED,
         PP_MACRO_UNTERMINATED_INVOCATION,
+        PP_MACRO_UNDEFINED_WHILE_INVOKED,
         PP_MACRO_TOO_FEW_ARGS,
         PP_MACRO_TOO_MANY_ARGS,
         PP_MACRO_DUPLICATE_PARAMETER_NAMES,

src/compiler/preprocessor/DirectiveParser.cpp

@@ -445,6 +445,13 @@ void DirectiveParser::parseUndef(Token *token)
         {
             mDiagnostics->report(Diagnostics::PP_MACRO_PREDEFINED_UNDEFINED,
                                  token->location, token->text);
+            return;
+        }
+        else if (iter->second.expansionCount > 0)
+        {
+            mDiagnostics->report(Diagnostics::PP_MACRO_UNDEFINED_WHILE_INVOKED, token->location,
+                                 token->text);
+            return;
         }
         else
         {

src/compiler/preprocessor/Macro.h

@@ -26,16 +26,12 @@ struct Macro
     typedef std::vector<std::string> Parameters;
     typedef std::vector<Token> Replacements;
 
-    Macro()
-        : predefined(false),
-          disabled(false),
-          type(kTypeObj)
-    {
-    }
+    Macro() : predefined(false), disabled(false), expansionCount(0), type(kTypeObj) {}
     bool equals(const Macro &other) const;
 
     bool predefined;
     mutable bool disabled;
+    mutable int expansionCount;
 
     Type type;
     std::string name;

src/compiler/preprocessor/MacroExpander.cpp

@@ -151,6 +151,8 @@ bool MacroExpander::pushMacro(const Macro &macro, const Token &identifier)
     assert(identifier.type == Token::IDENTIFIER);
     assert(identifier.text == macro.name);
 
+    macro.expansionCount++;
+
     std::vector<Token> replacements;
     if (!expandMacro(macro, identifier, &replacements))
         return false;
@@ -174,7 +176,9 @@ void MacroExpander::popMacro()
 
     assert(context->empty());
     assert(context->macro->disabled);
+    assert(context->macro->expansionCount > 0);
     context->macro->disabled = false;
+    context->macro->expansionCount--;
     delete context;
 }
 

src/tests/preprocessor_tests/define_test.cpp

@@ -938,4 +938,20 @@ TEST_F(DefineTest, NegativeShiftInLineDirective)
     EXPECT_CALL(mDiagnostics, print(pp::Diagnostics::PP_UNDEFINED_SHIFT, _, _)).Times(4);
     EXPECT_CALL(mDiagnostics, print(pp::Diagnostics::PP_INVALID_LINE_NUMBER, _, _)).Times(2);
     preprocess(input, expected);
-}
+}
\ No newline at end of file
+
+// Undefining a macro in its invocation parameters produces and error
+TEST_F(DefineTest, UndefineInInvocation)
+{
+    const char *input =
+        "#define G(a, b) a b\n"
+        "G(\n"
+        "#undef G\n"
+        "1, 2)\n";
+    const char *expected = "\n\n\n1 2\n";
+
+    EXPECT_CALL(mDiagnostics, print(pp::Diagnostics::PP_MACRO_UNDEFINED_WHILE_INVOKED,
+                                    pp::SourceLocation(0, 3), _));
+
+    preprocess(input, expected);
+}