#include <cerrno>
#include <cinttypes>
#include <cstddef>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>

#include <linux/auxvec.h>
#include <linux/rseq.h>
#include <sys/auxv.h>
#include <sys/prctl.h>
#include <sys/syscall.h>
#include <sys/utsname.h>
#include <unistd.h>

#ifndef ENOTSUPP
#define ENOTSUPP 524
#endif

#if defined(__x86_64__)
#ifndef ARCH_GET_FS
#define ARCH_GET_FS 0x1003
#endif
#endif

#ifndef PR_RSEQ_SLICE_EXTENSION
#define PR_RSEQ_SLICE_EXTENSION 79
#define PR_RSEQ_SLICE_EXTENSION_GET 1
#define PR_RSEQ_SLICE_EXTENSION_SET 2
#define PR_RSEQ_SLICE_EXT_ENABLE 0x01
#endif

#ifndef RSEQ_CS_FLAG_SLICE_EXT_AVAILABLE
#define RSEQ_CS_FLAG_SLICE_EXT_AVAILABLE (1U << 4)
#define RSEQ_CS_FLAG_SLICE_EXT_ENABLED (1U << 5)
#endif

#ifndef RSEQ_SIG
#if defined(__x86_64__) || defined(__i386__)
#define RSEQ_SIG 0x53053053
#elif defined(__aarch64__)
#define RSEQ_SIG 0x00bc28d4
#else
#error "Add RSEQ_SIG for this architecture"
#endif
#endif

#ifndef SYS_rseq
#ifdef __NR_rseq
#define SYS_rseq __NR_rseq
#endif
#endif

#ifndef SYS_getcpu
#ifdef __NR_getcpu
#define SYS_getcpu __NR_getcpu
#endif
#endif

extern "C" {
extern __attribute__((weak)) ptrdiff_t __rseq_offset;
extern __attribute__((weak)) unsigned int __rseq_size;
extern __attribute__((weak)) unsigned int __rseq_flags;
}

struct rseq_slice_ctrl_fields {
	uint8_t request;
	uint8_t granted;
	uint16_t reserved;
};

struct rseq_slice_ctrl_compat {
	union {
		uint32_t all;
		rseq_slice_ctrl_fields parts;
	};
};

struct rseq_compat {
	uint32_t cpu_id_start;
	int32_t cpu_id;
	uint64_t rseq_cs;
	uint32_t flags;
	uint32_t node_id;
	uint32_t mm_cid;
	struct rseq_slice_ctrl_compat slice_ctrl;
	uint8_t reserved;
} __attribute__((aligned(32)));

struct prctl_probe_result {
	bool ok;
	int value;
	int err;
};

alignas(32) static thread_local unsigned char local_rseq_storage[512];

static unsigned int g_rseq_feature_size;
static unsigned int g_rseq_alloc_size;
static unsigned int g_rseq_align;
static unsigned long g_aux_rseq_feature_size;
static unsigned long g_aux_rseq_align;
static struct rseq *g_registered_rseq;
static struct rseq_compat *g_registered_rseq_compat;
static bool g_own_registration;

static const char *yes_no(bool value)
{
	return value ? "yes" : "no";
}

static const char *set_clear(bool value)
{
	return value ? "set" : "clear";
}

static int get_thread_pointer(uintptr_t *tp_out)
{
#if defined(__x86_64__)
	unsigned long fsbase = 0;

	if (syscall(SYS_arch_prctl, ARCH_GET_FS, &fsbase) != 0)
		return -1;
	*tp_out = fsbase;
	return 0;
#elif defined(__aarch64__)
	void *tp = nullptr;

	__asm__ volatile("mrs %0, tpidr_el0" : "=r"(tp));
	*tp_out = reinterpret_cast<uintptr_t>(tp);
	return 0;
#else
	(void) tp_out;
	errno = ENOTSUP;
	return -1;
#endif
}

static int sys_rseq(struct rseq *rseq, uint32_t len, int flags, uint32_t sig)
{
#ifdef SYS_rseq
	return static_cast<int>(syscall(SYS_rseq, rseq, len, flags, sig));
#else
	(void) rseq;
	(void) len;
	(void) flags;
	(void) sig;
	errno = ENOSYS;
	return -1;
#endif
}

