Courgette: Add static method QuickDetect() to optimize program detection.

courgette/disassembler_win32.cc

 // Copyright 2016 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_win32.h"
 
 #include <stddef.h>
 #include <stdint.h>
 
 #include <algorithm>
 
 #include "base/logging.h"
 #include "courgette/assembly_program.h"
 #include "courgette/courgette.h"
 
 #if COURGETTE_HISTOGRAM_TARGETS
 #include <iostream>
 #endif
 
 namespace courgette {
 
-DisassemblerWin32::DisassemblerWin32(const void* start, size_t length)
+DisassemblerWin32::DisassemblerWin32(const uint8_t* start, size_t length)
     : Disassembler(start, length) {}
 
 RVA DisassemblerWin32::FileOffsetToRVA(FileOffset file_offset) const {
   for (int i = 0; i < number_of_sections_; ++i) {
     const Section* section = &sections_[i];
     if (file_offset >= section->file_offset_of_raw_data) {
       FileOffset offset_in_section =
           file_offset - section->file_offset_of_raw_data;
       if (offset_in_section < section->size_of_raw_data)
         return static_cast<RVA>(section->virtual_address + offset_in_section);
@@ skipping 298 matching lines @@
 
 std::string DisassemblerWin32::SectionName(const Section* section) {
   if (section == nullptr)
     return "<none>";
   char name[9];
   memcpy(name, section->name, 8);
   name[8] = '\0';  // Ensure termination.
   return name;
 }
 
+// static
+bool DisassemblerWin32::QuickDetect(const uint8_t* start,
+                                    size_t length,
+                                    uint16_t magic) {
+  if (length < kOffsetOfFileAddressOfNewExeHeader + 4 /* size */)
+    return false;
+
+  // Have 'MZ' magic for a DOS header?
+  if (start[0] != 'M' || start[1] != 'Z')
+    return false;
+
+  FileOffset file_offset = static_cast<FileOffset>(
+      ReadU32(start, kOffsetOfFileAddressOfNewExeHeader));
+  if (file_offset >= length || file_offset % 8 != 0)
+    return false;
+  const uint8_t* const pe_header = start + file_offset;
+  const size_t kMinPEHeaderSize = 4 /*signature*/ + kSizeOfCoffHeader;
+  if (pe_header <= start || pe_header + kMinPEHeaderSize >= start + length)
+    return false;
+
+  const uint8_t* optional_header = pe_header + 4 + kSizeOfCoffHeader;
+  // Check we can read the magic.
+  if (optional_header + 2 >= start + length)
+    return false;
+  if (magic != ReadU16(optional_header, 0))
+    return false;
+
+  return true;
+}
+
 RvaVisitor* DisassemblerWin32::CreateAbs32TargetRvaVisitor() {
   return new RvaVisitor_Abs32(abs32_locations_, *this);
 }
 
 RvaVisitor* DisassemblerWin32::CreateRel32TargetRvaVisitor() {
   return new RvaVisitor_Rel32(rel32_locations_, *this);
 }
 
 void DisassemblerWin32::RemoveUnusedRel32Locations(
     AssemblyProgram* program) {
@@ skipping 296 matching lines @@
     directory->size_ = static_cast<uint32_t>(size);
     return true;
   } else {
     directory->address_ = 0;
     directory->size_ = 0;
     return true;
   }
 }
 
 }  // namespace courgette