[clang] [Clang][Parser] Don't always destroy template annotations at the end of a declaration (PR #89494)

[Younan Zhang via cfe-commits]

https://github.com/zyn0217 created https://github.com/llvm/llvm-project/pull/89494

Since [6163aa9](https://github.com/llvm/llvm-project/commit/6163aa96799cbad7f2f58e02c5bebee9647056a5#diff-3a7ef0bff7d2b73b4100de636f09ea68b72eda191b39c8091a6a1765d917c1a2), we have introduced an optimization that almost always destroys TemplateIdAnnotations at the end of a function declaration. This doesn't always work properly: a lambda within a default template argument could also result in such deallocation and hence a use-after-free bug while building a type constraint on the template parameter.

Note the test doesn't always trigger a conspicuous bug/crash even with a debug build. But a sanitizer build can detect them, I believe.

Fixes https://github.com/llvm/llvm-project/issues/67235
Fixes https://github.com/llvm/llvm-project/issues/89127

>From e78b183748e52f4709db7f32d663a0a5c03c5af5 Mon Sep 17 00:00:00 2001
From: Younan Zhang <zyn7109 at gmail.com>
Date: Sat, 20 Apr 2024 02:52:16 +0800
Subject: [PATCH] [Clang][Parser] Don't always destroy template annotations at
 the end of a declaration

Since 6163aa9, we have introduced an optimization that almost always
destroys TemplateIdAnnotations at the end of a function declaration.
This doesn't always work properly: a lambda within a default template
argument could also result in such deallocation and hence a use-after-free
bug while building a type constraint on the template parameter.

Fixes https://github.com/llvm/llvm-project/issues/67235
Fixes https://github.com/llvm/llvm-project/issues/89127
---
 clang/docs/ReleaseNotes.rst                   |  1 +
 clang/include/clang/Parse/Parser.h            | 24 +++++++++++++++++++
 clang/lib/Parse/ParseTemplate.cpp             |  5 ++++
 .../cxx2a-constrained-template-param.cpp      | 20 +++++++++++++++-
 4 files changed, 49 insertions(+), 1 deletion(-)

diff --git a/clang/docs/ReleaseNotes.rst b/clang/docs/ReleaseNotes.rst
index efc32212f300cf..0cf81a59493faf 100644
--- a/clang/docs/ReleaseNotes.rst
+++ b/clang/docs/ReleaseNotes.rst
@@ -535,6 +535,7 @@ Bug Fixes to C++ Support
 - Fix a crash when deducing ``auto`` from an invalid dereference (#GH88329).
 - Fix a crash in requires expression with templated base class member function. Fixes (#GH84020).
 - Placement new initializes typedef array with correct size (#GH41441)
+- Fixed a use-after-free bug in parsing of type constraints with default arguments that involve lambdas. (#GH67235)
 
 Bug Fixes to AST Handling
 ^^^^^^^^^^^^^^^^^^^^^^^^^
diff --git a/clang/include/clang/Parse/Parser.h b/clang/include/clang/Parse/Parser.h
index 23b268126de4e0..2a08b3d56e0cd1 100644
--- a/clang/include/clang/Parse/Parser.h
+++ b/clang/include/clang/Parse/Parser.h
@@ -313,7 +313,15 @@ class Parser : public CodeCompletionHandler {
   /// top-level declaration is finished.
   SmallVector<TemplateIdAnnotation *, 16> TemplateIds;
 
+  /// Don't destroy template annotations in MaybeDestroyTemplateIds even if
+  /// we're at the end of a declaration. Instead, we defer the destruction until
+  /// after a top-level declaration.
+  /// Use DelayTemplateIdDestructionRAII rather than setting it directly.
+  bool DelayTemplateIdDestruction = false;
+
   void MaybeDestroyTemplateIds() {
+    if (DelayTemplateIdDestruction)
+      return;
     if (!TemplateIds.empty() &&
         (Tok.is(tok::eof) || !PP.mightHavePendingAnnotationTokens()))
       DestroyTemplateIds();
@@ -329,6 +337,22 @@ class Parser : public CodeCompletionHandler {
     ~DestroyTemplateIdAnnotationsRAIIObj() { Self.MaybeDestroyTemplateIds(); }
   };
 
+  struct DelayTemplateIdDestructionRAII {
+    Parser &Self;
+    bool PrevDelayTemplateIdDestruction;
+
+    DelayTemplateIdDestructionRAII(Parser &Self,
+                                   bool DelayTemplateIdDestruction) noexcept
+        : Self(Self),
+          PrevDelayTemplateIdDestruction(Self.DelayTemplateIdDestruction) {
+      Self.DelayTemplateIdDestruction = DelayTemplateIdDestruction;
+    }
+
+    ~DelayTemplateIdDestructionRAII() noexcept {
+      Self.DelayTemplateIdDestruction = PrevDelayTemplateIdDestruction;
+    }
+  };
+
   /// Identifiers which have been declared within a tentative parse.
   SmallVector<const IdentifierInfo *, 8> TentativelyDeclaredIdentifiers;
 
diff --git a/clang/lib/Parse/ParseTemplate.cpp b/clang/lib/Parse/ParseTemplate.cpp
index b07ce451e878eb..917690d42ad960 100644
--- a/clang/lib/Parse/ParseTemplate.cpp
+++ b/clang/lib/Parse/ParseTemplate.cpp
@@ -733,7 +733,12 @@ NamedDecl *Parser::ParseTypeParameter(unsigned Depth, unsigned Position) {
   // we introduce the type parameter into the local scope.
   SourceLocation EqualLoc;
   ParsedType DefaultArg;
+  std::optional<DelayTemplateIdDestructionRAII> DontDestructTemplateIds;
   if (TryConsumeToken(tok::equal, EqualLoc)) {
+    // The default argument might contain a lambda declaration; avoid destroying
+    // parsed template ids at the end of that declaration because they can be
+    // used in a type constraint later.
+    DontDestructTemplateIds.emplace(*this, false);
     // The default argument may declare template parameters, notably
     // if it contains a generic lambda, so we need to increase
     // the template depth as these parameters would not be instantiated
diff --git a/clang/test/Parser/cxx2a-constrained-template-param.cpp b/clang/test/Parser/cxx2a-constrained-template-param.cpp
index 6f14b66419c498..5afa99ea50cf91 100644
--- a/clang/test/Parser/cxx2a-constrained-template-param.cpp
+++ b/clang/test/Parser/cxx2a-constrained-template-param.cpp
@@ -49,4 +49,22 @@ namespace temp
   template<C1 TT, C1 UU = test1> // expected-error{{use of class template 'test1' requires template arguments}}
   // expected-error at -1 2{{concept named in type constraint is not a type concept}}
   using A = TT<int>; // expected-error{{expected ';' after alias declaration}}
-}
\ No newline at end of file
+}
+
+namespace PR67235 {
+
+template <class T>
+concept C = true;
+
+template <auto D>
+struct S {};
+
+// Don't destroy annotation 'C' at the end of the lambda; else we'll run into a
+// use-after-free while constructing the type constraint 'C' on 'Default'.
+template <typename Ret, C Default = decltype([] { return Ret(); })>
+void func() {}
+
+template <typename Ret, C Default = S<[] { return Ret(); }>>
+void func2() {}
+
+}