[clang] Fix use-after-free on code completion

[ilya-biryukov]

Summary

Found by asan. Fiddling with code completion AST after
FrontendAction::Exceute can lead to errors.
Calling the callback in ProcessCodeCompleteResults to make sure we
don't access uninitialized state.

This particular issue comes from the fact that Sema::TUScope is
deleted when destructor of ~Parser runs, but still present in
Sema::TUScope and accessed when building completion items.

I'm still struggling to come up with a small repro. The relevant
stackframes reported by asan are:
ERROR: AddressSanitizer: heap-use-after-free on address
READ of size 8 at 0x61400020d090 thread T175

#0 0x5632dff7821b in llvm::SmallPtrSetImplBase::isSmall() const include/llvm/ADT/SmallPtrSet.h:195:33
#1 0x5632e0335901 in llvm::SmallPtrSetImplBase::insert_imp(void const*) include/llvm/ADT/SmallPtrSet.h:127:9
#2 0x5632e067347d in llvm::SmallPtrSetImpl<clang::Decl*>::insert(clang::Decl*) include/llvm/ADT/SmallPtrSet.h:372:14
#3 0x5632e065df80 in clang::Scope::AddDecl(clang::Decl*) tools/clang/include/clang/Sema/Scope.h:287:18
#4 0x5632e0623eea in clang::ASTReader::pushExternalDeclIntoScope(clang::NamedDecl*, clang::DeclarationName) clang/lib/Serialization/ASTReader.cpp
#5 0x5632e062ce74 in clang::ASTReader::finishPendingActions() tools/clang/lib/Serialization/ASTReader.cpp:9164:9
....
#30 0x5632e02009c4 in clang::index::generateUSRForDecl(clang::Decl const*, llvm::SmallVectorImpl<char>&) tools/clang/lib/Index/USRGeneration.cpp:1037:6
#31 0x5632dff73eab in clang::clangd::(anonymous namespace)::getSymbolID(clang::CodeCompletionResult const&) tools/clang/tools/extra/clangd/CodeComplete.cpp:326:20
#32 0x5632dff6fe91 in clang::clangd::CodeCompleteFlow::mergeResults(std::vector<clang::CodeCompletionResult, std::allocator<clang::CodeCompletionResult> > const&, clang::clangd::SymbolSlab const&)::'lambda'(clang::CodeCompletionResult const&)::operator()(clang::CodeCompletionResult const&) tools/clang/tools/extra/clangd/CodeComplete.cpp:938:24
#33 0x5632dff6e426 in clang::clangd::CodeCompleteFlow::mergeResults(std::vector<clang::CodeCompletionResult, std::allocator<clang::CodeCompletionResult> > const&, clang::clangd::SymbolSlab const&) third_party/llvm/llvm/tools/clang/tools/extra/clangd/CodeComplete.cpp:949:38
#34 0x5632dff7a34d in clang::clangd::CodeCompleteFlow::runWithSema() llvm/tools/clang/tools/extra/clangd/CodeComplete.cpp:894:16
#35 0x5632dff6df6a in clang::clangd::CodeCompleteFlow::run(clang::clangd::(anonymous namespace)::SemaCompleteInput const&) &&::'lambda'()::operator()() const third_party/llvm/llvm/tools/clang/tools/extra/clangd/CodeComplete.cpp:858:35
#36 0x5632dff6cd42 in clang::clangd::(anonymous namespace)::semaCodeComplete(std::unique_ptr<clang::CodeCompleteConsumer, std::default_delete<clang::CodeCompleteConsumer> >, clang::CodeCompleteOptions const&, clang::clangd::(anonymous namespace)::SemaCompleteInput const&, llvm::function_ref<void ()>) tools/clang/tools/extra/clangd/CodeComplete.cpp:735:5

0x61400020d090 is located 80 bytes inside of 432-byte region [0x61400020d040,0x61400020d1f0)
freed by thread T175 here:

