[clangd] Fix use-after-free in HeaderIncluderCache

[kadircet]

Summary

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.

[kadircet]

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

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

[sammccall]

clang-tools-extra/clangd/TUScheduler.cpp

289

Hmm, this leaves MainToFirst pointing to an invalidated entry.
Would it make more sense to just null out the value (i.e. First here)?

Then I think we need no code changes elsewhere, as update() already checks if First already exists before invalidating.

(This special case would occur on a rarer path, but the real reason for my preference is that I think the data structure is easier to interpret)

[kadircet]

clang-tools-extra/clangd/TUScheduler.cpp

289

Thanks! Yeah that makes sense + now there's a point for First being a reference :P

clang-tools-extra/clangd/TUScheduler.cpp

void update(PathRef MainFile, llvm::ArrayRef<std::string> Headers) {
     updateMemoryUsage();
   }
 
   // Mark MainFile as gone.
   // This will *not* disassociate headers with MainFile immediately, but they
   // 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.
   std::string get(PathRef Header) const {
     std::lock_guard<std::mutex> Lock(Mu);
     return HeaderToMain.lookup(Header).MainFile.str();
   }
 

clang-tools-extra/clangd/unittests/TUSchedulerTests.cpp

 // it's included by another open file, then we parse it using that files flags.
 TEST_F(TUSchedulerTests, IncluderCache) {
   static std::string Main = testPath("main.cpp"), Main2 = testPath("main2.cpp"),
                      Main3 = testPath("main3.cpp"),
                      NoCmd = testPath("no_cmd.h"),
                      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();
       if (File == Unreliable) {
         Basic.Heuristic = "not reliable";
       } else if (File == Main) {
         Basic.CommandLine.push_back("-DMAIN");
       } else if (File == Main2) {
         Basic.CommandLine.push_back("-DMAIN2");
       } else if (File == Main3) {
         Basic.CommandLine.push_back("-DMAIN3");
       }
       return Basic;
     }
+
+    std::atomic<bool> FailAll{false};
   } CDB;
   TUScheduler S(CDB, optsForTest());
   auto GetFlags = [&](PathRef Header) {
     S.update(Header, getInputs(Header, ";"), WantDiagnostics::Yes);
     EXPECT_TRUE(S.blockUntilIdle(timeoutSeconds(10)));
     tooling::CompileCommand Cmd;
     S.runWithPreamble("GetFlags", Header, TUScheduler::StaleOrAbsent,
                       [&](llvm::Expected<InputsAndPreamble> Inputs) {
TEST_F(TUSchedulerTests, IncluderCache) {
   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";
   EXPECT_THAT(GetFlags(Unreliable), Contains("-DMAIN"))
       << "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) {
   for (bool Sync : {false, true}) {
     auto Opts = optsForTest();
     if (Sync)
       Opts.AsyncThreadsCount = 0;
     TUScheduler S(CDB, Opts);