// Copyright 2019 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "base/profiler/suspendable_thread_delegate_mac.h"

#include <mach/mach.h>
#include <mach/thread_act.h>
#include <pthread.h>

#include <vector>

#include "base/apple/mach_logging.h"
#include "base/check.h"
#include "base/profiler/profile_builder.h"
#include "base/profiler/register_context_registers.h"
#include "build/build_config.h"

// IMPORTANT NOTE: Some functions within this implementation are invoked while
// the target thread is suspended so it must not do any allocation from the
// heap, including indirectly via use of DCHECK/CHECK or other logging
// statements. Otherwise this code can deadlock on heap locks acquired by the
// target thread before it was suspended. These functions are commented with "NO
// HEAP ALLOCATIONS".

namespace base {

namespace {

#if defined(ARCH_CPU_X86_64)
constexpr mach_msg_type_number_t kThreadStateCount = x86_THREAD_STATE64_COUNT;
constexpr thread_state_flavor_t kThreadStateFlavor = x86_THREAD_STATE64;
#elif defined(ARCH_CPU_ARM64)
constexpr mach_msg_type_number_t kThreadStateCount = ARM_THREAD_STATE64_COUNT;
constexpr thread_state_flavor_t kThreadStateFlavor = ARM_THREAD_STATE64;
#endif

// Fills |state| with |target_thread|'s context. NO HEAP ALLOCATIONS.
bool GetThreadContextImpl(thread_act_t target_thread, RegisterContext* state) {
  auto count = kThreadStateCount;
  return thread_get_state(target_thread, kThreadStateFlavor,
                          reinterpret_cast<thread_state_t>(state),
                          &count) == KERN_SUCCESS;
}

}  // namespace

// ScopedSuspendThread --------------------------------------------------------

// NO HEAP ALLOCATIONS after thread_suspend.
SuspendableThreadDelegateMac::ScopedSuspendThread::ScopedSuspendThread(
    mach_port_t thread_port)
    : thread_port_(thread_suspend(thread_port) == KERN_SUCCESS
                       ? thread_port
                       : MACH_PORT_NULL) {}

// NO HEAP ALLOCATIONS. The MACH_CHECK is OK because it provides a more noisy
// failure mode than deadlocking.
SuspendableThreadDelegateMac::ScopedSuspendThread::~ScopedSuspendThread() {
  if (!WasSuccessful()) {
    return;
  }

  kern_return_t kr = thread_resume(thread_port_);
  MACH_CHECK(kr == KERN_SUCCESS, kr) << "thread_resume";
}

bool SuspendableThreadDelegateMac::ScopedSuspendThread::WasSuccessful() const {
  return thread_port_ != MACH_PORT_NULL;
}

// SuspendableThreadDelegateMac -----------------------------------------------

SuspendableThreadDelegateMac::SuspendableThreadDelegateMac(
    SamplingProfilerThreadToken thread_token)
    : thread_id_(thread_token.id),
      thread_port_(pthread_mach_thread_np(thread_token.pthread_id)),
      thread_stack_base_address_(reinterpret_cast<uintptr_t>(
          pthread_get_stackaddr_np(thread_token.pthread_id))) {
  // This class suspends threads, and those threads might be suspended in dyld.
  // Therefore, for all the system functions that might be linked in dynamically
  // that are used while threads are suspended, make calls to them to make sure
  // that they are linked up.
  RegisterContext thread_context;
  GetThreadContextImpl(thread_port_, &thread_context);
}

SuspendableThreadDelegateMac::~SuspendableThreadDelegateMac() = default;

std::unique_ptr<SuspendableThreadDelegate::ScopedSuspendThread>
SuspendableThreadDelegateMac::CreateScopedSuspendThread() {
  return std::make_unique<ScopedSuspendThread>(thread_port_);
}

PlatformThreadId SuspendableThreadDelegateMac::GetThreadId() const {
  return thread_id_;
}

// NO HEAP ALLOCATIONS.
bool SuspendableThreadDelegateMac::GetThreadContext(
    RegisterContext* thread_context) {
  return GetThreadContextImpl(thread_port_, thread_context);
}

// NO HEAP ALLOCATIONS.
uintptr_t SuspendableThreadDelegateMac::GetStackBaseAddress() const {
  return thread_stack_base_address_;
}

// NO HEAP ALLOCATIONS.
bool SuspendableThreadDelegateMac::CanCopyStack(uintptr_t stack_pointer) {
  return true;
}

std::vector<uintptr_t> SuspendableThreadDelegateMac::GetRegisters(
    RegisterContext* thread_context) {
#if defined(ARCH_CPU_X86_64)
  return {AsUintPtr(&thread_context->__rbx), AsUintPtr(&thread_context->__rbp),
          AsUintPtr(&thread_context->__rsp), AsUintPtr(&thread_context->__r12),
          AsUintPtr(&thread_context->__r13), AsUintPtr(&thread_context->__r14),
          AsUintPtr(&thread_context->__r15)};
#elif defined(ARCH_CPU_ARM64)  // defined(ARCH_CPU_X86_64)
  return {
      RegisterContextFramePointer(thread_context),
      RegisterContextStackPointer(thread_context),
      thread_context->__x[19],
      thread_context->__x[20],
      thread_context->__x[21],
      thread_context->__x[22],
      thread_context->__x[23],
      thread_context->__x[24],
      thread_context->__x[25],
      thread_context->__x[26],
      thread_context->__x[27],
      thread_context->__x[28],
  };
#endif                         // defined(ARCH_CPU_ARM64)
}

void SuspendableThreadDelegateMac::SetRegisters(
    RegisterContext* thread_context,
    const std::vector<uintptr_t>& registers) {
#if defined(ARCH_CPU_X86_64)
  CHECK_EQ(registers.size(), 7u);
  AsUintPtr(&thread_context->__rbx) = registers[0];
  AsUintPtr(&thread_context->__rbp) = registers[1];
  AsUintPtr(&thread_context->__rsp) = registers[2];
  AsUintPtr(&thread_context->__r12) = registers[3];
  AsUintPtr(&thread_context->__r13) = registers[4];
  AsUintPtr(&thread_context->__r14) = registers[5];
  AsUintPtr(&thread_context->__r15) = registers[6];
#elif defined(ARCH_CPU_ARM64)  // defined(ARCH_CPU_X86_64)
  CHECK_EQ(registers.size(), 12u);
  SetRegisterContextFramePointer(thread_context, registers[0]);
  SetRegisterContextStackPointer(thread_context, registers[1]);
  thread_context->__x[19] = registers[2];
  thread_context->__x[20] = registers[3];
  thread_context->__x[21] = registers[4];
  thread_context->__x[22] = registers[5];
  thread_context->__x[23] = registers[6];
  thread_context->__x[24] = registers[7];
  thread_context->__x[25] = registers[8];
  thread_context->__x[26] = registers[9];
  thread_context->__x[27] = registers[10];
  thread_context->__x[28] = registers[11];
#endif                         // defined(ARCH_CPU_ARM64)
}

}  // namespace base