#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 <unistd.h>

#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)));

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 struct rseq *g_registered_rseq;
static struct rseq_compat *g_registered_rseq_compat;
static bool g_own_registration;

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 ENXIO: return "ENXIO";
	case EPERM: return "EPERM";
	case EBUSY: return "EBUSY";
	default: return "UNKNOWN";
	}
}

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)
{
	unsigned long aux_feature_size = getauxval(AT_RSEQ_FEATURE_SIZE);
	unsigned long aux_align = getauxval(AT_RSEQ_ALIGN);

	g_rseq_feature_size = aux_feature_size ? static_cast<unsigned int>(aux_feature_size) : 20U;
	g_rseq_align = aux_align ? static_cast<unsigned int>(aux_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);
		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 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));
	unsigned int cpu = 0;
	unsigned int node = 0;

	std::printf("AT_RSEQ_FEATURE_SIZE: %u\n", g_rseq_feature_size);
	std::printf("AT_RSEQ_ALIGN: %u\n", g_rseq_align);
	std::printf("registered rseq size: %u\n", g_rseq_alloc_size);
	print_registration_source();
	std::printf("registered rseq addr: %p\n", static_cast<void *>(g_registered_rseq));
	std::printf("struct rseq has slice_ctrl field available: %s\n",
		    has_slice_ctrl ? "yes" : "no");

	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",
			    (g_registered_rseq_compat->flags & RSEQ_CS_FLAG_SLICE_EXT_AVAILABLE) ? "set" : "clear");
		std::printf("slice ext enabled bit: %s\n",
			    (g_registered_rseq_compat->flags & RSEQ_CS_FLAG_SLICE_EXT_ENABLED) ? "set" : "clear");
	}

	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);
	}

	errno = 0;
	int rc = prctl(PR_RSEQ_SLICE_EXTENSION, PR_RSEQ_SLICE_EXTENSION_GET, 0UL, 0UL, 0UL);
	int saved_errno = errno;

	if (rc >= 0) {
		std::printf("prctl(GET) result=%d\n", rc);
		std::printf("slice extension runtime availability: yes\n");
		std::printf("slice extension enabled for this thread: %s\n",
			    (rc & PR_RSEQ_SLICE_EXT_ENABLE) ? "yes" : "no");
		return;
	}

	std::printf("prctl(GET) failed: errno=%d (%s)\n",
		    saved_errno, errno_name(saved_errno));
	std::printf("slice extension runtime availability: %s\n",
		    saved_errno == EINVAL ? "no" : "unknown");
}

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

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