// Copyright 2016 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include <elf.h>
#include <inttypes.h>
#include <link.h>
#include <stdarg.h>
#include <stdbool.h>
#include <stdlib.h>
#include <string.h>

#include <fbl/algorithm.h>
#include <test-utils/test-utils.h>
#include <unittest/unittest.h>
#include <zircon/compiler.h>
#include <zircon/process.h>
#include <zircon/processargs.h>
#include <zircon/syscalls.h>
#include <zircon/syscalls/debug.h>
#include <zircon/syscalls/object.h>
#include <zircon/syscalls/port.h>

#include "utils.h"

// argv[0]
const char* g_program_path;

uint64_t extract_pc_reg(const zx_thread_state_general_regs_t* regs) {
#if defined(__x86_64__)
  return regs->rip;
#elif defined(__aarch64__)
  return regs->pc;
#endif
}

uint64_t extract_sp_reg(const zx_thread_state_general_regs_t* regs) {
#if defined(__x86_64__)
  return regs->rsp;
#elif defined(__aarch64__)
  return regs->sp;
#endif
}

static __NO_RETURN void fatal(const char* msg, ...) {
  fprintf(stderr, "%s: encountered fatal error:\n", g_program_path);
  va_list args;
  va_start(args, msg);
  vfprintf(stderr, msg, args);
  va_end(args);
  exit(1);
}

uint32_t get_uint32(char* buf) {
  uint32_t value = 0;
  memcpy(&value, buf, sizeof(value));
  return value;
}

uint64_t get_uint64(char* buf) {
  uint64_t value = 0;
  memcpy(&value, buf, sizeof(value));
  return value;
}

void set_uint64(char* buf, uint64_t value) { memcpy(buf, &value, sizeof(value)); }

uint32_t get_uint32_property(zx_handle_t handle, uint32_t prop) {
  uint32_t value;
  zx_status_t status = zx_object_get_property(handle, prop, &value, sizeof(value));
  if (status != ZX_OK)
    tu_fatal("zx_object_get_property failed", status);
  return value;
}

void send_request(zx_handle_t handle, const request_message_t& rqst) {
  unittest_printf("sending request %d on handle %u\n", rqst.type, handle);
  tu_channel_write(handle, 0, &rqst, sizeof(rqst), NULL, 0);
}

void send_simple_request(zx_handle_t handle, request_t type) {
  unittest_printf("sending request %d on handle %u\n", type, handle);
  tu_channel_write(handle, 0, &type, sizeof(type), NULL, 0);
}

void send_response(zx_handle_t handle, const response_message_t& resp) {
  unittest_printf("sending response %d on handle %u\n", resp.type, handle);
  tu_channel_write(handle, 0, &resp, sizeof(resp), NULL, 0);
}

void send_response_with_handle(zx_handle_t handle, const response_message_t& resp,
                               zx_handle_t resp_handle) {
  unittest_printf("sending response %d on handle %u\n", resp.type, handle);
  tu_channel_write(handle, 0, &resp, sizeof(resp), &resp_handle, 1);
}

void send_simple_response(zx_handle_t handle, response_t type) {
  unittest_printf("sending response %d on handle %u\n", type, handle);
  tu_channel_write(handle, 0, &type, sizeof(type), NULL, 0);
}

// This returns "bool" because it uses ASSERT_*.

bool recv_request(zx_handle_t handle, request_message_t* rqst) {
  BEGIN_HELPER;

  unittest_printf("waiting for request on handle %u\n", handle);

  ASSERT_TRUE(tu_channel_wait_readable(handle), "peer closed while trying to read message");

  uint32_t num_bytes = sizeof(*rqst);
  tu_channel_read(handle, 0, rqst, &num_bytes, NULL, NULL);
  ASSERT_LE(num_bytes, sizeof(*rqst), "unexpected request size");

  END_HELPER;
}

// This returns "bool" because it uses ASSERT_*.

bool recv_response(zx_handle_t handle, response_message_t* resp) {
  BEGIN_HELPER;

  unittest_printf("waiting for response on handle %u\n", handle);

  ASSERT_TRUE(tu_channel_wait_readable(handle), "peer closed while trying to read message");

  uint32_t num_bytes = sizeof(*resp);
  zx_handle_t resp_handle = ZX_HANDLE_INVALID;
  uint32_t num_handles = 1;
  tu_channel_read(handle, 0, resp, &num_bytes, &resp_handle, &num_handles);
  ASSERT_LE(num_bytes, sizeof(*resp), "unexpected response size");
  if (num_handles > 0) {
    EXPECT_EQ(num_handles, 1, "");
    EXPECT_NE(resp_handle, ZX_HANDLE_INVALID, "");
    resp->handle = resp_handle;
  }

  END_HELPER;
}

// This returns "bool" because it uses ASSERT_*.

bool recv_simple_response(zx_handle_t handle, response_t expected_type) {
  BEGIN_HELPER;

  response_message_t response;
  ASSERT_TRUE(recv_response(handle, &response), "");
  unittest_printf("received message %d\n", response.type);
  EXPECT_EQ(response.type, expected_type, "");

  END_HELPER;
}

