1    :  // Copyright 2015 Google Inc. All Rights Reserved.
   2    :  //
   3    :  // Licensed under the Apache License, Version 2.0 (the "License");
   4    :  // you may not use this file except in compliance with the License.
   5    :  // You may obtain a copy of the License at
   6    :  //
   7    :  //     http://www.apache.org/licenses/LICENSE-2.0
   8    :  //
   9    :  // Unless required by applicable law or agreed to in writing, software
  10    :  // distributed under the License is distributed on an "AS IS" BASIS,
  11    :  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12    :  // See the License for the specific language governing permissions and
  13    :  // limitations under the License.
  14    :  //
  15    :  // The HotPatchingWriter allocates executable memory pages and writes blocks
  16    :  // into this memory area, where they can be executed.
  17    :  //
  18    :  // First the |Init| function must be called. This allocates a new executable
  19    :  // virtual memory of a given size using VirtualAlloc.
  20    :  //
  21    :  // The |Write| function can be used to write a block into this memory area.
  22    :  // It does the following:
  23    :  // - Copies the block data into the executable virtual memory.
  24    :  // - Changes the data of the block to point into the new location (the block
  25    :  //     will not own the data)
  26    :  // - Finalizes the bytes of inter-block references in the block data.
  27    :  //
  28    :  // NOTE: To write a block with this class, the |data| of all referred blocks
  29    :  //     must be backed by executable memory: they must be either blocks
  30    :  //     decomposed by the hot patching decomposer or blocks already written by
  31    :  //     the writer. The reason for this restriction is because the references
  32    :  //     will be calculated using the |data| pointers of the blocks.
  33    :  // TODO(cseri): The precondition is currently not checked. We could introduce
  34    :  //     a new flag for in-memory executable blocks.
  35    :  //
  36    :  // TODO(cseri): Implement some page protection logic, the write permission
  37    :  //     should be removed after the writes are finished.
  38    :  // TODO(cseri): Consider freeing the allocated virtual memory in destructor.
  39    :  
  40    :  #ifndef SYZYGY_PE_HOT_PATCHING_WRITER_H_
  41    :  #define SYZYGY_PE_HOT_PATCHING_WRITER_H_
  42    :  
  43    :  #include "syzygy/block_graph/block_graph.h"
  44    :  
  45  m :  namespace pe {
  46    :  
  47  m :  class HotPatchingWriter {
  48  m :   public:
  49  m :    typedef block_graph::BlockGraph BlockGraph;
  50  m :    typedef void* FunctionPointer;
  51    :  
  52  m :    HotPatchingWriter();
  53  m :    ~HotPatchingWriter();
  54    :  
  55    :    // Queries how much of the virtual memory of the writer has used so far.
  56    :    // @returns the size of the written code, in bytes.
  57  m :    size_t GetUsedMemory() const;
  58    :  
  59    :    // Allocates an executable virtual page with a given size that will be used to
  60    :    // write the code into.
  61    :    // @param virtual_memory_size The size of the requested virtual memory, in
  62    :    //    bytes.
  63  m :    bool Init(size_t virtual_memory_size);
  64    :  
  65    :    // Writes a block into the executable memory.
  66    :    // @param block The block to write.
  67    :    // @returns A pointer to the written function on success, nullptr if there was
  68    :    //     not enough space to write the function.
  69    :    // @pre |Init| must be called before a call to this function.
  70  m :    FunctionPointer Write(BlockGraph::Block* block);
  71    :  
  72    :    // Returns the size of the allocated virtual memory. Valid after the |Init|
  73    :    // function is called.
  74    :    // @returns the size of the virtual memory, in bytes.
  75  m :    size_t virtual_memory_size() const {
  76  m :      return virtual_memory_size_;
  77  m :    }
  78    :  
  79  m :   protected:
  80    :    // The pointer to the virtual memory.
  81  m :    LPVOID virtual_memory_;
  82    :  
  83    :    // The size of the allocated virtual memory. Valid after the |Init| function
  84    :    // is called.
  85  m :    size_t virtual_memory_size_;
  86    :  
  87    :    // The pointer to the current position in the virtual memory.
  88  m :    uint8* virtual_memory_cursor_;
  89    :  
  90  m :   private:
  91  m :    DISALLOW_COPY_AND_ASSIGN(HotPatchingWriter);
  92  m :  };
  93    :  
  94  m :  }  // namespace pe
  95    :  
  96    :  #endif  // SYZYGY_PE_HOT_PATCHING_WRITER_H_