[clang-tools-extra] 60e19f6 - [clangd] Fix use-after-free in HeaderIncluderCache

[Kadir Cetinkaya via cfe-commits]

Author: Kadir Cetinkaya
Date: 2021-10-20T16:36:07+02:00
New Revision: 60e19f6752b7f19b022d8bdca0937c59b49811f9

URL: https://github.com/llvm/llvm-project/commit/60e19f6752b7f19b022d8bdca0937c59b49811f9
DIFF: https://github.com/llvm/llvm-project/commit/60e19f6752b7f19b022d8bdca0937c59b49811f9.diff

LOG: [clangd] Fix use-after-free in HeaderIncluderCache

Includer cache could get into a bad state when a main file went bad and
added back afterwards. This patch adds a check to invalidate to prevent
that.

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

Added: 
    

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

Removed: 
    


################################################################################
diff  --git a/clang-tools-extra/clangd/TUScheduler.cpp b/clang-tools-extra/clangd/TUScheduler.cpp
index 4dd7a06b8bce..bbcda5bc0dfb 100644
--- a/clang-tools-extra/clangd/TUScheduler.cpp
+++ b/clang-tools-extra/clangd/TUScheduler.cpp
@@ -286,8 +286,12 @@ class TUScheduler::HeaderIncluderCache {
   void remove(PathRef MainFile) {
     std::lock_guard<std::mutex> Lock(Mu);
     Association *&First = MainToFirst[MainFile];
-    if (First)
+    if (First) {
       invalidate(First);
+      First = nullptr;
+    }
+    // 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.

diff  --git a/clang-tools-extra/clangd/unittests/TUSchedulerTests.cpp b/clang-tools-extra/clangd/unittests/TUSchedulerTests.cpp
index 306392fb2731..833df03a20ea 100644
--- a/clang-tools-extra/clangd/unittests/TUSchedulerTests.cpp
+++ b/clang-tools-extra/clangd/unittests/TUSchedulerTests.cpp
@@ -1200,10 +1200,10 @@ TEST_F(TUSchedulerTests, IncluderCache) {
                      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();
@@ -1218,6 +1218,8 @@ TEST_F(TUSchedulerTests, IncluderCache) {
       }
       return Basic;
     }
+
+    std::atomic<bool> FailAll{false};
   } CDB;
   TUScheduler S(CDB, optsForTest());
   auto GetFlags = [&](PathRef Header) {
@@ -1288,6 +1290,21 @@ TEST_F(TUSchedulerTests, IncluderCache) {
       << "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) {