tinyinst.h

/*
Copyright 2020 Google LLC

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

https ://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.
*/

#ifndef TINYINST_H
#define TINYINST_H

#include <list>
#include <set>
#include <map>
#include <unordered_map>
#include <unordered_set>

#if defined(WIN32) || defined(_WIN32) || defined(__WIN32)
  #include "Windows/debugger.h"
#elif __APPLE__
  #include "macOS/debugger.h"
#else
  #include "Linux/debugger.h"
#endif

#include "common.h"
#include "assembler.h"
#include "instruction.h"
#include "unwind.h"

class Hook;

#if defined(_WIN64)

#include "Windows/winunwind.h"

#elif __APPLE__

#include "macOS/unwindmacos.h"
class UnwindGeneratorMacOS;

#endif

// must be a power of two
#define JUMPTABLE_SIZE 0x2000

// we will allocate
// original_code_size * CODE_SIZE_MULTIPLIER +
// JUMPTABLE_SIZE * child_ptr_size
// for instrumented code
#ifdef ARM64
#define CODE_SIZE_MULTIPLIER 8
#else
#define CODE_SIZE_MULTIPLIER 4
#endif

typedef struct xed_decoded_inst_s xed_decoded_inst_t;

class ModuleInfo;


enum InstructionResult {
  INST_HANDLED,
  INST_NOTHANDLED,
  INST_STOPBB
};

class TinyInst : public Debugger {
public:
  virtual void Init(int argc, char **argv) override;

  void EnableInstrumentation() {
    instrumentation_disabled = false;
  }

  void DisableInstrumentation() {
    instrumentation_disabled = true;
  }

protected:

  enum IndirectInstrumentation {
    II_NONE,
    II_GLOBAL,
    II_LOCAL,
    II_AUTO
  };

  enum PatchModuleEntriesValue {
    OFF = 0,
    DATA = 1,
    CODE = 2,
    ALL = DATA | CODE
  };

  std::list<ModuleInfo *> instrumented_modules;

  struct CrossModuleLink {
    ModuleInfo *module1;
    ModuleInfo *module2;
    size_t offset1;
    size_t offset2;
  };

  virtual void OnEntrypoint() override;
  virtual void OnProcessCreated() override;
  virtual void OnProcessExit() override;
  virtual void OnModuleLoaded(void *module, char *module_name) override;
  virtual void OnModuleUnloaded(void *module) override;
  virtual bool OnException(Exception *exception_record) override;
  virtual void OnTargetMethodReached() override;
  virtual void OnCrashed(Exception *exception_record) override;

  virtual size_t GetTranslatedAddress(size_t address) override;

  void WriteCode(ModuleInfo *module, void *data, size_t size);
  void WriteCodeAtOffset(ModuleInfo *module, size_t offset, void *data, size_t size);
  void WritePointer(ModuleInfo *module, size_t value);
  void WritePointerAtOffset(ModuleInfo *module, size_t value, size_t offset);
  size_t ReadPointer(ModuleInfo *module, size_t offset);

  inline void FixDisp4(ModuleInfo *module, int32_t disp); 

  size_t GetCurrentInstrumentedAddress(ModuleInfo *module);
  void CommitCode(ModuleInfo *module, size_t start_offset, size_t size);

  ModuleInfo *GetModuleByName(const char *name);
  ModuleInfo *GetModule(size_t address);
  ModuleInfo *GetModuleFromInstrumented(size_t address);
  AddressRange *GetRegion(ModuleInfo *module, size_t address);

  // instrumentation API
  virtual void OnModuleEntered(ModuleInfo *module, size_t entry_address) {}
  virtual void InstrumentBasicBlock(ModuleInfo *module, size_t bb_address) {}
  virtual void InstrumentEdge(ModuleInfo *previous_module,
                              ModuleInfo *next_module,
                              size_t previous_address,
                              size_t next_address) {}

  virtual InstructionResult InstrumentInstruction(ModuleInfo *module,
                                                  Instruction& inst,
                                                  size_t bb_address,
                                                  size_t instruction_address);

  virtual void OnModuleInstrumented(ModuleInfo* module);
  virtual void OnModuleUninstrumented(ModuleInfo* module);
  virtual void OnBasicBlcokTranslated(ModuleInfo *module, size_t start_offset, size_t end_offset) { }

  int32_t sp_offset;
  Assembler* assembler_;

  UnwindGenerator* unwind_generator;
  virtual void OnReturnAddress(ModuleInfo *module, size_t original_address, size_t translated_address);
  
  void RegisterHook(Hook *hook);
  
