Switch out use of std::string and std::vector for large allocations for a buffer class that doesn...

courgette/disassembler.cc

 // Copyright (c) 2010 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #include "courgette/disassembler.h"
 
 #include <algorithm>
 #include <string>
 #include <vector>
 
@@ skipping 18 matching lines @@
         incomplete_disassembly_(false) {
   }
 
   virtual bool Disassemble(AssemblyProgram* target);
 
   virtual void Destroy() { delete this; }
 
  protected:
   PEInfo& pe_info() { return *pe_info_; }
 
-  void ParseFile(AssemblyProgram* target);
+  CheckBool ParseFile(AssemblyProgram* target);
   bool ParseAbs32Relocs();
   void ParseRel32RelocsFromSections();
   void ParseRel32RelocsFromSection(const Section* section);
 
-  void ParseNonSectionFileRegion(uint32 start_file_offset,
-                                 uint32 end_file_offset,
-                                 AssemblyProgram* program);
-  void ParseFileRegion(const Section* section,
-                       uint32 start_file_offset, uint32 end_file_offset,
-                       AssemblyProgram* program);
+  CheckBool ParseNonSectionFileRegion(uint32 start_file_offset,
+                                      uint32 end_file_offset,
+                                      AssemblyProgram* program);
+  CheckBool ParseFileRegion(const Section* section,
+                            uint32 start_file_offset, uint32 end_file_offset,
+                            AssemblyProgram* program);
 
 #if COURGETTE_HISTOGRAM_TARGETS
   void HistogramTargets(const char* kind, const std::map<RVA, int>& map);
 #endif
 
   PEInfo* pe_info_;
   bool incomplete_disassembly_;  // 'true' if can leave out 'uninteresting' bits
 
   std::vector<RVA> abs32_locations_;
   std::vector<RVA> rel32_locations_;
 
 #if COURGETTE_HISTOGRAM_TARGETS
   std::map<RVA, int> abs32_target_rvas_;
   std::map<RVA, int> rel32_target_rvas_;
 #endif
 };
 
 bool DisassemblerWin32X86::Disassemble(AssemblyProgram* target) {
   if (!pe_info().ok())
     return false;
 
   target->set_image_base(pe_info().image_base());
 
   if (!ParseAbs32Relocs())
     return false;
 
   ParseRel32RelocsFromSections();
 
-  ParseFile(target);
+  if (!ParseFile(target))
+    return false;
 
   target->DefaultAssignIndexes();
+
   return true;
 }
 
 static uint32 Read32LittleEndian(const void* address) {
   return *reinterpret_cast<const uint32*>(address);
 }
 
 bool DisassemblerWin32X86::ParseAbs32Relocs() {
   abs32_locations_.clear();
   if (!pe_info().ParseRelocs(&abs32_locations_))
@@ skipping 130 matching lines @@
         ++rel32_target_rvas_[target_rva];
 #endif
         p += 4;
         continue;
       }
     }
     p += 1;
   }
 }
 
