Fix use-after-free in dex2oat mini-debug-info compression.

This might happen if the device runs out of disk space,
which causes the WriteOutputFiles to return early and
dex2oat starts to shutdown and free memory before exit.

However, the mini-debug-info compression thread still
keeps running, so there is short window of time where
it might access freed memory before dex2oat exits.

This CL uses RAII to ensure that the compression task
finishes before return from the frame that started it.

Bug: 375314886
Test: test.py --host
Change-Id: I69b679551477de5e3df9c0867daefd8e22c52833

3 files changed, 27 insertions, 19 deletions

diff --git a/dex2oat/dex2oat.cc b/dex2oat/dex2oat.cc
index bcc51182b7..09dd6bc6e3 100644
--- a/dex2oat/dex2oat.cc
+++ b/dex2oat/dex2oat.cc
@@ -2154,7 +2154,10 @@ class Dex2Oat final {
         // We need to mirror the layout of the ELF file in the compressed debug-info.
         // Therefore PrepareDebugInfo() relies on the SetLoadedSectionSizes() call further above.
         debug::DebugInfo debug_info = oat_writer->GetDebugInfo();  // Keep the variable alive.
-        elf_writer->PrepareDebugInfo(debug_info);  // Processes the data on background thread.
+        // This will perform the compression on background thread while we do other I/O below.
+        // If we hit any ERROR path below, the destructor of this variable will wait for the
+        // task to finish (since it accesses the 'debug_info' above and other 'Dex2Oat' data).
+        std::unique_ptr<ThreadPool> compression_job = elf_writer->PrepareDebugInfo(debug_info);
 
         OutputStream* rodata = rodata_[i];
         DCHECK(rodata != nullptr);
diff --git a/dex2oat/linker/elf_writer.h b/dex2oat/linker/elf_writer.h
index 35e3565592..d27d4fa92c 100644
--- a/dex2oat/linker/elf_writer.h
+++ b/dex2oat/linker/elf_writer.h
@@ -18,6 +18,7 @@
 #define ART_DEX2OAT_LINKER_ELF_WRITER_H_
 
 #include <stdint.h>
+
 #include <cstddef>
 #include <string>
 #include <vector>
@@ -27,6 +28,7 @@
 #include "base/mutex.h"
 #include "base/os.h"
 #include "debug/debug_info.h"
+#include "thread_pool.h"
 
 namespace art {
 
@@ -66,7 +68,7 @@ class ElfWriter {
                                      size_t bss_methods_offset,
                                      size_t bss_roots_offset,
                                      size_t dex_section_size) = 0;
-  virtual void PrepareDebugInfo(const debug::DebugInfo& debug_info) = 0;
+  virtual std::unique_ptr<ThreadPool> PrepareDebugInfo(const debug::DebugInfo& debug_info) = 0;
   virtual OutputStream* StartRoData() = 0;
   virtual void EndRoData(OutputStream* rodata) = 0;
   virtual OutputStream* StartText() = 0;
diff --git a/dex2oat/linker/elf_writer_quick.cc b/dex2oat/linker/elf_writer_quick.cc
index f87ca6d81e..425f9bd554 100644
--- a/dex2oat/linker/elf_writer_quick.cc
+++ b/dex2oat/linker/elf_writer_quick.cc
@@ -40,21 +40,22 @@ namespace linker {
 
 class DebugInfoTask : public Task {
  public:
-  DebugInfoTask(InstructionSet isa,
+  DebugInfoTask(ThreadPool* owner,
+                InstructionSet isa,
                 const InstructionSetFeatures* features,
                 uint64_t text_section_address,
                 size_t text_section_size,
                 uint64_t dex_section_address,
                 size_t dex_section_size,
                 const debug::DebugInfo& debug_info)
-      : isa_(isa),
+      : owner_(owner),
+        isa_(isa),
         instruction_set_features_(features),
         text_section_address_(text_section_address),
         text_section_size_(text_section_size),
         dex_section_address_(dex_section_address),
         dex_section_size_(dex_section_size),
-        debug_info_(debug_info) {
-  }
+        debug_info_(debug_info) {}
 
   void Run(Thread*) override {
     result_ = debug::MakeMiniDebugInfo(isa_,
@@ -66,11 +67,13 @@ class DebugInfoTask : public Task {
                                        debug_info_);
   }
 
-  std::vector<uint8_t>* GetResult() {
+  std::vector<uint8_t>* WaitAndGetMiniDebugInfo() {
+    owner_->Wait(Thread::Current(), true, false);
     return &result_;
   }
 
  private:
+  ThreadPool* owner_;
   InstructionSet isa_;
   const InstructionSetFeatures* instruction_set_features_;
   uint64_t text_section_address_;
@@ -97,7 +100,7 @@ class ElfWriterQuick final : public ElfWriter {
                              size_t bss_methods_offset,
                              size_t bss_roots_offset,
                              size_t dex_section_size) override;
-  void PrepareDebugInfo(const debug::DebugInfo& debug_info) override;
+  std::unique_ptr<ThreadPool> PrepareDebugInfo(const debug::DebugInfo& debug_info) override;
   OutputStream* StartRoData() override;
   void EndRoData(OutputStream* rodata) override;
   OutputStream* StartText() override;
@@ -127,7 +130,6 @@ class ElfWriterQuick final : public ElfWriter {
   std::unique_ptr<BufferedOutputStream> output_stream_;
   std::unique_ptr<ElfBuilder<ElfTypes>> builder_;
   std::unique_ptr<DebugInfoTask> debug_info_task_;
-  std::unique_ptr<ThreadPool> debug_info_thread_pool_;
 
   void ComputeFileBuildId(uint8_t (*build_id)[ElfBuilder<ElfTypes>::kBuildIdLen]);
 
@@ -245,11 +247,15 @@ void ElfWriterQuick<ElfTypes>::WriteDynamicSection() {
 }
 
 template <typename ElfTypes>
-void ElfWriterQuick<ElfTypes>::PrepareDebugInfo(const debug::DebugInfo& debug_info) {
+std::unique_ptr<ThreadPool> ElfWriterQuick<ElfTypes>::PrepareDebugInfo(
+    const debug::DebugInfo& debug_info) {
+  std::unique_ptr<ThreadPool> thread_pool;
   if (compiler_options_.GetGenerateMiniDebugInfo()) {
+    thread_pool.reset(ThreadPool::Create("Mini-debug-info writer", 1));
     // Prepare the mini-debug-info in background while we do other I/O.
     Thread* self = Thread::Current();
     debug_info_task_ = std::make_unique<DebugInfoTask>(
+        thread_pool.get(),
         builder_->GetIsa(),
         compiler_options_.GetInstructionSetFeatures(),
         builder_->GetText()->GetAddress(),
@@ -257,24 +263,21 @@ void ElfWriterQuick<ElfTypes>::PrepareDebugInfo(const debug::DebugInfo& debug_in
         builder_->GetDex()->Exists() ? builder_->GetDex()->GetAddress() : 0,
         dex_section_size_,
         debug_info);
-    debug_info_thread_pool_.reset(ThreadPool::Create("Mini-debug-info writer", 1));
-    debug_info_thread_pool_->AddTask(self, debug_info_task_.get());
-    debug_info_thread_pool_->StartWorkers(self);
+    thread_pool->AddTask(self, debug_info_task_.get());
+    thread_pool->StartWorkers(self);
   }
+  return thread_pool;
 }
 
 template <typename ElfTypes>
 void ElfWriterQuick<ElfTypes>::WriteDebugInfo(const debug::DebugInfo& debug_info) {
+  std::unique_ptr<ThreadPool> thread_pool;
   if (compiler_options_.GetGenerateMiniDebugInfo()) {
     // If mini-debug-info wasn't explicitly created so far, create it now (happens in tests).
     if (debug_info_task_ == nullptr) {
-      PrepareDebugInfo(debug_info);
+      thread_pool = PrepareDebugInfo(debug_info);
     }
-    // Wait for the mini-debug-info generation to finish and write it to disk.
-    Thread* self = Thread::Current();
-    DCHECK(debug_info_thread_pool_ != nullptr);
-    debug_info_thread_pool_->Wait(self, true, false);
-    builder_->WriteSection(".gnu_debugdata", debug_info_task_->GetResult());
+    builder_->WriteSection(".gnu_debugdata", debug_info_task_->WaitAndGetMiniDebugInfo());
   }
   // The Strip method expects debug info to be last (mini-debug-info is not stripped).
   if (!debug_info.Empty() && compiler_options_.GetGenerateDebugInfo()) {