  void InstrumentAddressRange(const char *name, size_t min_address, size_t max_address);

private:
  void AddInstrumentedModule(char* name, bool do_protect);
  bool HandleBreakpoint(void *address);
  void OnInstrumentModuleLoaded(void *module, ModuleInfo *target_module);
  ModuleInfo *IsInstrumentModule(char *module_name);
  void InstrumentAllLoadedModules();
  void InstrumentModule(ModuleInfo *module);
  void ClearInstrumentation(ModuleInfo *module);
  bool TryExecuteInstrumented(char * address);
  size_t GetTranslatedAddress(ModuleInfo *module, size_t address);
  void TranslateBasicBlock(char *address,
                           ModuleInfo *module,
                           std::set<char *> *queue,
                           std::list<std::pair<uint32_t, uint32_t>> *offset_fixes);
  void TranslateBasicBlockRecursive(char *address, ModuleInfo *module);
  void FixOffsetOrEnqueue(ModuleInfo *module,
                          uint32_t bb,
                          uint32_t jmp_offset,
                          std::set<char *> *queue,
                          std::list<std::pair<uint32_t, uint32_t>> *offset_fixes);
  void InvalidInstruction(ModuleInfo* module);

  // relative jump outside of current module
  void OutsideJump(ModuleInfo* module, size_t address);

  // needed to support cross-module linking
  // on module unloads / reloads
  void InvalidateCrossModuleLink(CrossModuleLink *link);
  void FixCrossModuleLink(CrossModuleLink *link);
  void FixCrossModuleLinks(ModuleInfo *module);
  void InvalidateCrossModuleLinks(ModuleInfo *module);
  void InvalidateCrossModuleLinks();
  void ClearCrossModuleLinks(ModuleInfo *module);
  void ClearCrossModuleLinks();

  void PatchModuleEntries(ModuleInfo* module);

  // functions related to indirect jump/call instrumentation
  void InitGlobalJumptable(ModuleInfo *module);
  void InstrumentIndirect(ModuleInfo *module,
                          Instruction& inst,
                          size_t instruction_address,
                          IndirectInstrumentation mode,
                          size_t bb_address);

  // returns the indirect instrumentation mode that should be used for a particular call
  // can be overriden
  virtual IndirectInstrumentation ShouldInstrumentIndirect(ModuleInfo *module,
                                                           Instruction& inst,
                                                           size_t instruction_address);

  size_t AddTranslatedJump(ModuleInfo *module,
                           ModuleInfo *target_module,
                           size_t original_target,
                           size_t actual_target,
                           size_t list_head_offset,
                           IndirectBreakpoinInfo& breakpoint_info,
                           bool global_indirect);
  bool HandleIndirectJMPBreakpoint(void *address);

  void PatchPointersLocal(char* buf, size_t size,
                          std::unordered_map<size_t, size_t>& search_replace,
                          bool commit_code, ModuleInfo* module);
  template<typename T>
  void PatchPointersLocalT(char* buf, size_t size,
                           std::unordered_map<size_t, size_t>& search_replace,
                           bool commit_code, ModuleInfo* module);

  IndirectInstrumentation indirect_instrumentation_mode;

  bool instrument_cross_module_calls;
  bool patch_return_addresses;
  bool persist_instrumentation_data;

  bool trace_basic_blocks;
  bool trace_module_entries;

  bool generate_unwind;

  bool page_extend_modules;

  // these could be indexed by module1 and module2 for performance
  // but the assumption for now is that there won't be too many of
  // them so a flat structure shoudl be ok for now
  std::list<CrossModuleLink> cross_module_links;

  bool instrumentation_disabled;
  bool instrument_modules_on_load;
  
  bool full_address_map;

  PatchModuleEntriesValue patch_module_entries;

  std::list<Hook *> hooks;
  std::unordered_map<uint64_t, Hook *> resolved_hooks;
  
  friend class Asssembler;
  friend class X86Assembler;
  friend class Arm64Assembler;
  friend class ModuleInfo;
  friend class UnwindGenerator;
  friend class Hook;
#if defined(_WIN64)
  friend class WinUnwindGenerator;
#elif __APPLE__
  friend class UnwindGeneratorMacOS;
#endif
};

struct IndirectBreakpoinInfo {
  size_t list_head;
  size_t source_bb;
#ifdef ARM64
  uint8_t branch_register;
#endif
};

class ModuleInfo {
 public:
  ModuleInfo();
  void ClearInstrumentation();

  std::string module_name;
  void *module_header;
  size_t min_address;
  size_t max_address;
  size_t code_size;
  bool loaded;
  bool instrumented;
  std::list<AddressRange> executable_ranges;

  size_t instrumented_code_size;
  size_t instrumented_code_allocated;
  char *instrumented_code_local;
  char *instrumented_code_remote;
  char *instrumented_code_remote_previous;

  std::unordered_map<uint32_t, uint32_t> basic_blocks;

  // instrumented address to original address
  std::map<size_t, size_t> address_map;

  size_t br_indirect_newtarget_global;

  // per callsite jumplist breakpoint
  // from breakpoint address to list head offset
  std::unordered_map<size_t, IndirectBreakpoinInfo> br_indirect_newtarget_list;

  size_t jumptable_offset;
  size_t jumptable_address_offset;

  std::unordered_set<size_t> invalid_instructions;
  std::unordered_map<size_t, size_t> outside_jumps;
  std::unordered_map<size_t, size_t> tracepoints;

  std::unordered_set<size_t> entry_offsets;

  UnwindData *unwind_data;
  
  bool do_protect;

  // clients can use this to store additional data
  // about the module
  void *client_data;
};

#endif // TINYINST_H