-void DisassemblerWin32X86::ParseFile(AssemblyProgram* program) {
+CheckBool DisassemblerWin32X86::ParseFile(AssemblyProgram* program) {
+  bool ok = true;
   // Walk all the bytes in the file, whether or not in a section.
   uint32 file_offset = 0;
-  while (file_offset < pe_info().length()) {
+  while (ok && file_offset < pe_info().length()) {
     const Section* section = pe_info().FindNextSection(file_offset);
     if (section == NULL) {
       // No more sections.  There should not be extra stuff following last
       // section.
       //   ParseNonSectionFileRegion(file_offset, pe_info().length(), program);
       break;
     }
     if (file_offset < section->file_offset_of_raw_data) {
       uint32 section_start_offset = section->file_offset_of_raw_data;
-      ParseNonSectionFileRegion(file_offset, section_start_offset, program);
+      ok = ParseNonSectionFileRegion(file_offset, section_start_offset,
+                                     program);
       file_offset = section_start_offset;
     }
-    uint32 end = file_offset + section->size_of_raw_data;
-    ParseFileRegion(section, file_offset, end, program);
-    file_offset = end;
+    if (ok) {
+      uint32 end = file_offset + section->size_of_raw_data;
+      ok = ParseFileRegion(section, file_offset, end, program);
+      file_offset = end;
+    }
   }
 
 #if COURGETTE_HISTOGRAM_TARGETS
   HistogramTargets("abs32 relocs", abs32_target_rvas_);
   HistogramTargets("rel32 relocs", rel32_target_rvas_);
 #endif
+
+  return ok;
 }
 
-void DisassemblerWin32X86::ParseNonSectionFileRegion(
+CheckBool DisassemblerWin32X86::ParseNonSectionFileRegion(
     uint32 start_file_offset,
     uint32 end_file_offset,
     AssemblyProgram* program) {
   if (incomplete_disassembly_)
-    return;
+    return true;
 
   const uint8* start = pe_info().FileOffsetToPointer(start_file_offset);
   const uint8* end = pe_info().FileOffsetToPointer(end_file_offset);
 
   const uint8* p = start;
 
-  while (p < end) {
-    program->EmitByteInstruction(*p);
+  bool ok = true;
+  while (p < end && ok) {
+    ok = program->EmitByteInstruction(*p);
     ++p;
   }
+
+  return ok;
 }
 
-void DisassemblerWin32X86::ParseFileRegion(
+CheckBool DisassemblerWin32X86::ParseFileRegion(
     const Section* section,
     uint32 start_file_offset, uint32 end_file_offset,
     AssemblyProgram* program) {
   RVA relocs_start_rva = pe_info().base_relocation_table().address_;
 
   const uint8* start_pointer = pe_info().FileOffsetToPointer(start_file_offset);
   const uint8* end_pointer = pe_info().FileOffsetToPointer(end_file_offset);
 
   RVA start_rva = pe_info().FileOffsetToRVA(start_file_offset);
   RVA end_rva = start_rva + section->virtual_size;
 
   // Quick way to convert from Pointer to RVA within a single Section is to
   // subtract 'pointer_to_rva'.
   const uint8* const adjust_pointer_to_rva = start_pointer - start_rva;
 
   std::vector<RVA>::iterator rel32_pos = rel32_locations_.begin();
   std::vector<RVA>::iterator abs32_pos = abs32_locations_.begin();
 
-  program->EmitOriginInstruction(start_rva);
+  bool ok = program->EmitOriginInstruction(start_rva);
 
   const uint8* p = start_pointer;
 
-  while (p < end_pointer) {
+  while (ok && p < end_pointer) {
     RVA current_rva = static_cast<RVA>(p - adjust_pointer_to_rva);
 
     // The base relocation table is usually in the .relocs section, but it could
     // actually be anywhere.  Make sure we skip it because we will regenerate it
     // during assembly.
     if (current_rva == relocs_start_rva) {
-      program->EmitMakeRelocsInstruction();
+      ok = program->EmitMakeRelocsInstruction();
+      if (!ok)
+        break;
       uint32 relocs_size = pe_info().base_relocation_table().size_;
       if (relocs_size) {
         p += relocs_size;
         continue;
       }
     }
 
     while (abs32_pos != abs32_locations_.end() && *abs32_pos < current_rva)
       ++abs32_pos;
 
     if (abs32_pos != abs32_locations_.end() && *abs32_pos == current_rva) {
       uint32 target_address = Read32LittleEndian(p);
       RVA target_rva = target_address - pe_info().image_base();
       // TODO(sra): target could be Label+offset.  It is not clear how to guess
       // which it might be.  We assume offset==0.
-      program->EmitAbs32(program->FindOrMakeAbs32Label(target_rva));
+      ok = program->EmitAbs32(program->FindOrMakeAbs32Label(target_rva));
+      if (!ok)
+        break;
       p += 4;
       continue;
     }
 
     while (rel32_pos != rel32_locations_.end() && *rel32_pos < current_rva)
       ++rel32_pos;
 
     if (rel32_pos != rel32_locations_.end() && *rel32_pos == current_rva) {
       RVA target_rva = current_rva + 4 + Read32LittleEndian(p);
-      program->EmitRel32(program->FindOrMakeRel32Label(target_rva));
+      ok = program->EmitRel32(program->FindOrMakeRel32Label(target_rva));
       p += 4;
       continue;
     }
 
     if (incomplete_disassembly_) {
       if ((abs32_pos == abs32_locations_.end() || end_rva <= *abs32_pos) &&
           (rel32_pos == rel32_locations_.end() || end_rva <= *rel32_pos) &&
           (end_rva <= relocs_start_rva || current_rva >= relocs_start_rva)) {
         // No more relocs in this section, don't bother encoding bytes.
         break;
       }
     }
 
-    program->EmitByteInstruction(*p);
+    ok = program->EmitByteInstruction(*p);
     p += 1;
   }
+
+  return ok;
 }
 
 #if COURGETTE_HISTOGRAM_TARGETS
 // Histogram is printed to std::cout.  It is purely for debugging the algorithm
 // and is only enabled manually in 'exploration' builds.  I don't want to add
 // command-line configuration for this feature because this code has to be
 // small, which means compiled-out.
 void DisassemblerWin32X86::HistogramTargets(const char* kind,
                                             const std::map<RVA, int>& map) {
   int total = 0;
@@ skipping 67 matching lines @@
   delete pe_info;
   *output = program;
   return C_OK;
 }
 
 void DeleteAssemblyProgram(AssemblyProgram* program) {
   delete program;
 }
 
 }  // namespace courgette