//===-- X86AsmInstrumentation.cpp - Instrument X86 inline assembly C++ -*-===//
//
//                     The LLVM Compiler Infrastructure
//
// This file is distributed under the University of Illinois Open Source
// License. See LICENSE.TXT for details.
//
//===----------------------------------------------------------------------===//

#include "X86AsmInstrumentation.h"
#include "MCTargetDesc/X86BaseInfo.h"
#include "X86Operand.h"
#include "llvm/ADT/StringExtras.h"
#include "llvm/ADT/Triple.h"
#include "llvm/MC/MCAsmInfo.h"
#include "llvm/MC/MCContext.h"
#include "llvm/MC/MCInst.h"
#include "llvm/MC/MCInstBuilder.h"
#include "llvm/MC/MCInstrInfo.h"
#include "llvm/MC/MCParser/MCParsedAsmOperand.h"
#include "llvm/MC/MCParser/MCTargetAsmParser.h"
#include "llvm/MC/MCStreamer.h"
#include "llvm/MC/MCSubtargetInfo.h"
#include "llvm/MC/MCTargetOptions.h"
#include <algorithm>
#include <cassert>
#include <vector>

//#include <iostream>

// Following comment describes how assembly instrumentation works.
// Currently we have only AddressSanitizer instrumentation, but we're
// planning to implement MemorySanitizer for inline assembly too. If
// you're not familiar with AddressSanitizer algorithm, please, read
// https://code.google.com/p/address-sanitizer/wiki/AddressSanitizerAlgorithm.
//
// When inline assembly is parsed by an instance of X86AsmParser, all
// instructions are emitted via EmitInstruction method. That's the
// place where X86AsmInstrumentation analyzes an instruction and
// decides, whether the instruction should be emitted as is or
// instrumentation is required. The latter case happens when an
// instruction reads from or writes to memory. Now instruction opcode
// is explicitly checked, and if an instruction has a memory operand
// (for instance, movq (%rsi, %rcx, 8), %rax) - it should be
// instrumented.  There're also exist instructions that modify
// memory but don't have an explicit memory operands, for instance,
// movs.
//
// Let's consider at first 8-byte memory accesses when an instruction
// has an explicit memory operand. In this case we need two registers -
// AddressReg to compute address of a memory cells which are accessed
// and ShadowReg to compute corresponding shadow address. So, we need
// to spill both registers before instrumentation code and restore them
// after instrumentation. Thus, in general, instrumentation code will
// look like this:
// PUSHF  # Store flags, otherwise they will be overwritten
// PUSH AddressReg  # spill AddressReg
// PUSH ShadowReg   # spill ShadowReg
// LEA MemOp, AddressReg  # compute address of the memory operand
// MOV AddressReg, ShadowReg
// SHR ShadowReg, 3
// # ShadowOffset(AddressReg >> 3) contains address of a shadow
// # corresponding to MemOp.
// CMP ShadowOffset(ShadowReg), 0  # test shadow value
// JZ .Done  # when shadow equals to zero, everything is fine
// MOV AddressReg, RDI
// # Call __asan_report function with AddressReg as an argument
// CALL __asan_report
// .Done:
// POP ShadowReg  # Restore ShadowReg
// POP AddressReg  # Restore AddressReg
// POPF  # Restore flags
//
// Memory accesses with different size (1-, 2-, 4- and 16-byte) are
// handled in a similar manner, but small memory accesses (less than 8
// byte) require an additional ScratchReg, which is used for shadow value.
//
// If, suppose, we're instrumenting an instruction like movs, only
// contents of RDI, RDI + AccessSize * RCX, RSI, RSI + AccessSize *
// RCX are checked.  In this case there're no need to spill and restore
// AddressReg , ShadowReg or flags four times, they're saved on stack
// just once, before instrumentation of these four addresses, and restored
// at the end of the instrumentation.
//
// There exist several things which complicate this simple algorithm.
// * Instrumented memory operand can have RSP as a base or an index
//   register.  So we need to add a constant offset before computation
//   of memory address, since flags, AddressReg, ShadowReg, etc. were
//   already stored on stack and RSP was modified.
// * Debug info (usually, DWARF) should be adjusted, because sometimes
//   RSP is used as a frame register. So, we need to select some
//   register as a frame register and temprorary override current CFA
//   register.

namespace llvm_ks {

X86AsmInstrumentation::X86AsmInstrumentation(const MCSubtargetInfo *&STI)
    : STI(STI), InitialFrameReg(0) {}

X86AsmInstrumentation::~X86AsmInstrumentation() {}

void X86AsmInstrumentation::InstrumentAndEmitInstruction(
    MCInst &Inst, OperandVector &Operands, MCContext &Ctx,
    const MCInstrInfo &MII, MCStreamer &Out, unsigned int &KsError) {
  EmitInstruction(Out, Inst, KsError);
}

void X86AsmInstrumentation::EmitInstruction(MCStreamer &Out,
                                            MCInst &Inst,
                                            unsigned int &KsError) {
  Out.EmitInstruction(Inst, *STI, KsError);
}

unsigned X86AsmInstrumentation::GetFrameRegGeneric(const MCContext &Ctx,
                                                   MCStreamer &Out) {
  if (!Out.getNumFrameInfos()) // No active dwarf frame
    return X86::NoRegister;
  const MCDwarfFrameInfo &Frame = Out.getDwarfFrameInfos().back();
  if (Frame.End) // Active dwarf frame is closed
    return X86::NoRegister;
  const MCRegisterInfo *MRI = Ctx.getRegisterInfo();
  if (!MRI) // No register info
    return X86::NoRegister;

  if (InitialFrameReg) {
    // FrameReg is set explicitly, we're instrumenting a MachineFunction.
    return InitialFrameReg;
  }

  return MRI->getLLVMRegNum(Frame.CurrentCfaRegister, true /* IsEH */);
}

X86AsmInstrumentation *
CreateX86AsmInstrumentation(const MCTargetOptions &MCOptions,
                            const MCContext &Ctx, const MCSubtargetInfo *&STI) {
  return new X86AsmInstrumentation(STI);
}

} // end llvm namespace