[PATCH] D112130: [clangd] Fix use-after-free in HeaderIncluderCache

[Kadir Cetinkaya via Phabricator via cfe-commits]

kadircet created this revision.
kadircet added a reviewer: sammccall.
Herald added subscribers: usaxena95, arphaman, javed.absar.
kadircet requested review of this revision.
Herald added subscribers: cfe-commits, MaskRay, ilya-biryukov.
Herald added a project: clang-tools-extra.

Includer cache could get into a bad state when a main file went bad and
added back afterwards. This fixes the issue by getting rid of the main file to
first pointer completely to make sure we don't try invalidating again.


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D112130

Files:
  clang-tools-extra/clangd/TUScheduler.cpp
  clang-tools-extra/clangd/unittests/TUSchedulerTests.cpp


Index: clang-tools-extra/clangd/unittests/TUSchedulerTests.cpp
===================================================================
--- clang-tools-extra/clangd/unittests/TUSchedulerTests.cpp
+++ clang-tools-extra/clangd/unittests/TUSchedulerTests.cpp
@@ -1174,10 +1174,10 @@
                      Unreliable = testPath("unreliable.h"),
                      OK = testPath("ok.h"),
                      NotIncluded = testPath("not_included.h");
-  class NoHeadersCDB : public GlobalCompilationDatabase {
+  struct NoHeadersCDB : public GlobalCompilationDatabase {
     llvm::Optional<tooling::CompileCommand>
     getCompileCommand(PathRef File) const override {
-      if (File == NoCmd || File == NotIncluded)
+      if (File == NoCmd || File == NotIncluded || FailAll)
         return llvm::None;
       auto Basic = getFallbackCommand(File);
       Basic.Heuristic.clear();
@@ -1192,6 +1192,8 @@
       }
       return Basic;
     }
+
+    std::atomic<bool> FailAll{false};
   } CDB;
   TUScheduler S(CDB, optsForTest());
   auto GetFlags = [&](PathRef Header) {
@@ -1262,6 +1264,21 @@
       << "association invalidated but not reclaimed";
   EXPECT_THAT(GetFlags(NotIncluded), Contains("-DMAIN2"))
       << "association still valid";
+
+  // Delete the file from CDB, it should invalidate the associations.
+  CDB.FailAll = true;
+  EXPECT_THAT(GetFlags(NoCmd), Not(Contains("-DMAIN3")))
+      << "association should've been invalidated.";
+  // Also run update for Main3 to invalidate the preeamble to make sure next
+  // update populates include cache associations.
+  S.update(Main3, getInputs(Main3, SomeIncludes), WantDiagnostics::Yes);
+  EXPECT_TRUE(S.blockUntilIdle(timeoutSeconds(10)));
+  // Re-add the file and make sure nothing crashes.
+  CDB.FailAll = false;
+  S.update(Main3, getInputs(Main3, SomeIncludes), WantDiagnostics::Yes);
+  EXPECT_TRUE(S.blockUntilIdle(timeoutSeconds(10)));
+  EXPECT_THAT(GetFlags(NoCmd), Contains("-DMAIN3"))
+      << "association invalidated and then claimed by main3";
 }
 
 TEST_F(TUSchedulerTests, PreservesLastActiveFile) {
Index: clang-tools-extra/clangd/TUScheduler.cpp
===================================================================
--- clang-tools-extra/clangd/TUScheduler.cpp
+++ clang-tools-extra/clangd/TUScheduler.cpp
@@ -285,9 +285,12 @@
   // will be eligible for association with other files that get update()d.
   void remove(PathRef MainFile) {
     std::lock_guard<std::mutex> Lock(Mu);
-    Association *&First = MainToFirst[MainFile];
-    if (First)
-      invalidate(First);
+    auto It = MainToFirst.find(MainFile);
+    if (It == MainToFirst.end())
+      return;
+    if (It->second)
+      invalidate(It->second);
+    MainToFirst.erase(It);
   }
 
   /// Get the mainfile associated with Header, or the empty string if none.


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D112130.380890.patch
Type: text/x-patch
Size: 2813 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/cfe-commits/attachments/20211020/c2b51278/attachment.bin>