// SPDX-License-Identifier: GPL-2.0
#include "util/jitdump.h"

#include <errno.h>
#include <fcntl.h>
#include <stddef.h>
#include <stdint.h>
#include <stdio.h>
#include <string.h>
#include <time.h>

#include <sys/mman.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <unistd.h>

#include "../tests.h"

#ifndef HAVE_GETTID
#include <sys/syscall.h>
static inline pid_t gettid(void)
{
	return (pid_t)syscall(SYS_gettid);
}
#endif

#define CHK_BYTE 0x5a

static inline uint64_t get_timestamp(void)
{
#if defined(__x86_64__) || defined(__i386__)
	unsigned int low, high;

	asm volatile("rdtsc" : "=a"(low), "=d"(high));

	return low | ((uint64_t)high) << 32;
#else
	struct timespec ts;
	int ret;

	ret = clock_gettime(CLOCK_MONOTONIC, &ts);
	if (ret)
		return 0;

	return ((uint64_t)ts.tv_sec * 1000000000) + ts.tv_nsec;
#endif
}

static FILE *open_jitdump(void)
{
	struct jitheader header = {
		.magic = JITHEADER_MAGIC,
		.version = JITHEADER_VERSION,
		.total_size = sizeof(header),
		.pid = getpid(),
		.timestamp = get_timestamp(),
		.flags =
#if defined(__x86_64__) || defined(__i386__)
			JITDUMP_FLAGS_ARCH_TIMESTAMP,
#else
			0,
#endif
	};
	char filename[256];
	int fd;
	FILE *f;
	void *m;

	snprintf(filename, sizeof(filename), "jit-%d.dump", getpid());
	/* Securely open using O_CREAT | O_EXCL to prevent symlink attacks. */
	fd = open(filename, O_CREAT | O_EXCL | O_RDWR, 0644);
	if (fd < 0) {
		pr_err("Failed to open jitdump '%s': %s\n", filename, strerror(errno));
		return NULL;
	}
	f = fdopen(fd, "w+");
	if (!f) {
		pr_err("Failed to associate stream with fd for '%s'\n", filename);
		close(fd);
		unlink(filename);
		return NULL;
	}
	/* Create an MMAP event for the jitdump file. That is how perf tool finds it. */
	m = mmap(0, getpagesize(), PROT_READ | PROT_EXEC, MAP_PRIVATE, fileno(f), 0);
	if (m == MAP_FAILED) {
		pr_err("mmap failed: %s\n", strerror(errno));
		fclose(f);
		unlink(filename);
		return NULL;
	}
	munmap(m, getpagesize());

	if (fwrite(&header, sizeof(header), 1, f) != 1) {
		pr_err("Error writing jitdump header\n");
		fclose(f);
		unlink(filename);
		return NULL;
	}
	return f;
}

static int write_jitdump(FILE *f, void *addr, const void *dat, size_t sz, uint64_t *idx)
{
	const char *sym = "jit_workload";
	size_t sym_len = strlen(sym) + 1;
	struct jr_code_load rec = {
		.p.id = JIT_CODE_LOAD,
		.p.total_size = sizeof(rec) + sym_len + sz,
		.p.timestamp = get_timestamp(),
		.pid = getpid(),
		.tid = gettid(),
		.vma = (unsigned long)addr,
		.code_addr = (unsigned long)addr,
		.code_size = sz,
		.code_index = ++*idx,
	};

	if (fwrite(&rec, sizeof(rec), 1, f) != 1 ||
	    fwrite(sym, sym_len, 1, f) != 1 ||
	    fwrite(dat, sz, 1, f) != 1)
		return -1;
	return 0;
}

static void close_jitdump(FILE *f)
{
	fclose(f);
}

static int jitdump(int argc __maybe_unused, const char **argv __maybe_unused)
{
#if defined(__x86_64__) || defined(__i386__)
	/* Code to execute: mov CHK_BYTE, %eax ; ret */
	uint8_t dat[] = { 0xb8, CHK_BYTE, 0x00, 0x00, 0x00, 0xc3 };
#elif defined(__aarch64__)
	/* Code to execute: mov w0, #CHK_BYTE ; ret */
	uint8_t dat[] = {
		(CHK_BYTE << 5) & 0xff, (CHK_BYTE >> 3) & 0xff, 0x80, 0x52,
		0xc0, 0x03, 0x5f, 0xd6
	};
#elif defined(__riscv)
	/* Code to execute: li a0, CHK_BYTE ; ret */
	uint8_t dat[] = {
		0x13, 0x05, (CHK_BYTE << 4) & 0xff, (CHK_BYTE >> 4) & 0xff,
		0x67, 0x80, 0x00, 0x00
	};
#elif defined(__powerpc__)
	/* Code to execute: li r3, CHK_BYTE ; blr */
	uint32_t dat[] = { 0x38600000 | (CHK_BYTE & 0xffff), 0x4e800020 };
#elif defined(__s390x__)
	/* Code to execute: lhi %r2, CHK_BYTE ; br %r14 */
	uint8_t dat[] = { 0xa7, 0x28, (CHK_BYTE >> 8) & 0xff, CHK_BYTE & 0xff, 0x07, 0xfe };
#elif defined(__arm__)
	/* Code to execute: mov r0, #CHK_BYTE ; bx lr */
	uint8_t dat[] = {
		CHK_BYTE & 0xff, 0x00, 0xa0, 0xe3,
		0x1e, 0xff, 0x2f, 0xe1
	};
#elif defined(__mips__)
	/* Code to execute: addiu $v0, $zero, CHK_BYTE ; jr $ra ; nop */
	uint32_t dat[] = { 0x24020000 | (CHK_BYTE & 0xffff), 0x03e00008, 0x00000000 };
#elif defined(__loongarch__)
	/* Code to execute: addi.w $a0, $zero, CHK_BYTE ; jirl $zero, $ra, 0 */
	uint32_t dat[] = { 0x02800004 | ((CHK_BYTE & 0xfff) << 10), 0x4c000020 };
#else
	uint32_t dat[0];
#endif
	void *addr;
	FILE *f;
	uint64_t idx = 0;
	int ret = 1;

	/* Reachable fallback check for unsupported architectures right at start. */
	if (sizeof(dat) == 0) {
		pr_err("JITDUMP workload not supported on this architecture\n");
		return 1;
	}

	/* Get a memory page to store executable code. */
	addr = mmap(0, getpagesize(), PROT_READ | PROT_WRITE | PROT_EXEC,
		    MAP_ANONYMOUS | MAP_PRIVATE, -1, 0);
	if (addr == MAP_FAILED) {
		pr_err("Failed to map 1 -rwx page\n");
		return 1;
	}

	f = open_jitdump();
	if (!f) {
		pr_err("Failed to open JITDUMP\n");
		munmap(addr, getpagesize());
		return 1;
	}
	/* Copy executable code to executable memory page. */
	memcpy(addr, dat, sizeof(dat));
	/* Synchronize the Instruction and Data caches. */
	__builtin___clear_cache(addr, (char *)addr + sizeof(dat));

	/* Record it in the jitdump file */
	if (write_jitdump(f, addr, dat, sizeof(dat), &idx) == 0) {
		int (*fn)(void) = addr;

		/* Call the function. */
		ret = fn() - CHK_BYTE;
	}
	close_jitdump(f);
	munmap(addr, getpagesize());
	return ret;
}

DEFINE_WORKLOAD(jitdump);