//===-- DynamicLoaderWindowsDYLD.cpp --------------------------------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//

#include "DynamicLoaderWindowsDYLD.h"

#include "MSVCRTCFrameRecognizer.h"
#include "Plugins/ObjectFile/PECOFF/ObjectFilePECOFF.h"
#include "lldb/Core/Module.h"
#include "lldb/Core/PluginManager.h"
#include "lldb/Symbol/ObjectFile.h"
#include "lldb/Target/ExecutionContext.h"
#include "lldb/Target/Platform.h"
#include "lldb/Target/Process.h"
#include "lldb/Target/RegisterContext.h"
#include "lldb/Target/Target.h"
#include "lldb/Target/ThreadPlanStepInstruction.h"
#include "lldb/Utility/LLDBLog.h"
#include "lldb/Utility/Log.h"

#include "llvm/TargetParser/Triple.h"

using namespace lldb;
using namespace lldb_private;

LLDB_PLUGIN_DEFINE(DynamicLoaderWindowsDYLD)

DynamicLoaderWindowsDYLD::DynamicLoaderWindowsDYLD(Process *process)
    : DynamicLoader(process) {}

DynamicLoaderWindowsDYLD::~DynamicLoaderWindowsDYLD() = default;

void DynamicLoaderWindowsDYLD::Initialize() {
  PluginManager::RegisterPlugin(GetPluginNameStatic(),
                                GetPluginDescriptionStatic(), CreateInstance);
}

void DynamicLoaderWindowsDYLD::Terminate() {
  PluginManager::UnregisterPlugin(CreateInstance);
}

llvm::StringRef DynamicLoaderWindowsDYLD::GetPluginDescriptionStatic() {
  return "Dynamic loader plug-in that watches for shared library "
         "loads/unloads in Windows processes.";
}

DynamicLoader *DynamicLoaderWindowsDYLD::CreateInstance(Process *process,
                                                        bool force) {
  bool should_create = force;
  if (!should_create) {
    const llvm::Triple &triple_ref =
        process->GetTarget().GetArchitecture().GetTriple();
    if (triple_ref.getOS() == llvm::Triple::Win32)
      should_create = true;
  }

  if (should_create)
    return new DynamicLoaderWindowsDYLD(process);

  return nullptr;
}

void DynamicLoaderWindowsDYLD::OnLoadModule(lldb::ModuleSP module_sp,
                                            const ModuleSpec module_spec,
                                            lldb::addr_t module_addr) {
  // Resolve the module unless we already have one.
  if (!module_sp) {
    Status error;
    module_sp = m_process->GetTarget().GetOrCreateModule(
        module_spec, true /* notify */, &error);
    if (error.Fail())
      return;
  }

  m_loaded_modules.insert({module_addr, lldb::ModuleWP(module_sp)});
  UpdateLoadedSectionsCommon(module_sp, module_addr, false);
  ModuleList module_list;
  module_list.Append(module_sp);
  m_process->GetTarget().ModulesDidLoad(module_list);
}

void DynamicLoaderWindowsDYLD::OnUnloadModule(lldb::addr_t module_addr) {
  auto it = m_loaded_modules.find(module_addr);
  if (it == m_loaded_modules.end())
    return;

  ModuleSP module_sp = it->second.lock();
  m_loaded_modules.erase(it);

  if (!module_sp)
    return;

  bool full_unload = true;
  for (const auto &entry : m_loaded_modules) {
    ModuleSP other_sp = entry.second.lock();
    if (other_sp == module_sp) {
      UpdateLoadedSectionsCommon(module_sp, entry.first, false);
      full_unload = false;
    }
  }

  if (full_unload) {
    UnloadSectionsCommon(module_sp);
    ModuleList module_list;
    module_list.Append(module_sp);
    m_process->GetTarget().ModulesDidUnload(module_list, false);
  }
}

lldb::addr_t DynamicLoaderWindowsDYLD::GetLoadAddress(ModuleSP executable) {
  // First, see if the load address is already cached.
  for (const auto &entry : m_loaded_modules) {
    ModuleSP mod = entry.second.lock();
    if (mod == executable && entry.first != LLDB_INVALID_ADDRESS)
      return entry.first;
  }

  lldb::addr_t load_addr = LLDB_INVALID_ADDRESS;

  // Second, try to get it through the process plugins.  For a remote process,
  // the remote platform will be responsible for providing it.
  FileSpec file_spec(executable->GetPlatformFileSpec());
  bool is_loaded = false;
  Status status =
      m_process->GetFileLoadAddress(file_spec, is_loaded, load_addr);
  // Servers other than lldb server could respond with a bogus address.
  if (status.Success() && is_loaded && load_addr != LLDB_INVALID_ADDRESS) {
    m_loaded_modules.insert({load_addr, lldb::ModuleWP(executable)});
    return load_addr;
  }

  return LLDB_INVALID_ADDRESS;
}