#0 0x5632df74e115 in operator delete(void*, unsigned long) projects/compiler-rt/lib/asan/asan_new_delete.cc:161:3
#1 0x5632e0b06973 in clang::Parser::~Parser() tools/clang/lib/Parse/Parser.cpp:410:3
#2 0x5632e0b06ddd in clang::Parser::~Parser() clang/lib/Parse/Parser.cpp:408:19
#3 0x5632e0b03286 in std::unique_ptr<clang::Parser, std::default_delete<clang::Parser> >::~unique_ptr() .../bits/unique_ptr.h:236:4
#4 0x5632e0b021c4 in clang::ParseAST(clang::Sema&, bool, bool) tools/clang/lib/Parse/ParseAST.cpp:182:1
#5 0x5632e0726544 in clang::FrontendAction::Execute() tools/clang/lib/Frontend/FrontendAction.cpp:904:8
#6 0x5632dff6cd05 in clang::clangd::(anonymous namespace)::semaCodeComplete(std::unique_ptr<clang::CodeCompleteConsumer, std::default_delete<clang::CodeCompleteConsumer> >, clang::CodeCompleteOptions const&, clang::clangd::(anonymous namespace)::SemaCompleteInput const&, llvm::function_ref<void ()>) tools/clang/tools/extra/clangd/CodeComplete.cpp:728:15

[ilya-biryukov]

Remove code refactoring from the patch.

[ilya-biryukov]

Remove a trace for "sema cleanup", it is not very informative now that we run the callback under the "sema completion" trace

[sammccall]

This feels like a bug in the underlying clang libraries, but since none of the lifetimes are documented and everyone just does it this way...

clangd/CodeComplete.cpp

431 ↗