bool verify_inferior_running(zx_handle_t channel) {
  BEGIN_HELPER;

  send_simple_request(channel, RQST_PING);
  EXPECT_TRUE(recv_simple_response(channel, RESP_PONG), "");

  END_HELPER;
}

bool get_inferior_thread_handle(zx_handle_t channel, zx_handle_t* thread) {
  BEGIN_HELPER;

  send_simple_request(channel, RQST_GET_THREAD_HANDLE);
  response_message_t response;
  ASSERT_TRUE(recv_response(channel, &response), "");
  ASSERT_EQ(response.type, RESP_THREAD_HANDLE, "");
  ASSERT_NE(response.handle, ZX_HANDLE_INVALID, "");
  *thread = response.handle;

  END_HELPER;
}

static int phdr_info_callback(dl_phdr_info* info, size_t size, void* argp) {
  dl_phdr_info* arg = reinterpret_cast<dl_phdr_info*>(argp);
  if (info->dlpi_addr == arg->dlpi_addr) {
    *arg = *info;
    return 1;
  }
  return 0;
}

// Fetch the [inclusive] range of the executable segment of the vdso.

bool get_vdso_exec_range(uintptr_t* start, uintptr_t* end) {
  BEGIN_HELPER;

  char msg[128];

  uintptr_t prop_vdso_base;
  zx_status_t status = zx_object_get_property(zx_process_self(), ZX_PROP_PROCESS_VDSO_BASE_ADDRESS,
                                              &prop_vdso_base, sizeof(prop_vdso_base));
  snprintf(msg, sizeof(msg), "zx_object_get_property failed: %d", status);
  ASSERT_EQ(status, 0, msg);

  dl_phdr_info info;
  info.dlpi_addr = prop_vdso_base;
  int ret = dl_iterate_phdr(&phdr_info_callback, &info);
  ASSERT_EQ(ret, 1, "dl_iterate_phdr didn't see vDSO?");

  uintptr_t vdso_code_start = 0;
  size_t vdso_code_len = 0;
  for (unsigned i = 0; i < info.dlpi_phnum; ++i) {
    if (info.dlpi_phdr[i].p_type == PT_LOAD && (info.dlpi_phdr[i].p_flags & PF_X)) {
      vdso_code_start = info.dlpi_addr + info.dlpi_phdr[i].p_vaddr;
      vdso_code_len = info.dlpi_phdr[i].p_memsz;
      break;
    }
  }
  ASSERT_NE(vdso_code_start, 0u, "vDSO has no code segment?");
  ASSERT_NE(vdso_code_len, 0u, "vDSO has no code segment?");

  *start = vdso_code_start;
  *end = vdso_code_start + vdso_code_len - 1;

  END_HELPER;
}

struct find_so_callback_arg_t {
  const char* so_name;
  dl_phdr_info* info;
};

static int find_so_callback(dl_phdr_info* info, size_t size, void* argp) {
  auto arg = reinterpret_cast<find_so_callback_arg_t*>(argp);
  if (strcmp(info->dlpi_name, arg->so_name) == 0) {
    *arg->info = *info;
    return 1;
  }
  return 0;
}

void find_so(const char* so_name, dl_phdr_info* info) {
  find_so_callback_arg_t callback_arg;
  callback_arg.so_name = so_name;
  callback_arg.info = info;

  int ret = dl_iterate_phdr(&find_so_callback, &callback_arg);
  if (ret != 1) {
    fatal("dl_iterate_phdr didn't find SO?\n");
  }

  uint64_t base_addr = info->dlpi_addr;
  auto ehdr = reinterpret_cast<const Elf64_Ehdr*>(base_addr);
  if (ehdr->e_ident[0] != 0x7f || ehdr->e_ident[1] != 'E' || ehdr->e_ident[2] != 'L' ||
      ehdr->e_ident[3] != 'F') {
    fatal("Unexpected executable ELF header contents\n");
  }
}

zx_vaddr_t get_exec_load_addr() {
  dl_phdr_info info;
  // The executable doesn't have an SO name.
  find_so("", &info);
  return info.dlpi_addr;
}

zx_vaddr_t get_libc_load_addr() {
  dl_phdr_info info;
  find_so("libc.so", &info);
  return info.dlpi_addr;
}

bool get_inferior_load_addrs(zx_handle_t channel, zx_vaddr_t* libc_load_addr,
                             zx_vaddr_t* exec_load_addr) {
  BEGIN_HELPER;

  send_simple_request(channel, RQST_GET_LOAD_ADDRS);
  response_message_t response;
  ASSERT_TRUE(recv_response(channel, &response), "");
  ASSERT_EQ(response.type, RESP_LOAD_ADDRS, "");
  *libc_load_addr = response.payload.load_addrs.libc_load_addr;
  *exec_load_addr = response.payload.load_addrs.exec_load_addr;

  END_HELPER;
}

zx_vaddr_t get_libc_entry_point() {
  dl_phdr_info info;
  find_so("libc.so", &info);

  uint64_t base_addr = info.dlpi_addr;
  auto ehdr = reinterpret_cast<const Elf64_Ehdr*>(base_addr);
  return ehdr->e_entry;
}