Added all code for integrity check transform

syzygy/pe/pe_file_writer.cc

 // Copyright 2012 Google Inc. All Rights Reserved.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
 // You may obtain a copy of the License at
 //
 //     http://www.apache.org/licenses/LICENSE-2.0
 //
 // Unless required by applicable law or agreed to in writing, software
 // distributed under the License is distributed on an "AS IS" BASIS,
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
 #include "syzygy/pe/pe_file_writer.h"
 
 #include <windows.h>
 #include <winnt.h>
 #include <imagehlp.h>  // NOLINT
 
 #include "base/logging.h"
 #include "base/files/file_util.h"
 #include "base/win/scoped_handle.h"
-#include "syzygy/common/buffer_parser.h"
 #include "syzygy/common/com_utils.h"
 #include "syzygy/pe/pe_utils.h"
 
 namespace pe {
 
 using block_graph::BlockGraph;
 using common::BinaryBufferParser;
 using core::AbsoluteAddress;
 using core::FileOffsetAddress;
 using core::RelativeAddress;
 using pe::ImageLayout;
 
 namespace {
 
-template <class Type>
-bool UpdateReference(size_t start, Type new_value, std::vector<uint8_t>* data) {
-  BinaryBufferParser parser(&data->at(0), data->size());
-
-  Type* ref_ptr = NULL;
-  if (!parser.GetAtIgnoreAlignment(start,
-                                   const_cast<const Type**>(&ref_ptr))) {
-    LOG(ERROR) << "Reference data not in block";
-    return false;
-  }
-  *ref_ptr = new_value;
-
-  return true;
-}
-
 // Returns the type of padding byte to use for a given section. Int3s will be
 // used for executable sections, nulls for everything else.
 uint8_t GetSectionPaddingByte(const ImageLayout& image_layout,
                               size_t section_index) {
   const uint8_t kZero = 0;
   const uint8_t kInt3 = 0xCC;
 
   if (section_index == BlockGraph::kInvalidSectionId)
     return kZero;
   DCHECK_GT(image_layout.sections.size(), section_index);
@@ skipping 254 matching lines @@
       return false;
     }
 
     previous_section_end = section_start + section_size;
     previous_section_file_end = section_file_start + section_file_size;
   }
 
   return true;
 }
 
+AbsoluteAddress* PEFileWriter::GetImageBase() {
+  DCHECK(nt_headers_ != nullptr);
+  AbsoluteAddress *addr =
+    new AbsoluteAddress(nt_headers_->OptionalHeader.ImageBase);
+  return addr;
+}
+
+const size_t PEFileWriter::GetImageSize() {
+  uint32_t last_section_addr = image_layout_.sections.back().addr.value();
+  size_t last_section_size = image_layout_.sections.back().size;
+  return last_section_addr + last_section_size;
+}
+
 bool PEFileWriter::WriteBlocks(FILE* file) {
   DCHECK(file != NULL);
 
   AbsoluteAddress image_base(nt_headers_->OptionalHeader.ImageBase);
 
   // Create the output buffer, reserving enough room for the whole file.
   DCHECK(!image_layout_.sections.empty());
   size_t last_section_index = image_layout_.sections.size() - 1;
   size_t image_size = section_file_range_map_[last_section_index].end().value();
   std::vector<uint8_t> buffer;
@@ skipping 16 matching lines @@
     BlockGraph::Block* block = const_cast<BlockGraph::Block*>(block_it->second);
 
     // If we're jumping to a new section output the necessary padding.
     if (block->section() != section_id) {
       FlushSection(section_index, &buffer);
       section_id = block->section();
       section_index++;
       DCHECK_GT(image_layout_.sections.size(), section_index);
     }
 
-    if (!WriteOneBlock(image_base, section_index, block, &buffer)) {
+    if (!WriteOneBlock(image_base, section_index, block, &buffer, nullptr)) {
       LOG(ERROR) << "Failed to write block \"" << block->name() << "\".";
       return false;
     }
   }
 
   FlushSection(last_section_index, &buffer);
   DCHECK_EQ(image_size, buffer.size());
 
   // Write the whole image to disk in one go.
   if (::fwrite(&buffer[0], sizeof(buffer[0]), buffer.size(), file) !=
@@ skipping 20 matching lines @@
 
   uint8_t padding_byte = GetSectionPaddingByte(image_layout_, section_index);
   buffer->resize(section_file_end, padding_byte);
 
   return;
 }
 
 bool PEFileWriter::WriteOneBlock(AbsoluteAddress image_base,
                                  size_t section_index,
                                  const BlockGraph::Block* block,
-                                 std::vector<uint8_t>* buffer) {
+                                 std::vector<uint8_t>* buffer,
+                                 FileOffsetAddress* offset_block) {
   // This function walks through the data referred by the input block, and
   // patches it to reflect the addresses and offsets of the blocks
   // referenced before writing the block's data to the file.
   DCHECK(block != NULL);
   DCHECK(buffer != NULL);
 
   RelativeAddress addr;
   if (!image_layout_.blocks.GetAddressOf(block, &addr)) {
     LOG(ERROR) << "All blocks must have an address.";
     return false;
@@ skipping 47 matching lines @@
   // and the explicit block data).
   if (file_offs + inited_data_size > section_file_range.end()) {
     LOG(ERROR) << "Initialized portion of block data lies outside of section.";
     return false;
   }
 
   // Add any necessary padding to get us to the block offset.
   if (buffer->size() < file_offs.value())
     buffer->resize(file_offs.value(), padding_byte);
 
+  if (offset_block != nullptr)
+    offset_block->set_value(buffer->size());
+
   // Copy the block data into the buffer.
   buffer->insert(buffer->end(),
                  block->data(),
                  block->data() + block->data_size());
 
   // We now want to append zeros for the implicit portion of the block data.
   size_t trailing_zeros = block->size() - block->data_size();
   if (trailing_zeros > 0) {
     // It is possible for a block to be laid out at the end of a section such
     // that part of its data lies within the virtual portion of the section.
@@ skipping 98 matching lines @@
       default:
         LOG(ERROR) << "Unsupported reference size.";
         return false;
     }
   }
 
   return true;
 }
 
 }  // namespace pe