[clang-tools-extra] 76910d4 - [pseudo] Share the underly payload when stripping comments for a token stream

[Haojian Wu via cfe-commits]

Author: Haojian Wu
Date: 2022-07-15T15:20:48+02:00
New Revision: 76910d4a56c8dba000f198bba13e71cf0492c8cb

URL: https://github.com/llvm/llvm-project/commit/76910d4a56c8dba000f198bba13e71cf0492c8cb
DIFF: https://github.com/llvm/llvm-project/commit/76910d4a56c8dba000f198bba13e71cf0492c8cb.diff

LOG: [pseudo] Share the underly payload when stripping comments for a token stream

`stripComments(cook(...))` is a common pattern being written.
Without this patch, this has a use-after-free issue (cook returns a temporary
TokenStream object which has its own payload, but the payload is not
shared with the one returned by stripComments).

Reviewed By: sammccall

Differential Revision: https://reviews.llvm.org/D125311

Added: 
    

Modified: 
    clang-tools-extra/pseudo/include/clang-pseudo/Token.h
    clang-tools-extra/pseudo/lib/Lex.cpp
    clang-tools-extra/pseudo/lib/Token.cpp

Removed: 
    


################################################################################
diff  --git a/clang-tools-extra/pseudo/include/clang-pseudo/Token.h b/clang-tools-extra/pseudo/include/clang-pseudo/Token.h
index b558891f0a86..36e5221a0d30 100644
--- a/clang-tools-extra/pseudo/include/clang-pseudo/Token.h
+++ b/clang-tools-extra/pseudo/include/clang-pseudo/Token.h
@@ -170,6 +170,18 @@ class TokenStream {
     return Storage[1];
   }
 
+  /// Returns the shared payload.
+  std::shared_ptr<void> getPayload() const { return Payload; }
+  /// Adds the given payload to the stream.
+  void addPayload(std::shared_ptr<void> P) {
+    if (!Payload)
+      Payload = std::move(P);
+    else
+      Payload = std::make_shared<
+          std::pair<std::shared_ptr<void>, std::shared_ptr<void>>>(
+          std::move(P), std::move(Payload));
+  }
+
   /// Print the tokens in this stream to the output stream.
   ///
   /// The presence of newlines/spaces is preserved, but not the quantity.

diff  --git a/clang-tools-extra/pseudo/lib/Lex.cpp b/clang-tools-extra/pseudo/lib/Lex.cpp
index 6a5a10f1b979..c96e2f27cba9 100644
--- a/clang-tools-extra/pseudo/lib/Lex.cpp
+++ b/clang-tools-extra/pseudo/lib/Lex.cpp
@@ -77,7 +77,7 @@ TokenStream cook(const TokenStream &Code, const LangOptions &LangOpts) {
   auto CleanedStorage = std::make_shared<llvm::BumpPtrAllocator>();
   clang::IdentifierTable Identifiers(LangOpts);
   TokenStream Result(CleanedStorage);
-
+  Result.addPayload(Code.getPayload());
   for (auto Tok : Code.tokens()) {
     if (Tok.flag(LexFlags::NeedsCleaning)) {
       // Remove escaped newlines and trigraphs.

diff  --git a/clang-tools-extra/pseudo/lib/Token.cpp b/clang-tools-extra/pseudo/lib/Token.cpp
index b58c8e4a862e..5b07a62f37fb 100644
--- a/clang-tools-extra/pseudo/lib/Token.cpp
+++ b/clang-tools-extra/pseudo/lib/Token.cpp
@@ -116,7 +116,7 @@ clang::LangOptions genericLangOpts(clang::Language Lang,
 }
 
 TokenStream stripComments(const TokenStream &Input) {
-  TokenStream Out;
+  TokenStream Out(Input.getPayload());
   for (const Token &T : Input.tokens()) {
     if (T.Kind == tok::comment)
       continue;