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

[Ben Langmuir via cfe-commits]

Author: Ben Langmuir
Date: 2022-07-07T10:23:57-07:00
New Revision: 67a84ec8105e590159b6303a1f0e3cb77c02b5fe

URL: https://github.com/llvm/llvm-project/commit/67a84ec8105e590159b6303a1f0e3cb77c02b5fe
DIFF: https://github.com/llvm/llvm-project/commit/67a84ec8105e590159b6303a1f0e3cb77c02b5fe.diff

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

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. Add an assertion so it won't regress.

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

rdar://96525032

Added: 
    

Modified: 
    clang/lib/Frontend/CompilerInstance.cpp
    clang/lib/Frontend/FrontendAction.cpp

Removed: 
    


################################################################################
diff  --git a/clang/lib/Frontend/CompilerInstance.cpp b/clang/lib/Frontend/CompilerInstance.cpp
index b982ca72c78ce..ba006b5d49f87 100644
--- a/clang/lib/Frontend/CompilerInstance.cpp
+++ b/clang/lib/Frontend/CompilerInstance.cpp
@@ -757,6 +757,8 @@ void CompilerInstance::createSema(TranslationUnitKind TUKind,
 // 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) {
@@ -1235,8 +1237,7 @@ compileModuleImpl(CompilerInstance &ImportingInstance, SourceLocation ImportLoc,
 
   // 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);
@@ -1249,9 +1250,15 @@ compileModuleImpl(CompilerInstance &ImportingInstance, SourceLocation ImportLoc,
                                             diag::remark_module_build_done)
     << ModuleName;
 
-  // Delete any remaining temporary files related to Instance, in case the
-  // module generation thread crashed.
-  Instance.clearOutputFiles(/*EraseFiles=*/true);
+  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.

diff  --git a/clang/lib/Frontend/FrontendAction.cpp b/clang/lib/Frontend/FrontendAction.cpp
index 81915e6330b03..ed3e314cc73bb 100644
--- a/clang/lib/Frontend/FrontendAction.cpp
+++ b/clang/lib/Frontend/FrontendAction.cpp
@@ -581,6 +581,7 @@ bool FrontendAction::BeginSourceFile(CompilerInstance &CI,
   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());