void DynamicLoaderWindowsDYLD::DidAttach() {
  Log *log = GetLog(LLDBLog::DynamicLoader);
  LLDB_LOGF(log, "DynamicLoaderWindowsDYLD::%s()", __FUNCTION__);

  RegisterMSVCRTCFrameRecognizer(m_process->GetTarget());

  ModuleSP executable = GetTargetExecutable();

  if (!executable)
    return;

  // Try to fetch the load address of the file from the process, since there
  // could be randomization of the load address.
  lldb::addr_t load_addr = GetLoadAddress(executable);
  if (load_addr == LLDB_INVALID_ADDRESS) {
    // The process does not know the executable under the target's path: it
    // was started through another path to the same file, or runs a different
    // one. Take the modules it reports as loaded, which replaces the target's
    // executable with the process's.
    auto error = m_process->LoadModules();
    if (!error)
      return;
    LLDB_LOG_ERROR(log, std::move(error), "failed to load modules: {0}");

    // If the load address couldn't be determined (e.g. for a remote gdbstub
    // target), fall back to the preferred load address/image base from the
    // file if the process is live.
    // This is a best-effort guess, if the image isn't loaded there the
    // sections will have incorrect addresses and the slide will have to be
    // specified manually.
    if (!m_process->IsLiveDebugSession())
      return;

    ObjectFile *object_file = executable->GetObjectFile();
    if (!object_file)
      return;

    load_addr = object_file->GetBaseAddress().GetFileAddress();
    if (load_addr == LLDB_INVALID_ADDRESS)
      return;

    executable->ReportWarning(
        "failed to get load address, assuming image base at {0:x}", load_addr);
  } else {
    // Request the process base address.
    lldb::addr_t image_base = m_process->GetImageInfoAddress();
    if (image_base == load_addr)
      return;
  }

  // Rebase the process's modules if there is a mismatch.
  UpdateLoadedSections(executable, LLDB_INVALID_ADDRESS, load_addr, false);

  ModuleList module_list;
  module_list.Append(executable);
  m_process->GetTarget().ModulesDidLoad(module_list);
  auto error = m_process->LoadModules();
  LLDB_LOG_ERROR(log, std::move(error), "failed to load modules: {0}");
}

void DynamicLoaderWindowsDYLD::DidLaunch() {
  Log *log = GetLog(LLDBLog::DynamicLoader);
  LLDB_LOGF(log, "DynamicLoaderWindowsDYLD::%s()", __FUNCTION__);

  RegisterMSVCRTCFrameRecognizer(m_process->GetTarget());

  ModuleSP executable = GetTargetExecutable();
  if (!executable)
    return;

  lldb::addr_t load_addr = GetLoadAddress(executable);
  if (load_addr != LLDB_INVALID_ADDRESS) {
    // Update the loaded sections so that the breakpoints can be resolved.
    UpdateLoadedSections(executable, LLDB_INVALID_ADDRESS, load_addr, false);

    ModuleList module_list;
    module_list.Append(executable);
    m_process->GetTarget().ModulesDidLoad(module_list);
    auto error = m_process->LoadModules();
    LLDB_LOG_ERROR(log, std::move(error), "failed to load modules: {0}");
  }
}

Status DynamicLoaderWindowsDYLD::CanLoadImage() { return Status(); }

ThreadPlanSP
DynamicLoaderWindowsDYLD::GetStepThroughTrampolinePlan(Thread &thread,
                                                       bool stop) {
  auto arch = m_process->GetTarget().GetArchitecture();
  if (arch.GetMachine() != llvm::Triple::x86) {
    return ThreadPlanSP();
  }

  RegisterContextSP reg_ctx_sp = thread.GetRegisterContext();
  if (!reg_ctx_sp)
    return ThreadPlanSP();

  uint64_t pc = reg_ctx_sp->GetPC();
  // Max size of an instruction in x86 is 15 bytes.
  AddressRange range(pc, 2 * 15);

  DisassemblerSP disassembler_sp = Disassembler::DisassembleRange(
      arch, nullptr, nullptr, nullptr, nullptr, m_process->GetTarget(), range);
  if (!disassembler_sp) {
    return ThreadPlanSP();
  }

  InstructionList &insn_list = disassembler_sp->GetInstructionList();

  // First instruction in a x86 Windows trampoline is going to be an indirect
  // jump through the IAT and the next one will be a nop (usually there for
  // alignment purposes). e.g.:
  //     0x70ff4cfc <+956>: jmpl   *0x7100c2a8
  //     0x70ff4d02 <+962>: nop

  auto first_insn = insn_list.GetInstructionAtIndex(0);
  auto second_insn = insn_list.GetInstructionAtIndex(1);

  ExecutionContext exe_ctx(m_process->GetTarget());
  if (first_insn == nullptr || second_insn == nullptr)
    return ThreadPlanSP();

  const char *first_mnemonic = first_insn->GetMnemonic(&exe_ctx);
  const char *second_mnemonic = second_insn->GetMnemonic(&exe_ctx);
  if (!first_mnemonic || !second_mnemonic ||
      strcmp(first_mnemonic, "jmpl") != 0 ||
      strcmp(second_mnemonic, "nop") != 0)
    return ThreadPlanSP();

  assert(first_insn->DoesBranch() && !second_insn->DoesBranch());

  return ThreadPlanSP(new ThreadPlanStepInstruction(
      thread, false, false, eVoteNoOpinion, eVoteNoOpinion));
}

