// Copyright 2017 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 <dlfcn.h>
#include <lib/fdio/directory.h>
#include <lib/fdio/fd.h>
#include <lib/fdio/fdio.h>
#include <link.h>
#include <stdio.h>
#include <string.h>
#include <sys/param.h>
#include <zircon/dlfcn.h>
#include <zircon/process.h>
#include <zircon/processargs.h>
#include <zircon/syscalls.h>

#include <zxtest/zxtest.h>

TEST(VdsoBaseTests, vdso_base_test) {
  char msg[128];

  struct link_map* lm = dlopen("libzircon.so", RTLD_NOLOAD);
  snprintf(msg, sizeof(msg), "dlopen(\"libzircon.so\") failed: %s", dlerror());
  EXPECT_NOT_NULL(lm, "%s", msg);
  uintptr_t rtld_vdso_base = lm->l_addr;
  int ok = dlclose(lm);
  snprintf(msg, sizeof(msg), "dlclose failed: %s", dlerror());
  EXPECT_EQ(ok, 0, "%s", msg);

  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);
  EXPECT_EQ(status, 0, "%s", msg);

  EXPECT_EQ(rtld_vdso_base, prop_vdso_base,
            "rtld reported address != process property reported address");
}

static int phdr_info_callback(struct dl_phdr_info* info, size_t size, void* data) {
  struct dl_phdr_info* key = data;
  if (info->dlpi_addr == key->dlpi_addr) {
    *key = *info;
    return 1;
  }
  return 0;
}

TEST(VdsoBaseTests, vdso_unmap_test) {
  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, "%s", msg);

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

  uintptr_t vdso_code_start = 0;
  size_t vdso_code_len = 0;
  for (uint_fast16_t 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?");

  // Removing the vDSO code mapping is not allowed.
  status = zx_vmar_unmap(zx_vmar_root_self(), vdso_code_start, vdso_code_len);
  EXPECT_EQ(status, ZX_ERR_ACCESS_DENIED, "unmap vDSO code");

  // Nor is removing a whole range overlapping the vDSO code.
  status = zx_vmar_unmap(zx_vmar_root_self(), vdso_code_start - PAGE_SIZE, PAGE_SIZE * 2);
  EXPECT_EQ(status, ZX_ERR_ACCESS_DENIED, "unmap range overlapping vDSO code");
}

TEST(VdsoBaseTests, vdso_map_test) {
  zx_handle_t vmo = zx_take_startup_handle(PA_HND(PA_VMO_VDSO, 0));
  ASSERT_NE(vmo, ZX_HANDLE_INVALID, "zx_take_startup_handle(PA_HND(PA_VMO_VDSO, 0))");

  // Since we already have a vDSO mapping, loading it again should fail.
  void* h = dlopen_vmo(vmo, RTLD_LOCAL);
  EXPECT_NULL(h, "dlopen_vmo on vDSO VMO succeeded");

  // Create a fresh process that doesn't already have a vDSO mapping.
  // We can't meaningfully test the other constraints on our own
  // process, because the "there can be only one" constraint trumps them.
  const char* name = "vdso_map_test";
  zx_handle_t proc, vmar;
  ASSERT_EQ(zx_process_create(zx_job_default(), name, strlen(name), 0, &proc, &vmar), ZX_OK,
            "zx_process_create failed");

  // This should fail because it's an executable mapping of
  // the wrong portion of the vDSO image (the first page is
  // rodata including the ELF headers).  Only the actual code
  // segment can be mapped executable.
  uintptr_t addr;
  zx_status_t status =
      zx_vmar_map(vmar, ZX_VM_PERM_READ | ZX_VM_PERM_EXECUTE, 0, vmo, 0, PAGE_SIZE, &addr);
  EXPECT_EQ(status, ZX_ERR_ACCESS_DENIED, "map vDSO data as executable");

  zx_handle_close(proc);
  zx_handle_close(vmar);
}