[clang] Cleanup ASTContext before output files in crash recovery for modules

[benlangmuir]

Summary

When we recover from a crash in a module compilation thread, we need to ensure any output streams owned by the ASTConsumer (e.g. in RawPCHContainerGenerator) are deleted before we call clearOutputFiles(). This has the same theoretical issues with proxy streams that Duncan discusses in the commit 2d133867833fe8eb. In practice, this was observed as a use-after-free crash on a downstream branch that uses such a proxy stream in this code path.

[benlangmuir]

Updated per out-of-band feedback from @steven_wu :

• Added an assert to clearOutputFiles that the ASTConsumer is removed. Normally this would be done in FrontendAction::EndSourceFile. This should help avoid regressions and also means it is covered by existing tests.
• Updated FailureCleanup in FrontendAction::BeginSourceFile to reset the ASTConsumer. This is needed due to the assertion, but I don't think it changes much in practice since there would also be no output files at this stage.

[steven_wu]

LGTM. Thanks!

clang/lib/Frontend/CompilerInstance.cpp

if (ExternalSemaSrc) {
     TheSema->addExternalSource(ExternalSemaSrc.get());
     ExternalSemaSrc->InitializeSema(*TheSema);
   }
 }
 
 // Output Files
 
 void CompilerInstance::clearOutputFiles(bool EraseFiles) {
+  // The ASTConsumer can own streams that write to the output files.
+  assert(!hasASTConsumer() && "ASTConsumer should be reset");
   // Ignore errors that occur when trying to discard the temp file.
   for (OutputFile &OF : OutputFiles) {
     if (EraseFiles) {
       if (OF.File)
         consumeError(OF.File->discard());
       if (!OF.Filename.empty())
         llvm::sys::fs::remove(OF.Filename);
       continue;
compileModuleImpl(CompilerInstance &ImportingInstance, SourceLocation ImportLoc,
   ImportingInstance.getDiagnostics().Report(ImportLoc,
                                             diag::remark_module_build)
     << ModuleName << ModuleFileName;
 
   PreBuildStep(Instance);
 
   // Execute the action to actually build the module in-place. Use a separate
   // thread so that we get a stack large enough.
-  llvm::CrashRecoveryContext CRC;
-  CRC.RunSafelyOnThread(
+  bool Crashed = !llvm::CrashRecoveryContext().RunSafelyOnThread(
       [&]() {
         GenerateModuleFromModuleMapAction Action;
         Instance.ExecuteAction(Action);
       },
       DesiredStackSize);
 
   PostBuildStep(Instance);
 
   ImportingInstance.getDiagnostics().Report(ImportLoc,
                                             diag::remark_module_build_done)
     << ModuleName;
 
-  // Delete any remaining temporary files related to Instance, in case the
-  // module generation thread crashed.
+  if (Crashed) {
+    // Clear the ASTConsumer if it hasn't been already, in case it owns streams
+    // that must be closed before clearing output files.
+    Instance.setSema(nullptr);
+    Instance.setASTConsumer(nullptr);
+
+    // Delete any remaining temporary files related to Instance.
     Instance.clearOutputFiles(/*EraseFiles=*/true);
+  }
 
   // If \p AllowPCMWithCompilerErrors is set return 'success' even if errors
   // occurred.
   return !Instance.getDiagnostics().hasErrorOccurred() ||
          Instance.getFrontendOpts().AllowPCMWithCompilerErrors;
 }
 
 static const FileEntry *getPublicModuleMap(const FileEntry *File,

clang/lib/Frontend/FrontendAction.cpp

bool FrontendAction::BeginSourceFile(CompilerInstance &CI,
   bool ReplayASTFile = Input.getKind().getFormat() == InputKind::Precompiled &&
                        usesPreprocessorOnly();
 
   // If we fail, reset state since the client will not end up calling the
   // matching EndSourceFile(). All paths that return true should release this.
   auto FailureCleanup = llvm::make_scope_exit([&]() {
     if (HasBegunSourceFile)
       CI.getDiagnosticClient().EndSourceFile();
+    CI.setASTConsumer(nullptr);
     CI.clearOutputFiles(/*EraseFiles=*/true);
     CI.getLangOpts().setCompilingModule(LangOptions::CMK_None);
     setCurrentInput(FrontendInputFile());
     setCompilerInstance(nullptr);
   });
 
   if (!BeginInvocation(CI))
     return false;