static int sys_getcpu(unsigned int *cpu, unsigned int *node)
{
#ifdef SYS_getcpu
	return static_cast<int>(syscall(SYS_getcpu, cpu, node, nullptr));
#else
	(void) cpu;
	(void) node;
	errno = ENOSYS;
	return -1;
#endif
}

static const char *errno_name(int err)
{
	switch (err) {
	case 0: return "0";
	case EINVAL: return "EINVAL";
	case ENOSYS: return "ENOSYS";
	case ENOTSUP: return "ENOTSUP";
	case ENOTSUPP: return "ENOTSUPP";
	case ENXIO: return "ENXIO";
	case EPERM: return "EPERM";
	case EBUSY: return "EBUSY";
	default: return "UNKNOWN";
	}
}

static const char *prctl_failure_meaning(int err)
{
	switch (err) {
	case EINVAL:
		return "the prctl operation or argument is not accepted by this kernel";
	case ENOTSUPP:
		return "the prctl operation exists, but the slice extension is not supported here";
	case EPERM:
		return "the kernel denied the requested operation";
	default:
		return "the kernel returned an unclassified failure";
	}
}

static void print_errno_status(const char *label, int err)
{
	std::printf("%s: errno=%d (%s: %s)\n",
		    label, err, errno_name(err), std::strerror(err));
}

static unsigned int max_u32(unsigned int a, unsigned int b)
{
	return a > b ? a : b;
}

static bool feature_present(size_t end_offset)
{
	return g_rseq_feature_size >= end_offset;
}

static void print_registration_source(void)
{
	if (g_own_registration)
		std::printf("rseq registration: local syscall registration\n");
	else
		std::printf("rseq registration: existing libc-owned registration\n");
}

static int setup_rseq(void)
{
	g_aux_rseq_feature_size = getauxval(AT_RSEQ_FEATURE_SIZE);
	g_aux_rseq_align = getauxval(AT_RSEQ_ALIGN);

	g_rseq_feature_size = g_aux_rseq_feature_size ?
		static_cast<unsigned int>(g_aux_rseq_feature_size) : 20U;
	g_rseq_align = g_aux_rseq_align ?
		static_cast<unsigned int>(g_aux_rseq_align) : 32U;
	g_rseq_alloc_size = max_u32(g_rseq_feature_size, 32U);

	if (&__rseq_size != nullptr && __rseq_size != 0) {
		uintptr_t tp = 0;

		if (get_thread_pointer(&tp) != 0) {
			std::perror("get_thread_pointer");
			return -1;
		}

		g_registered_rseq = reinterpret_cast<struct rseq *>(tp + __rseq_offset);
		g_registered_rseq_compat = reinterpret_cast<struct rseq_compat *>(tp + __rseq_offset);
		g_own_registration = false;
		if (__rseq_size < g_rseq_feature_size)
			g_rseq_feature_size = __rseq_size;
		if (__rseq_size > g_rseq_alloc_size)
			g_rseq_alloc_size = __rseq_size;
		return 0;
	}

	if (g_rseq_alloc_size > sizeof(local_rseq_storage)) {
		std::fprintf(stderr,
			     "local rseq area too small: need %u bytes, have %zu\n",
			     g_rseq_alloc_size, sizeof(local_rseq_storage));
		errno = EOVERFLOW;
		return -1;
	}
	if ((reinterpret_cast<uintptr_t>(local_rseq_storage) % g_rseq_align) != 0) {
		std::fprintf(stderr, "local rseq area alignment mismatch: need %u\n",
			     g_rseq_align);
		errno = EINVAL;
		return -1;
	}

	auto *local_rseq = reinterpret_cast<struct rseq *>(local_rseq_storage);
	auto *local_rseq_compat = reinterpret_cast<struct rseq_compat *>(local_rseq_storage);
	std::memset(local_rseq_storage, 0, sizeof(local_rseq_storage));
	local_rseq_compat->cpu_id = RSEQ_CPU_ID_UNINITIALIZED;

	if (sys_rseq(local_rseq, g_rseq_alloc_size, 0, RSEQ_SIG) != 0) {
		std::perror("rseq register");
		return -1;
	}

	g_registered_rseq = local_rseq;
	g_registered_rseq_compat = local_rseq_compat;
	g_own_registration = true;
	return 0;
}

