// -*- Mode: C++; c-basic-offset: 2; indent-tabs-mode: nil -*-
/* Copyright (c) 2007, Google Inc.
 * All rights reserved.
 * 
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are
 * met:
 * 
 *     * Redistributions of source code must retain the above copyright
 * notice, this list of conditions and the following disclaimer.
 *     * Redistributions in binary form must reproduce the above
 * copyright notice, this list of conditions and the following disclaimer
 * in the documentation and/or other materials provided with the
 * distribution.
 *     * Neither the name of Google Inc. nor the names of its
 * contributors may be used to endorse or promote products derived from
 * this software without specific prior written permission.
 * 
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
 * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
 * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
 * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
 * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
 * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *
 * ---
 * Author: Joi Sigurdsson
 * Author: Scott Francis
 *
 * Implementation of PreamblePatcher
 */

#include "preamble_patcher.h"

#include "mini_disassembler.h"

// Definitions of assembly statements we need
#define ASM_JMP32REL 0xE9
#define ASM_INT3 0xCC
#define ASM_NOP 0x90
// X64 opcodes
#define ASM_MOVRAX_IMM 0xB8
#define ASM_REXW 0x48
#define ASM_JMP 0xFF
#define ASM_JMP_RAX 0xE0
#define ASM_PUSH 0x68
#define ASM_RET 0xC3

