// Copyright 2016 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 <string>

#include "absl/strings/ascii.h"

#include "absl/strings/escaping.h"

#include "absl/strings/str_cat.h"

#include "absl/strings/string_view.h"

#include "absl/strings/substitute.h"

#include "bloaty.h"

#include "capstone/capstone.h"

#include "re.h"

#include "util.h"

using absl::string_view;

namespace bloaty {

namespace {

static std::string RightPad(const std::string& input, size_t size) {

  std::string ret = input;

  while (ret.size() < size) {

    ret += " ";

  }

  return ret;

}

}  // anonymous namespace

void DisassembleFindReferences(const DisassemblyInfo& info, RangeSink* sink) {

  if (info.arch != CS_ARCH_X86) {

    // x86 only for now.

    return;

  }

  csh handle;

  if (cs_open(info.arch, info.mode, &handle) != CS_ERR_OK ||

      cs_option(handle, CS_OPT_DETAIL, CS_OPT_ON) != CS_ERR_OK) {

    THROW("Couldn't initialize Capstone");

  }

  if (info.text.size() == 0) {

    cs_close(&handle);

    THROW("Tried to disassemble empty function.");

  }

  cs_insn *in = cs_malloc(handle);

  uint64_t address = info.start_address;

  const uint8_t* ptr = reinterpret_cast<const uint8_t*>(info.text.data());

  size_t size = info.text.size();

  while (size > 0) {

    if (!cs_disasm_iter(handle, &ptr, &size, &address, in)) {

      // Some symbols that end up in the .text section aren't really functions

      // but data.  Not sure why this happens.

      if (verbose_level > 1) {

        printf("Error disassembling function at address: %" PRIx64 "\n",

               address);

      }

      goto cleanup;

    }

    size_t count = in->detail->x86.op_count;

    for (size_t i = 0; i < count; i++) {

      cs_x86_op* op = &in->detail->x86.operands[i];

      if (op->type == X86_OP_MEM && op->mem.base == X86_REG_RIP &&

          op->mem.segment == X86_REG_INVALID &&

          op->mem.index == X86_REG_INVALID) {

        uint64_t to_address = in->address + in->size + op->mem.disp;

        if (to_address) {

          sink->AddVMRangeForVMAddr("x86_disassemble", in->address, to_address,

                                    RangeSink::kUnknownSize);

        }

      }

    }

  }

cleanup:

  cs_free(in, 1);

  cs_close(&handle);

}

bool TryGetJumpTarget(cs_arch arch, cs_insn *in, uint64_t* target) {

  switch (arch) {

    case CS_ARCH_X86:

      switch (in->id) {

        case X86_INS_JAE:

        case X86_INS_JA:

        case X86_INS_JBE:

        case X86_INS_JB:

        case X86_INS_JCXZ:

        case X86_INS_JECXZ:

        case X86_INS_JE:

        case X86_INS_JGE:

        case X86_INS_JG:

        case X86_INS_JLE:

        case X86_INS_JL:

        case X86_INS_JMP:

        case X86_INS_JNE:

        case X86_INS_JNO:

        case X86_INS_JNP:

        case X86_INS_JNS:

        case X86_INS_JO:

        case X86_INS_JP:

        case X86_INS_JS:

        case X86_INS_CALL: {

          auto op0 = in->detail->x86.operands[0];

          if (op0.type == X86_OP_IMM) {

            *target = op0.imm;

            return true;

          }

          return false;

        }

        default:

          return false;

      }

    default:

      return false;

  }

}

std::string DisassembleFunction(const DisassemblyInfo& info) {

  std::string ret;

  csh handle;

  if (cs_open(info.arch, info.mode, &handle) != CS_ERR_OK ||

      cs_option(handle, CS_OPT_DETAIL, CS_OPT_ON) != CS_ERR_OK) {

    THROW("Couldn't initialize Capstone");

  }

  if (info.text.size() == 0) {

    THROW("Tried to disassemble empty function.");

  }

  cs_insn *insn;

  size_t count =

      cs_disasm(handle, reinterpret_cast<const uint8_t *>(info.text.data()),

                info.text.size(), info.start_address, 0, &insn);

  if (count == 0) {

    THROW("Error disassembling function.");

  }

  std::map<uint64_t, int> local_labels;

  for (size_t i = 0; i < count; i++) {

    cs_insn *in = insn + i;

    uint64_t target;

    if (TryGetJumpTarget(info.arch, in, &target) &&

        target >= info.start_address &&

        target < info.start_address + info.text.size()) {

      local_labels[target] = 0;  // Fill in real value later.

    }

  }

  int label = 0;

  for (auto& pair : local_labels) {

    pair.second = label++;

  }

  for (size_t i = 0; i < count; i++) {

    cs_insn *in = insn + i;

    std::string bytes = absl::BytesToHexString(

        string_view(reinterpret_cast<const char*>(in->bytes), in->size));

    string_view mnemonic(in->mnemonic);

    std::string op_str(in->op_str);

    std::string match;

    std::string label;

    if (info.arch == CS_ARCH_X86) {

      if (in->id == X86_INS_LEA) {

        ReImpl::GlobalReplace(&op_str, "\\w?word ptr ", "");

      } else if (in->id == X86_INS_NOP) {

        op_str.clear();

      } else {

        // qword ptr => QWORD

        while (ReImpl::PartialMatch(op_str, "(\\w?word) ptr", &match)) {

          std::string upper_match = match;

          absl::AsciiStrToUpper(&upper_match);

          ReImpl::Replace(&op_str, match + " ptr", upper_match);

        }

      }

    }

    ReImpl::GlobalReplace(&op_str, " ", "");

    auto iter = local_labels.find(in->address);

    if (iter != local_labels.end()) {

      label = std::to_string(iter->second) + ":";

    }

    uint64_t target;

    if (TryGetJumpTarget(info.arch, in, &target)) {

      auto iter = local_labels.find(target);

      std::string label;

      if (iter != local_labels.end()) {

        if (target > in->address) {

          op_str = ">" + std::to_string(iter->second);

        } else {

          op_str = "<" + std::to_string(iter->second);

        }

      } else if (info.symbol_map.vm_map.TryGetLabel(target, &label)) {

        op_str = label;

      }

    }

    absl::StrAppend(&ret, " ", RightPad(label, 4),

                    RightPad(std::string(mnemonic), 8), " ", op_str, "\n");

  }

  cs_close(&handle);

  return ret;

}

}  // namespace bloaty