lldb::addr_t
DynamicLoaderWindowsDYLD::GetThreadLocalData(const lldb::ModuleSP module,
                                             lldb::ThreadSP thread,
                                             lldb::addr_t tls_file_addr) {
  Log *log = GetLog(LLDBLog::DynamicLoader);
  ObjectFile *obj = module->GetObjectFile();
  if (!obj) {
    LLDB_LOG(log, "no object file for module");
    return LLDB_INVALID_ADDRESS;
  }
  std::string object_name;
  if (log)
    object_name = obj->GetObjectName();

  StructuredData::ObjectSP extended_info = thread->GetExtendedInfo();
  if (!extended_info) {
    LLDB_LOG(log, "missing extended info of thread");
    return LLDB_INVALID_ADDRESS;
  }
  StructuredData::Dictionary *dict = extended_info->GetAsDictionary();
  if (!dict) {
    LLDB_LOG(log, "extended info of thread is not a dictionary");
    return LLDB_INVALID_ADDRESS;
  }
  addr_t teb_base = 0;
  if (!dict->GetValueForKeyAsInteger("teb_address", teb_base) || !teb_base ||
      teb_base == LLDB_INVALID_ADDRESS) {
    LLDB_LOG(log, "missing 'teb_address' in extended thread info");
    return LLDB_INVALID_ADDRESS;
  }

  ProcessSP process = thread->GetProcess();
  if (!process) {
    LLDB_LOG(log, "missing process");
    return LLDB_INVALID_ADDRESS;
  }
  uint32_t addr_byte_size = process->GetAddressByteSize();
  addr_t tls_array_addr = teb_base + 11 * addr_byte_size;
  llvm::Expected<addr_t> tls_array =
      process->ReadPointerFromMemory(tls_array_addr);
  if (!tls_array || *tls_array == LLDB_INVALID_ADDRESS) {
    LLDB_LOG_ERROR(log, tls_array.takeError(),
                   "failed to read TLS array from {1:x16}: {0}",
                   tls_array_addr);
    return LLDB_INVALID_ADDRESS;
  }

  Address tls_index_addr;
  if (!GetTlsIndexAddressForModule(module, tls_index_addr)) {
    LLDB_LOG(log, "failed to resolve tls index for {0}", object_name);
    return LLDB_INVALID_ADDRESS;
  }

  Target &target = process->GetTarget();
  addr_t tls_index_load_addr = tls_index_addr.GetLoadAddress(&target);
  if (tls_index_load_addr == LLDB_INVALID_ADDRESS) {
    LLDB_LOG(log, "failed to resolve load address of '_tls_index' in {0}",
             object_name);
    return LLDB_INVALID_ADDRESS;
  }

  Status error;
  uint64_t tls_index =
      process->ReadUnsignedIntegerFromMemory(tls_index_load_addr, 4, 0, error);
  if (error.Fail()) {
    LLDB_LOG_ERROR(log, error.takeError(),
                   "failed to read '_tls_index' in {1} from {2:x8}: {0}",
                   object_name, tls_index_load_addr);
    return LLDB_INVALID_ADDRESS;
  }

  addr_t tls_index_item = *tls_array + tls_index * addr_byte_size;
  LLDB_LOG(log, "tls base address for {0} is at {1:x8} + {2:x8} * {3} = {4:x8}",
           object_name, *tls_array, tls_index, addr_byte_size, tls_index_item);

  llvm::Expected<addr_t> tls_base =
      process->ReadPointerFromMemory(tls_index_item);
  if (!tls_base || *tls_base == LLDB_INVALID_ADDRESS || *tls_base == 0) {
    LLDB_LOG_ERROR(log, tls_base.takeError(),
                   "failed to read tls array in {1} at index {2} ({3:x8}): {0}",
                   object_name, tls_index, tls_index_item);
    return LLDB_INVALID_ADDRESS;
  }

  return *tls_base + tls_file_addr;
}

bool DynamicLoaderWindowsDYLD::GetTlsIndexAddressForModule(
    const ModuleSP &module, Address &addr) const {
  ObjectFilePECOFF *obj =
      llvm::dyn_cast_or_null<ObjectFilePECOFF>(module->GetObjectFile());
  if (obj) {
    addr_t tls_index_file_addr = obj->GetTlsIndexAddress();
    if (tls_index_file_addr != LLDB_INVALID_ADDRESS &&
        module->ResolveFileAddress(tls_index_file_addr, addr))
      return true;
  }

  // If we can't get the tls index from the object file, try the _tls_index
  // symbol.
  Symtab *symtab = module->GetSymtab();
  if (!symtab)
    return false;

  Symbol *tls_index_sym = symtab->FindFirstSymbolWithNameAndType(
      ConstString("_tls_index"), eSymbolTypeAny, Symtab::eDebugAny,
      Symtab::eVisibilityAny);
  if (!tls_index_sym)
    return false;

  addr = tls_index_sym->GetAddress();
  return addr.IsValid();
}