namespace sidestep {

SideStepError PreamblePatcher::RawPatchWithStub(
    void* target_function,
    void* replacement_function,
    unsigned char* preamble_stub,
    unsigned long stub_size,
    unsigned long* bytes_needed) {
  if ((NULL == target_function) ||
      (NULL == replacement_function) ||
      (NULL == preamble_stub)) {
    SIDESTEP_ASSERT(false &&
                    "Invalid parameters - either pTargetFunction or "
                    "pReplacementFunction or pPreambleStub were NULL.");
    return SIDESTEP_INVALID_PARAMETER;
  }

  // TODO(V7:joi) Siggi and I just had a discussion and decided that both
  // patching and unpatching are actually unsafe.  We also discussed a
  // method of making it safe, which is to freeze all other threads in the
  // process, check their thread context to see if their eip is currently
  // inside the block of instructions we need to copy to the stub, and if so
  // wait a bit and try again, then unfreeze all threads once we've patched.
  // Not implementing this for now since we're only using SideStep for unit
  // testing, but if we ever use it for production code this is what we
  // should do.
  //
  // NOTE: Stoyan suggests we can write 8 or even 10 bytes atomically using
  // FPU instructions, and on newer processors we could use cmpxchg8b or
  // cmpxchg16b. So it might be possible to do the patching/unpatching
  // atomically and avoid having to freeze other threads.  Note though, that
  // doing it atomically does not help if one of the other threads happens
  // to have its eip in the middle of the bytes you change while you change
  // them.
  unsigned char* target = reinterpret_cast<unsigned char*>(target_function);
  unsigned int required_trampoline_bytes = 0;
  const unsigned int kRequiredStubJumpBytes = 5;
  const unsigned int kRequiredTargetPatchBytes = 5;

  // Initialize the stub with INT3's just in case.
  if (stub_size) {
    memset(preamble_stub, 0xcc, stub_size);
  }
  if (kIs64BitBinary) {
    // In 64-bit mode JMP instructions are always relative to RIP.  If the
    // replacement - target offset is > 2GB, we can't JMP to the replacement
    // function.  In this case, we're going to use a trampoline - that is,
    // we're going to do a relative jump to a small chunk of code in the stub
    // that will then do the absolute jump to the replacement function.  By
    // doing this, we only need to patch 5 bytes in the target function, as
    // opposed to patching 12 bytes if we were to do an absolute jump.
    //
    // Note that the first byte of the trampoline is a NOP instruction.  This
    // is used as a trampoline signature that will be detected when unpatching
    // the function.
    //
    // jmp <trampoline>
    //
    // trampoline:
    //    nop
    //    mov rax, <replacement_function>
    //    jmp rax
    //
    __int64 replacement_target_offset = reinterpret_cast<__int64>(
        replacement_function) - reinterpret_cast<__int64>(target) - 5;
    if (replacement_target_offset > INT_MAX
        || replacement_target_offset < INT_MIN) {
      // The stub needs to be within 2GB of the target for the trampoline to
      // work!
      __int64 trampoline_offset = reinterpret_cast<__int64>(preamble_stub)
          - reinterpret_cast<__int64>(target) - 5;
      if (trampoline_offset > INT_MAX || trampoline_offset < INT_MIN) {
        // We're screwed.
        SIDESTEP_ASSERT(false 
                       && "Preamble stub is too far from target to patch.");
        return SIDESTEP_UNEXPECTED;
      }
      required_trampoline_bytes = 13;
    }
  }

  // Let's disassemble the preamble of the target function to see if we can
  // patch, and to see how much of the preamble we need to take.  We need 5
  // bytes for our jmp instruction, so let's find the minimum number of
  // instructions to get 5 bytes.
  MiniDisassembler disassembler;
  unsigned int preamble_bytes = 0;
  unsigned int stub_bytes = 0;
  while (preamble_bytes < kRequiredTargetPatchBytes) {
    unsigned int cur_bytes = 0;
    InstructionType instruction_type =
        disassembler.Disassemble(target + preamble_bytes, cur_bytes);
    if (IT_JUMP == instruction_type) {
      unsigned int jump_bytes = 0;
      SideStepError jump_ret = SIDESTEP_JUMP_INSTRUCTION;
      if (IsShortConditionalJump(target + preamble_bytes, cur_bytes)) {
        jump_ret = PatchShortConditionalJump(target + preamble_bytes, cur_bytes,
                                             preamble_stub + stub_bytes,
                                             &jump_bytes,
                                             stub_size - stub_bytes);
      } else if (IsShortJump(target + preamble_bytes, cur_bytes)) {
        jump_ret = PatchShortJump(target + preamble_bytes, cur_bytes,
                                  preamble_stub + stub_bytes,
                                  &jump_bytes,
                                  stub_size - stub_bytes);
      } else if (IsNearConditionalJump(target + preamble_bytes, cur_bytes) ||
                 IsNearRelativeJump(target + preamble_bytes, cur_bytes) ||
                 IsNearAbsoluteCall(target + preamble_bytes, cur_bytes) ||
                 IsNearRelativeCall(target + preamble_bytes, cur_bytes)) {
         jump_ret = PatchNearJumpOrCall(target + preamble_bytes, cur_bytes,
                                        preamble_stub + stub_bytes, &jump_bytes,
                                        stub_size - stub_bytes);
      }
      if (jump_ret != SIDESTEP_SUCCESS) {
        SIDESTEP_ASSERT(false &&
                        "Unable to patch because there is an unhandled branch "
                        "instruction in the initial preamble bytes.");
        return SIDESTEP_JUMP_INSTRUCTION;
      }
      stub_bytes += jump_bytes;
    } else if (IT_RETURN == instruction_type) {
      SIDESTEP_ASSERT(false &&
                      "Unable to patch because function is too short");
      return SIDESTEP_FUNCTION_TOO_SMALL;
    } else if (IT_GENERIC == instruction_type) {
      if (IsMovWithDisplacement(target + preamble_bytes, cur_bytes)) {
        unsigned int mov_bytes = 0;
        if (PatchMovWithDisplacement(target + preamble_bytes, cur_bytes,
                                     preamble_stub + stub_bytes, &mov_bytes,
                                     stub_size - stub_bytes)
            != SIDESTEP_SUCCESS) {
          return SIDESTEP_UNSUPPORTED_INSTRUCTION;
        }
        stub_bytes += mov_bytes;
      } else {
        memcpy(reinterpret_cast<void*>(preamble_stub + stub_bytes),
               reinterpret_cast<void*>(target + preamble_bytes), cur_bytes);
        stub_bytes += cur_bytes;
      }
    } else {
      SIDESTEP_ASSERT(false &&
                      "Disassembler encountered unsupported instruction "
                      "(either unused or unknown");
      return SIDESTEP_UNSUPPORTED_INSTRUCTION;
    }
    preamble_bytes += cur_bytes;
  }

  if (NULL != bytes_needed)
    *bytes_needed = stub_bytes + kRequiredStubJumpBytes
        + required_trampoline_bytes;

  // Inv: cbPreamble is the number of bytes (at least 5) that we need to take
  // from the preamble to have whole instructions that are 5 bytes or more
  // in size total. The size of the stub required is cbPreamble +
  // kRequiredStubJumpBytes (5) + required_trampoline_bytes (0 or 13)
  if (stub_bytes + kRequiredStubJumpBytes + required_trampoline_bytes
      > stub_size) {
    SIDESTEP_ASSERT(false);
    return SIDESTEP_INSUFFICIENT_BUFFER;
  }

  // Now, make a jmp instruction to the rest of the target function (minus the
  // preamble bytes we moved into the stub) and copy it into our preamble-stub.
  // find address to jump to, relative to next address after jmp instruction
#ifdef _MSC_VER
#pragma warning(push)
#pragma warning(disable:4244)
#endif
  int relative_offset_to_target_rest
      = ((reinterpret_cast<unsigned char*>(target) + preamble_bytes) -
         (preamble_stub + stub_bytes + kRequiredStubJumpBytes));
#ifdef _MSC_VER
#pragma warning(pop)
#endif
  // jmp (Jump near, relative, displacement relative to next instruction)
  preamble_stub[stub_bytes] = ASM_JMP32REL;
  // copy the address
  memcpy(reinterpret_cast<void*>(preamble_stub + stub_bytes + 1),
         reinterpret_cast<void*>(&relative_offset_to_target_rest), 4);

  if (kIs64BitBinary && required_trampoline_bytes != 0) {
    // Construct the trampoline
    unsigned int trampoline_pos = stub_bytes + kRequiredStubJumpBytes;
    preamble_stub[trampoline_pos] = ASM_NOP;
    preamble_stub[trampoline_pos + 1] = ASM_REXW;
    preamble_stub[trampoline_pos + 2] = ASM_MOVRAX_IMM;
    memcpy(reinterpret_cast<void*>(preamble_stub + trampoline_pos + 3),
           reinterpret_cast<void*>(&replacement_function),
           sizeof(void *));
    preamble_stub[trampoline_pos + 11] = ASM_JMP;
    preamble_stub[trampoline_pos + 12] = ASM_JMP_RAX;

    // Now update replacement_function to point to the trampoline
    replacement_function = preamble_stub + trampoline_pos;
  }

  // Inv: preamble_stub points to assembly code that will execute the
  // original function by first executing the first cbPreamble bytes of the
  // preamble, then jumping to the rest of the function.

  // Overwrite the first 5 bytes of the target function with a jump to our
  // replacement function.
  // (Jump near, relative, displacement relative to next instruction)
  target[0] = ASM_JMP32REL;

  // Find offset from instruction after jmp, to the replacement function.
#ifdef _MSC_VER
#pragma warning(push)
#pragma warning(disable:4244)
#endif
  int offset_to_replacement_function =
      reinterpret_cast<unsigned char*>(replacement_function) -
      reinterpret_cast<unsigned char*>(target) - 5;
#ifdef _MSC_VER
#pragma warning(pop)
#endif
  // complete the jmp instruction
  memcpy(reinterpret_cast<void*>(target + 1),
         reinterpret_cast<void*>(&offset_to_replacement_function), 4);

  // Set any remaining bytes that were moved to the preamble-stub to INT3 so
  // as not to cause confusion (otherwise you might see some strange
  // instructions if you look at the disassembly, or even invalid
  // instructions). Also, by doing this, we will break into the debugger if
  // some code calls into this portion of the code.  If this happens, it
  // means that this function cannot be patched using this patcher without
  // further thought.
  if (preamble_bytes > kRequiredTargetPatchBytes) {
    memset(reinterpret_cast<void*>(target + kRequiredTargetPatchBytes),
           ASM_INT3, preamble_bytes - kRequiredTargetPatchBytes);
  }

  // Inv: The memory pointed to by target_function now points to a relative
  // jump instruction that jumps over to the preamble_stub.  The preamble
  // stub contains the first stub_size bytes of the original target
  // function's preamble code, followed by a relative jump back to the next
  // instruction after the first cbPreamble bytes.
  //
  // In 64-bit mode the memory pointed to by target_function *may* point to a
  // relative jump instruction that jumps to a trampoline which will then
  // perform an absolute jump to the replacement function.  The preamble stub
  // still contains the original target function's preamble code, followed by a
  // jump back to the instructions after the first preamble bytes.
  //
  return SIDESTEP_SUCCESS;
}

};  // namespace sidestep