static void teardown_rseq(void)
{
	if (!g_own_registration)
		return;
	auto *local_rseq = reinterpret_cast<struct rseq *>(local_rseq_storage);
	if (sys_rseq(local_rseq, g_rseq_alloc_size, RSEQ_FLAG_UNREGISTER, RSEQ_SIG) != 0)
		std::perror("rseq unregister");
}

static prctl_probe_result probe_prctl_get(void)
{
	errno = 0;
	int rc = prctl(PR_RSEQ_SLICE_EXTENSION, PR_RSEQ_SLICE_EXTENSION_GET,
		       0UL, 0UL, 0UL);
	int saved_errno = errno;

	if (rc >= 0)
		return { true, rc, 0 };
	return { false, -1, saved_errno };
}

static prctl_probe_result probe_prctl_set(unsigned long value)
{
	errno = 0;
	int rc = prctl(PR_RSEQ_SLICE_EXTENSION, PR_RSEQ_SLICE_EXTENSION_SET,
		       value, 0UL, 0UL);
	int saved_errno = errno;

	if (rc >= 0)
		return { true, rc, 0 };
	return { false, -1, saved_errno };
}

static void print_prctl_result(const char *label, const prctl_probe_result &result)
{
	if (result.ok) {
		std::printf("%s: ok, value=%d\n", label, result.value);
		return;
	}
	print_errno_status(label, result.err);
	std::printf("%s meaning: %s\n", label, prctl_failure_meaning(result.err));
}

static void print_kernel_version(void)
{
	struct utsname uts;

	if (uname(&uts) != 0) {
		std::perror("uname");
		return;
	}
	std::printf("kernel: %s %s %s %s\n",
		    uts.sysname, uts.release, uts.version, uts.machine);
}

static void print_slice_status_summary(bool has_slice_ctrl, bool flags_available,
				       const prctl_probe_result &get_result)
{
	bool flag_available = flags_available &&
		(g_registered_rseq_compat->flags & RSEQ_CS_FLAG_SLICE_EXT_AVAILABLE);
	bool flag_enabled = flags_available &&
		(g_registered_rseq_compat->flags & RSEQ_CS_FLAG_SLICE_EXT_ENABLED);
	bool prctl_enabled = get_result.ok &&
		(get_result.value & PR_RSEQ_SLICE_EXT_ENABLE);

	std::printf("status: rseq syscall registered: yes\n");
	std::printf("status: rseq extensible feature area: %s\n",
		    yes_no(g_rseq_feature_size > 20U));
	std::printf("status: rseq slice_ctrl field present: %s\n",
		    yes_no(has_slice_ctrl));
	std::printf("status: rseq slice extension available flag: %s\n",
		    flags_available ? set_clear(flag_available) : "unavailable");
	std::printf("status: rseq slice extension enabled flag: %s\n",
		    flags_available ? set_clear(flag_enabled) : "unavailable");
	std::printf("status: PR_RSEQ_SLICE_EXTENSION GET usable: %s\n",
		    yes_no(get_result.ok));

	if (get_result.ok) {
		std::printf("status: PR_RSEQ_SLICE_EXTENSION enabled: %s\n",
			    yes_no(prctl_enabled));
		std::printf("status: rseq slice extension availability: %s\n",
			    prctl_enabled || flag_available ? "available" : "available but disabled");
		return;
	}
	if (get_result.err == ENOTSUPP) {
		std::printf("status: rseq slice extension availability: not supported by this kernel/arch/config\n");
		return;
	}
	if (get_result.err == EINVAL) {
		std::printf("status: rseq slice extension availability: no accepted prctl API on this kernel\n");
		return;
	}
	std::printf("status: rseq slice extension availability: unknown\n");
}

