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

[Kadir Cetinkaya via Phabricator via cfe-commits]

kadircet updated this revision to Diff 380895.
kadircet added a comment.

- Rather than dropping the entry, perform an extra check during invalidate. As

the entry actually backs the data for main file strings in associations.


Repository:
  rG LLVM Github Monorepo

CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D112130/new/

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
@@ -224,6 +224,9 @@
   mutable std::mutex Mu;
 
   void invalidate(Association *First) {
+    // Association might've already been invalidated.
+    if (!First->Next)
+      return;
     Association *Current = First;
     do {
       Association *Next = Current->Next;
@@ -288,6 +291,8 @@
     Association *&First = MainToFirst[MainFile];
     if (First)
       invalidate(First);
+    // MainToFirst entry should stay alive, as Associations might be pointing at
+    // its key.
   }
 
   /// Get the mainfile associated with Header, or the empty string if none.


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