(On Diff #136707)

maybe comment: 'generally the fields/methods should only be used from within the callback'

470 ↗

(On Diff #136707)

I think this check is dead - remove?

660 ↗

(On Diff #136707)

remove this comment

665 ↗

(On Diff #136707)

this can be stack-allocated, as we don't need to reset it

737 ↗

(On Diff #136707)

this tracer isn't needed anymore as discussed

816 ↗

(On Diff #136707)

this comment should probably be ", and do all our work in its callback"

852 ↗

(On Diff #136707)

remove "these are saved for merging"

[ilya-biryukov]

Addressed review comments

clang-tools-extra/trunk/clangd/CodeComplete.cpp

std::vector<std::string> getQueryScopes(CodeCompletionContext &CCContext,
 
   return Info.scopesForIndexQuery();
 }
 
 // The CompletionRecorder captures Sema code-complete output, including context.
 // It filters out ignored results (but doesn't apply fuzzy-filtering yet).
 // It doesn't do scoring or conversion to CompletionItem yet, as we want to
 // merge with index results first.
+// Generally the fields and methods of this object should only be used from
+// within the callback.
 struct CompletionRecorder : public CodeCompleteConsumer {
-  CompletionRecorder(const CodeCompleteOptions &Opts)
+  CompletionRecorder(const CodeCompleteOptions &Opts,
+                     UniqueFunction<void()> ResultsCallback)
       : CodeCompleteConsumer(Opts.getClangCompleteOpts(),
                              /*OutputIsBinary=*/false),
         CCContext(CodeCompletionContext::CCC_Other), Opts(Opts),
         CCAllocator(std::make_shared<GlobalCodeCompletionAllocator>()),
-        CCTUInfo(CCAllocator) {}
+        CCTUInfo(CCAllocator), ResultsCallback(std::move(ResultsCallback)) {
+    assert(this->ResultsCallback);
+  }
+
   std::vector<CodeCompletionResult> Results;
   CodeCompletionContext CCContext;
   Sema *CCSema = nullptr; // Sema that created the results.
   // FIXME: Sema is scary. Can we store ASTContext and Preprocessor, instead?
 
   void ProcessCodeCompleteResults(class Sema &S, CodeCompletionContext Context,
                                   CodeCompletionResult *InResults,
                                   unsigned NumResults) override final {
for (unsigned I = 0; I < NumResults; ++I) {
            Result.Availability == CXAvailability_NotAccessible))
         continue;
       // Destructor completion is rarely useful, and works inconsistently.
       // (s.^ completes ~string, but s.~st^ is an error).
       if (dyn_cast_or_null<CXXDestructorDecl>(Result.Declaration))
         continue;
       Results.push_back(Result);
     }
+    ResultsCallback();
   }
 
   CodeCompletionAllocator &getAllocator() override { return *CCAllocator; }
   CodeCompletionTUInfo &getCodeCompletionTUInfo() override { return CCTUInfo; }
 
   // Returns the filtering/sorting name for Result, which must be from Results.
   // Returned string is owned by this recorder (or the AST).
   llvm::StringRef getName(const CodeCompletionResult &Result) {
CodeCompletionString *codeCompletionString(const CodeCompletionResult &R,
     return const_cast<CodeCompletionResult &>(R).CreateCodeCompletionString(
         *CCSema, CCContext, *CCAllocator, CCTUInfo, IncludeBriefComments);
   }
 
 private:
   CodeCompleteOptions Opts;
   std::shared_ptr<GlobalCodeCompletionAllocator> CCAllocator;
   CodeCompletionTUInfo CCTUInfo;
+  UniqueFunction<void()> ResultsCallback;
 };
 
 // Tracks a bounded number of candidates with the best scores.
 class TopN {
 public:
   using value_type = std::pair<CompletionCandidate, CompletionItemScores>;
   static constexpr size_t Unbounded = std::numeric_limits<size_t>::max();
 
struct SemaCompleteInput {
   PrecompiledPreamble const *Preamble;
   StringRef Contents;
   Position Pos;
   IntrusiveRefCntPtr<vfs::FileSystem> VFS;
   std::shared_ptr<PCHContainerOperations> PCHs;
 };
 
 // Invokes Sema code completion on a file.
-// Callback will be invoked once completion is done, but before cleaning up.
 bool semaCodeComplete(std::unique_ptr<CodeCompleteConsumer> Consumer,
                       const clang::CodeCompleteOptions &Options,
-                      const SemaCompleteInput &Input,
-                      llvm::function_ref<void()> Callback = nullptr) {
-  auto Tracer = llvm::make_unique<trace::Span>("Sema completion");
+                      const SemaCompleteInput &Input) {
+  trace::Span Tracer("Sema completion");
   std::vector<const char *> ArgStrs;
   for (const auto &S : Input.Command.CommandLine)
     ArgStrs.push_back(S.c_str());
 
   if (Input.VFS->setCurrentWorkingDirectory(Input.Command.Directory)) {
     log("Couldn't set working directory");
     // We run parsing anyway, our lit-tests rely on results for non-existing
     // working dirs.
if (!Action.BeginSourceFile(*Clang, Clang->getFrontendOpts().Inputs[0])) {
     log("BeginSourceFile() failed when running codeComplete for " +
         Input.FileName);
     return false;
   }
   if (!Action.Execute()) {
     log("Execute() failed when running codeComplete for " + Input.FileName);
     return false;
   }
-  Tracer.reset();
-
-  if (Callback)
-    Callback();
-
-  Tracer = llvm::make_unique<trace::Span>("Sema completion cleanup");
   Action.EndSourceFile();
 
   return true;
 }
 
 // Should we perform index-based completion in this context?
 // FIXME: consider allowing completion, but restricting the result types.
 bool allowIndex(enum CodeCompletionContext::Kind K) {

 //     namespaces to search in). So Sema must start first.
 //   - we only want to return the top results (Opts.Limit).
 //     Building CompletionItems for everything else is wasteful, so we want to
 //     preserve the "native" format until we're done with scoring.
 //   - the data underlying Sema completion items is owned by the AST and various
 //     other arenas, which must stay alive for us to build CompletionItems.
 //   - we may get duplicate results from Sema and the Index, we need to merge.
 //
-// So we start Sema completion first, but defer its cleanup until we're done.
+// So we start Sema completion first, and do all our work in its callback.
 // We use the Sema context information to query the index.
 // Then we merge the two result sets, producing items that are Sema/Index/Both.
 // These items are scored, and the top N are synthesized into the LSP response.
 // Finally, we can clean up the data structures created by Sema completion.
 //
 // Main collaborators are:
 //   - semaCodeComplete sets up the compiler machinery to run code completion.
 //   - CompletionRecorder captures Sema completion results, including context.
 //   - SymbolIndex (Opts.Index) provides index completion results as Symbols
 //   - CompletionCandidates are the result of merging Sema and Index results.
 //     Each candidate points to an underlying CodeCompletionResult (Sema), a
 //     Symbol (Index), or both. It computes the result quality score.
 //     CompletionCandidate also does conversion to CompletionItem (at the end).
 //   - FuzzyMatcher scores how the candidate matches the partial identifier.
 //     This score is combined with the result quality score for the final score.
 //   - TopN determines the results with the best score.
 class CodeCompleteFlow {
   PathRef FileName;
   const CodeCompleteOptions &Opts;
   // Sema takes ownership of Recorder. Recorder is valid until Sema cleanup.
-  std::unique_ptr<CompletionRecorder> RecorderOwner;
-  CompletionRecorder &Recorder;
+  CompletionRecorder *Recorder = nullptr;
   int NSema = 0, NIndex = 0, NBoth = 0; // Counters for logging.
   bool Incomplete = false; // Would more be available with a higher limit?
   llvm::Optional<FuzzyMatcher> Filter; // Initialized once Sema runs.
 
 public:
   // A CodeCompleteFlow object is only useful for calling run() exactly once.
   CodeCompleteFlow(PathRef FileName, const CodeCompleteOptions &Opts)
-      : FileName(FileName), Opts(Opts),
-        RecorderOwner(new CompletionRecorder(Opts)), Recorder(*RecorderOwner) {}
+      : FileName(FileName), Opts(Opts) {}
 
   CompletionList run(const SemaCompleteInput &SemaCCInput) && {
     trace::Span Tracer("CodeCompleteFlow");
     // We run Sema code completion first. It builds an AST and calculates:
-    //   - completion results based on the AST. These are saved for merging.
+    //   - completion results based on the AST.
     //   - partial identifier and context. We need these for the index query.
     CompletionList Output;
-    semaCodeComplete(std::move(RecorderOwner), Opts.getClangCompleteOpts(),
-                     SemaCCInput, [&] {
-                       if (Recorder.CCSema) {
-                         Output = runWithSema();
-                         SPAN_ATTACH(
-                             Tracer, "sema_completion_kind",
-                             getCompletionKindString(Recorder.CCContext.getKind()));
-                       } else
-                         log("Code complete: no Sema callback, 0 results");
-                     });
+    auto RecorderOwner = llvm::make_unique<CompletionRecorder>(Opts, [&]() {
+      assert(Recorder && "Recorder is not set");
+      SPAN_ATTACH(Tracer, "sema_completion_kind",
+                  getCompletionKindString(Recorder->CCContext.getKind()));
 
+    Recorder = RecorderOwner.get();
+    semaCodeComplete(std::move(RecorderOwner), Opts.getClangCompleteOpts(),
+                     SemaCCInput);
+
     SPAN_ATTACH(Tracer, "sema_results", NSema);
     SPAN_ATTACH(Tracer, "index_results", NIndex);
     SPAN_ATTACH(Tracer, "merged_results", NBoth);
     SPAN_ATTACH(Tracer, "returned_results", Output.items.size());
     SPAN_ATTACH(Tracer, "incomplete", Output.isIncomplete);
     log(llvm::formatv("Code complete: {0} results from Sema, {1} from Index, "
                       "{2} matched, {3} returned{4}.",
                       NSema, NIndex, NBoth, Output.items.size(),
                       Output.isIncomplete ? " (incomplete)" : ""));
     assert(!Opts.Limit || Output.items.size() <= Opts.Limit);
     // We don't assert that isIncomplete means we hit a limit.
     // Indexes may choose to impose their own limits even if we don't have one.
     return Output;
   }
 
 private:
   // This is called by run() once Sema code completion is done, but before the
   // Sema data structures are torn down. It does all the real work.
   CompletionList runWithSema() {
     Filter = FuzzyMatcher(
-        Recorder.CCSema->getPreprocessor().getCodeCompletionFilter());
+        Recorder->CCSema->getPreprocessor().getCodeCompletionFilter());
     // Sema provides the needed context to query the index.
     // FIXME: in addition to querying for extra/overlapping symbols, we should
     //        explicitly request symbols corresponding to Sema results.
     //        We can use their signals even if the index can't suggest them.
     // We must copy index results to preserve them, but there are at most Limit.
     auto IndexResults = queryIndex();
     // Merge Sema and Index results, score them, and pick the winners.
-    auto Top = mergeResults(Recorder.Results, IndexResults);
+    auto Top = mergeResults(Recorder->Results, IndexResults);
     // Convert the results to the desired LSP structs.
     CompletionList Output;
     for (auto &C : Top)
       Output.items.push_back(toCompletionItem(C.first, C.second));
     Output.isIncomplete = Incomplete;
     return Output;
   }
 
   SymbolSlab queryIndex() {
-    if (!Opts.Index || !allowIndex(Recorder.CCContext.getKind()))
+    if (!Opts.Index || !allowIndex(Recorder->CCContext.getKind()))
       return SymbolSlab();
     trace::Span Tracer("Query index");
     SPAN_ATTACH(Tracer, "limit", Opts.Limit);
 
     SymbolSlab::Builder ResultsBuilder;
     // Build the query.
     FuzzyFindRequest Req;
     if (Opts.Limit)
       Req.MaxCandidateCount = Opts.Limit;
     Req.Query = Filter->pattern();
-    Req.Scopes =
-        getQueryScopes(Recorder.CCContext, Recorder.CCSema->getSourceManager());
+    Req.Scopes = getQueryScopes(Recorder->CCContext,
+                                Recorder->CCSema->getSourceManager());
     log(llvm::formatv("Code complete: fuzzyFind(\"{0}\", scopes=[{1}])",
                       Req.Query,
                       llvm::join(Req.Scopes.begin(), Req.Scopes.end(), ",")));
     // Run the query against the index.
     if (Opts.Index->fuzzyFind(
             Req, [&](const Symbol &Sym) { ResultsBuilder.insert(Sym); }))
       Incomplete = true;
     return std::move(ResultsBuilder).build();
auto CorrespondingIndexResult =
         if (I != IndexResults.end()) {
           UsedIndexResults.insert(&*I);
           return &*I;
         }
       }
       return nullptr;
     };
     // Emit all Sema results, merging them with Index results if possible.
-    for (auto &SemaResult : Recorder.Results)
+    for (auto &SemaResult : Recorder->Results)
       addCandidate(Top, &SemaResult, CorrespondingIndexResult(SemaResult));
     // Now emit any Index-only results.
     for (const auto &IndexResult : IndexResults) {
       if (UsedIndexResults.count(&IndexResult))
         continue;
       addCandidate(Top, /*SemaResult=*/nullptr, &IndexResult);
     }
     return std::move(Top).items();
   }
 
   // Scores a candidate and adds it to the TopN structure.
   void addCandidate(TopN &Candidates, const CodeCompletionResult *SemaResult,
                     const Symbol *IndexResult) {
     CompletionCandidate C;
     C.SemaResult = SemaResult;
     C.IndexResult = IndexResult;
-    C.Name = IndexResult ? IndexResult->Name : Recorder.getName(*SemaResult);
+    C.Name = IndexResult ? IndexResult->Name : Recorder->getName(*SemaResult);
 
     CompletionItemScores Scores;
     if (auto FuzzyScore = Filter->match(C.Name))
       Scores.filterScore = *FuzzyScore;
     else
       return;
     Scores.symbolScore = C.score();
     // We score candidates by multiplying symbolScore ("quality" of the result)
     // with filterScore (how well it matched the query).
     // This is sensitive to the distribution of both component scores!
     Scores.finalScore = Scores.filterScore * Scores.symbolScore;
 
     NSema += bool(SemaResult);
     NIndex += bool(IndexResult);
     NBoth += SemaResult && IndexResult;
     if (Candidates.push({C, Scores}))
       Incomplete = true;
   }
 
   CompletionItem toCompletionItem(const CompletionCandidate &Candidate,
                                   const CompletionItemScores &Scores) {
     CodeCompletionString *SemaCCS = nullptr;
     if (auto *SR = Candidate.SemaResult)
-      SemaCCS = Recorder.codeCompletionString(*SR, Opts.IncludeBriefComments);
+      SemaCCS = Recorder->codeCompletionString(*SR, Opts.IncludeBriefComments);
     return Candidate.build(FileName, Scores, Opts, SemaCCS);
   }
 };
 
 CompletionList codeComplete(PathRef FileName,
                             const tooling::CompileCommand &Command,
                             PrecompiledPreamble const *Preamble,
                             StringRef Contents, Position Pos,