| Index: chrome/common/safe_browsing/binary_feature_extractor_win.cc
|
| diff --git a/chrome/common/safe_browsing/binary_feature_extractor_win.cc b/chrome/common/safe_browsing/binary_feature_extractor_win.cc
|
| index 9973f263356bbfe5e490b18a943519c497248a91..be7ceb4d1f7c5987e9b288a61d4583a6a3a8e3ca 100644
|
| --- a/chrome/common/safe_browsing/binary_feature_extractor_win.cc
|
| +++ b/chrome/common/safe_browsing/binary_feature_extractor_win.cc
|
| @@ -18,6 +18,64 @@
|
|
|
| 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() {}
|
| @@ -94,59 +152,20 @@ 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
|
|
|