Include attributes of zipped binaries in ClientDownloadRequests.

chrome/common/safe_browsing/binary_feature_extractor_win.cc

 // Copyright 2014 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 "chrome/common/safe_browsing/binary_feature_extractor.h"
 
 #include <windows.h>
 #include <softpub.h>
 #include <wintrust.h>
 
 #include "base/files/file_path.h"
 #include "base/files/memory_mapped_file.h"
 #include "base/logging.h"
 #include "chrome/common/safe_browsing/csd.pb.h"
 #include "chrome/common/safe_browsing/pe_image_reader_win.h"
 
 #pragma comment(lib, "wintrust.lib")
 
 namespace safe_browsing {
 
+namespace {
+
+bool ExtractImageHeadersImpl(
+    const base::MemoryMappedFile& file,
+    BinaryFeatureExtractor::ExtractHeadersOption options,
+    ClientDownloadRequest_ImageHeaders* image_headers) {
+  PeImageReader pe_image;
+  if (!pe_image.Initialize(file.data(), file.length()))
+    return false;
+
+  // Copy the headers.
+  ClientDownloadRequest_PEImageHeaders* pe_headers =
+      image_headers->mutable_pe_headers();
+  pe_headers->set_dos_header(pe_image.GetDosHeader(), sizeof(IMAGE_DOS_HEADER));
+  pe_headers->set_file_header(pe_image.GetCoffFileHeader(),
+                              sizeof(IMAGE_FILE_HEADER));
+  size_t optional_header_size = 0;
+  const uint8_t* optional_header_data =
+      pe_image.GetOptionalHeaderData(&optional_header_size);
+  if (pe_image.GetWordSize() == PeImageReader::WORD_SIZE_32) {
+    pe_headers->set_optional_headers32(optional_header_data,
+                                       optional_header_size);
+  } else {
+    pe_headers->set_optional_headers64(optional_header_data,
+                                       optional_header_size);
+  }
+  const size_t number_of_sections = pe_image.GetNumberOfSections();
+  for (size_t i = 0; i != number_of_sections; ++i) {
+    pe_headers->add_section_header(pe_image.GetSectionHeaderAt(i),
+                                   sizeof(IMAGE_SECTION_HEADER));
+  }
+  if (!(options & BinaryFeatureExtractor::kOmitExports)) {
+    size_t export_size = 0;
+    const uint8_t* export_section = pe_image.GetExportSection(&export_size);
+    if (export_section)
+      pe_headers->set_export_section_data(export_section, export_size);
+  }
+  size_t number_of_debug_entries = pe_image.GetNumberOfDebugEntries();
+  for (size_t i = 0; i != number_of_debug_entries; ++i) {
+    const uint8_t* raw_data = NULL;
+    size_t raw_data_size = 0;
+    const IMAGE_DEBUG_DIRECTORY* directory_entry =
+        pe_image.GetDebugEntry(i, &raw_data, &raw_data_size);
+    if (directory_entry) {
+      ClientDownloadRequest_PEImageHeaders_DebugData* debug_data =
+          pe_headers->add_debug_data();
+      debug_data->set_directory_entry(directory_entry,
+                                      sizeof(*directory_entry));
+      if (raw_data)
+        debug_data->set_raw_data(raw_data, raw_data_size);
+    }
+  }
+
+  return true;
+}
+
+}  // namespace
+
 BinaryFeatureExtractor::BinaryFeatureExtractor() {}
 
 BinaryFeatureExtractor::~BinaryFeatureExtractor() {}
 
 void BinaryFeatureExtractor::CheckSignature(
     const base::FilePath& file_path,
     ClientDownloadRequest_SignatureInfo* signature_info) {
   DVLOG(2) << "Checking signature for " << file_path.value();
 
   WINTRUST_FILE_INFO file_info = {0};
@@ skipping 56 matching lines @@
     wintrust_data.dwStateAction = WTD_STATEACTION_CLOSE;
     WinVerifyTrust(static_cast<HWND>(INVALID_HANDLE_VALUE),
                    &policy_guid, &wintrust_data);
   }
 }
 
 bool BinaryFeatureExtractor::ExtractImageHeaders(
     const base::FilePath& file_path,
     ExtractHeadersOption options,
     ClientDownloadRequest_ImageHeaders* image_headers) {
-  // Map the file into memory.
-  base::MemoryMappedFile file;
-  if (!file.Initialize(file_path))
+  base::MemoryMappedFile mapped_file;
+  if (!mapped_file.Initialize(file_path))
     return false;
+  return ExtractImageHeadersImpl(mapped_file, options, image_headers);
+}
 
-  PeImageReader pe_image;
-  if (!pe_image.Initialize(file.data(), file.length()))
+bool BinaryFeatureExtractor::ExtractImageHeadersFromFile(
+    base::File file,
+    ExtractHeadersOption options,
+    ClientDownloadRequest_ImageHeaders* image_headers) {
+  base::MemoryMappedFile mapped_file;
+  if (!mapped_file.Initialize(file.Pass()))
     return false;
-
-  // Copy the headers.
-  ClientDownloadRequest_PEImageHeaders* pe_headers =
-      image_headers->mutable_pe_headers();
-  pe_headers->set_dos_header(pe_image.GetDosHeader(), sizeof(IMAGE_DOS_HEADER));
-  pe_headers->set_file_header(pe_image.GetCoffFileHeader(),
-                              sizeof(IMAGE_FILE_HEADER));
-  size_t optional_header_size = 0;
-  const uint8_t* optional_header_data =
-      pe_image.GetOptionalHeaderData(&optional_header_size);
-  if (pe_image.GetWordSize() == PeImageReader::WORD_SIZE_32) {
-    pe_headers->set_optional_headers32(optional_header_data,
-                                       optional_header_size);
-  } else {
-    pe_headers->set_optional_headers64(optional_header_data,
-                                       optional_header_size);
-  }
-  const size_t number_of_sections = pe_image.GetNumberOfSections();
-  for (size_t i = 0; i != number_of_sections; ++i) {
-    pe_headers->add_section_header(pe_image.GetSectionHeaderAt(i),
-                                   sizeof(IMAGE_SECTION_HEADER));
-  }
-  if (!(options & kOmitExports)) {
-    size_t export_size = 0;
-    const uint8_t* export_section = pe_image.GetExportSection(&export_size);
-    if (export_section)
-      pe_headers->set_export_section_data(export_section, export_size);
-  }
-  size_t number_of_debug_entries = pe_image.GetNumberOfDebugEntries();
-  for (size_t i = 0; i != number_of_debug_entries; ++i) {
-    const uint8_t* raw_data = NULL;
-    size_t raw_data_size = 0;
-    const IMAGE_DEBUG_DIRECTORY* directory_entry =
-        pe_image.GetDebugEntry(i, &raw_data, &raw_data_size);
-    if (directory_entry) {
-      ClientDownloadRequest_PEImageHeaders_DebugData* debug_data =
-          pe_headers->add_debug_data();
-      debug_data->set_directory_entry(directory_entry,
-                                      sizeof(*directory_entry));
-      if (raw_data)
-        debug_data->set_raw_data(raw_data, raw_data_size);
-    }
-  }
-
-  return true;
+  return ExtractImageHeadersImpl(mapped_file, options, image_headers);
 }
 
 }  // namespace safe_browsing