static void probe_slice_extension(void)
{
	bool has_slice_ctrl = feature_present(offsetof(struct rseq_compat, slice_ctrl) +
					      sizeof(g_registered_rseq_compat->slice_ctrl));
	bool flags_available = feature_present(offsetof(struct rseq_compat, flags) +
					       sizeof(g_registered_rseq_compat->flags));
	prctl_probe_result prctl_get = probe_prctl_get();
	unsigned int cpu = 0;
	unsigned int node = 0;

	print_kernel_version();
	std::printf("AT_RSEQ_FEATURE_SIZE raw: %lu\n", g_aux_rseq_feature_size);
	std::printf("AT_RSEQ_ALIGN raw: %lu\n", g_aux_rseq_align);
	std::printf("effective rseq feature size: %u\n", g_rseq_feature_size);
	std::printf("effective rseq alignment: %u\n", g_rseq_align);
	std::printf("registered rseq size: %u\n", g_rseq_alloc_size);
	print_registration_source();
	if (&__rseq_size != nullptr) {
		std::printf("libc __rseq_size=%u __rseq_offset=%td __rseq_flags=0x%x\n",
			    __rseq_size, __rseq_offset, __rseq_flags);
	}
	std::printf("registered rseq addr: %p\n", static_cast<void *>(g_registered_rseq));
	std::printf("struct rseq has slice_ctrl field available: %s\n",
		    yes_no(has_slice_ctrl));

	if (sys_getcpu(&cpu, &node) == 0)
		std::printf("getcpu(): cpu=%u node=%u\n", cpu, node);

	std::printf("rseq cpu_id_start=%u cpu_id=%d\n",
		    g_registered_rseq->cpu_id_start,
		    static_cast<int32_t>(g_registered_rseq->cpu_id));

	if (feature_present(offsetof(struct rseq_compat, node_id) +
			    sizeof(g_registered_rseq_compat->node_id))) {
		std::printf("rseq node_id=%u\n", g_registered_rseq_compat->node_id);
	}
	if (feature_present(offsetof(struct rseq_compat, mm_cid) +
			    sizeof(g_registered_rseq_compat->mm_cid))) {
		std::printf("rseq mm_cid=%u\n", g_registered_rseq_compat->mm_cid);
	}

	if (flags_available) {
		std::printf("rseq flags=0x%x\n", g_registered_rseq_compat->flags);
		std::printf("slice ext available bit: %s\n",
			    set_clear(g_registered_rseq_compat->flags & RSEQ_CS_FLAG_SLICE_EXT_AVAILABLE));
		std::printf("slice ext enabled bit: %s\n",
			    set_clear(g_registered_rseq_compat->flags & RSEQ_CS_FLAG_SLICE_EXT_ENABLED));
	}

	if (has_slice_ctrl) {
		std::printf("slice_ctrl.request=%u granted=%u raw=0x%x\n",
			    g_registered_rseq_compat->slice_ctrl.parts.request,
			    g_registered_rseq_compat->slice_ctrl.parts.granted,
			    g_registered_rseq_compat->slice_ctrl.all);
	}

	print_prctl_result("prctl(PR_RSEQ_SLICE_EXTENSION, GET)", prctl_get);
	if (prctl_get.ok) {
		prctl_probe_result set_enable = probe_prctl_set(PR_RSEQ_SLICE_EXT_ENABLE);
		print_prctl_result("prctl(PR_RSEQ_SLICE_EXTENSION, SET enable)",
				   set_enable);
		prctl_probe_result after_enable = probe_prctl_get();
		print_prctl_result("prctl(PR_RSEQ_SLICE_EXTENSION, GET after enable)",
				   after_enable);
		prctl_probe_result set_disable = probe_prctl_set(0UL);
		print_prctl_result("prctl(PR_RSEQ_SLICE_EXTENSION, SET disable)",
				   set_disable);
	}
	print_slice_status_summary(has_slice_ctrl, flags_available, prctl_get);
}

int main()
{
	if (setup_rseq() != 0)
		return 1;

	probe_slice_extension();
	teardown_rseq();
	return 0;
}