[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–433

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

475

I think this check is dead - remove?

660

remove this comment

665

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

737

this tracer isn't needed anymore as discussed

816

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

852

remove "these are saved for merging"

[ilya-biryukov]

Addressed review comments

clangd/CodeComplete.cpp

if (!Info.UnresolvedQualifier->empty())
     *Info.UnresolvedQualifier += "::";
 
   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,