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

#include <errno.h>
#include <inttypes.h>
#include <string.h>

#include <dlfcn.h>
#include <elf.h>
#include <fcntl.h>
#include <linux/ctype.h>

#include <capstone/capstone.h>

#include "addr_location.h"
#include "annotate.h"
#include "debug.h"
#include "disasm.h"
#include "dso.h"
#include "machine.h"
#include "map.h"
#include "namespaces.h"
#include "print_insn.h"
#include "symbol.h"
#include "thread.h"

#if CS_VERSION_MAJOR < 5
#define CS_ARCH_RISCV    15
#define CS_MODE_RISCV32  1
#define CS_MODE_RISCV64  2
#define CS_MODE_RISCVC   4
#endif

#ifdef LIBCAPSTONE_DLOPEN
static void *perf_cs_dll_handle(void)
{
	static bool dll_handle_init;
	static void *dll_handle;

	if (!dll_handle_init) {
		dll_handle_init = true;
		dll_handle = dlopen("libcapstone.so", RTLD_LAZY);
		if (!dll_handle)
			pr_debug("dlopen failed for libcapstone.so\n");
	}
	return dll_handle;
}
#endif

static enum cs_err perf_cs_open(enum cs_arch arch, enum cs_mode mode, csh *handle)
{
#ifndef LIBCAPSTONE_DLOPEN
	return cs_open(arch, mode, handle);
#else
	static bool fn_init;
	static enum cs_err (*fn)(enum cs_arch arch, enum cs_mode mode, csh *handle);

	if (!fn_init) {
		fn = dlsym(perf_cs_dll_handle(), "cs_open");
		if (!fn)
			pr_debug("dlsym failed for cs_open\n");
		fn_init = true;
	}
	if (!fn)
		return CS_ERR_HANDLE;
	return fn(arch, mode, handle);
#endif
}

static enum cs_err perf_cs_option(csh handle, enum cs_opt_type type, size_t value)
{
#ifndef LIBCAPSTONE_DLOPEN
	return cs_option(handle, type, value);
#else
	static bool fn_init;
	static enum cs_err (*fn)(csh handle, enum cs_opt_type type, size_t value);

	if (!fn_init) {
		fn = dlsym(perf_cs_dll_handle(), "cs_option");
		if (!fn)
			pr_debug("dlsym failed for cs_option\n");
		fn_init = true;
	}
	if (!fn)
		return CS_ERR_HANDLE;
	return fn(handle, type, value);
#endif
}

static size_t perf_cs_disasm(csh handle, const uint8_t *code, size_t code_size,
			uint64_t address, size_t count, struct cs_insn **insn)
{
#ifndef LIBCAPSTONE_DLOPEN
	return cs_disasm(handle, code, code_size, address, count, insn);
#else
	static bool fn_init;
	static enum cs_err (*fn)(csh handle, const uint8_t *code, size_t code_size,
				 uint64_t address, size_t count, struct cs_insn **insn);

	if (!fn_init) {
		fn = dlsym(perf_cs_dll_handle(), "cs_disasm");
		if (!fn)
			pr_debug("dlsym failed for cs_disasm\n");
		fn_init = true;
	}
	if (!fn)
		return CS_ERR_HANDLE;
	return fn(handle, code, code_size, address, count, insn);
#endif
}

static void perf_cs_free(struct cs_insn *insn, size_t count)
{
#ifndef LIBCAPSTONE_DLOPEN
	cs_free(insn, count);
#else
	static bool fn_init;
	static void (*fn)(struct cs_insn *insn, size_t count);

	if (!fn_init) {
		fn = dlsym(perf_cs_dll_handle(), "cs_free");
		if (!fn)
			pr_debug("dlsym failed for cs_free\n");
		fn_init = true;
	}
	if (!fn)
		return;
	fn(insn, count);
#endif
}

static enum cs_err perf_cs_close(csh *handle)
{
#ifndef LIBCAPSTONE_DLOPEN
	return cs_close(handle);
#else
	static bool fn_init;
	static enum cs_err (*fn)(csh *handle);

	if (!fn_init) {
		fn = dlsym(perf_cs_dll_handle(), "cs_close");
		if (!fn)
			pr_debug("dlsym failed for cs_close\n");
		fn_init = true;
	}
	if (!fn)
		return CS_ERR_HANDLE;
	return fn(handle);
#endif
}

static bool e_machine_to_capstone(uint16_t e_machine, bool is64, bool is_big_endian,
				  enum cs_arch *arch, enum cs_mode *mode)
{
	*mode = is_big_endian ? CS_MODE_BIG_ENDIAN : CS_MODE_LITTLE_ENDIAN;

	switch (e_machine) {
	case EM_X86_64:
	case EM_386:
		*arch = CS_ARCH_X86;
		*mode |= is64 ? CS_MODE_64 : CS_MODE_32;
		return true;
	case EM_AARCH64:
		*arch = CS_ARCH_ARM64;
		*mode |= CS_MODE_ARM;
		return true;
	case EM_ARM:
		*arch = CS_ARCH_ARM;
		*mode |= CS_MODE_ARM | CS_MODE_V8;
		return true;
	case EM_S390:
		*arch = CS_ARCH_SYSZ;
		return true;
	case EM_MIPS:
		*arch = CS_ARCH_MIPS;
		*mode |= is64 ? CS_MODE_MIPS64 : CS_MODE_MIPS32;
		return true;
	case EM_PPC:
		*arch = CS_ARCH_PPC;
		return true;
	case EM_PPC64:
		*arch = CS_ARCH_PPC;
		*mode |= CS_MODE_64;
		return true;
	case EM_SPARC:
		*arch = CS_ARCH_SPARC;
		return true;
	case EM_SPARCV9:
		*arch = CS_ARCH_SPARC;
		*mode |= CS_MODE_V9;
		return true;
	case EM_RISCV:
		*arch = CS_ARCH_RISCV;
		*mode |= (is64 ? CS_MODE_RISCV64 : CS_MODE_RISCV32) | CS_MODE_RISCVC;
		return true;
	default:
		return false;
	}
}

static int capstone_init(uint16_t e_machine, csh *cs_handle, bool is64, bool is_big_endian,
			 bool disassembler_style)
{
	enum cs_arch arch;
	enum cs_mode mode;

	if (!e_machine_to_capstone(e_machine, is64, is_big_endian, &arch, &mode))
		return -1;

	if (perf_cs_open(arch, mode, cs_handle) != CS_ERR_OK) {
		pr_warning_once("cs_open failed\n");
		return -1;
	}

	if (arch == CS_ARCH_X86) {
		/*
		 * In case of using capstone_init while symbol__disassemble
		 * setting CS_OPT_SYNTAX_ATT depends if disassembler_style opts
		 * is set via annotation args
		 */
		if (disassembler_style)
			perf_cs_option(*cs_handle, CS_OPT_SYNTAX, CS_OPT_SYNTAX_ATT);
		/*
		 * Resolving address operands to symbols is implemented
		 * on x86 by investigating instruction details.
		 */
		perf_cs_option(*cs_handle, CS_OPT_DETAIL, CS_OPT_ON);
	}

	return 0;
}

static size_t print_insn_x86(struct thread *thread, u8 cpumode, struct cs_insn *insn,
			     int print_opts, FILE *fp)
{
	struct addr_location al;
	size_t printed = 0;

	if (insn->detail && insn->detail->x86.op_count == 1) {
		struct cs_x86_op *op = &insn->detail->x86.operands[0];

		addr_location__init(&al);
		if (op->type == X86_OP_IMM &&
		    thread__find_symbol(thread, cpumode, op->imm, &al)) {
			printed += fprintf(fp, "%s ", insn[0].mnemonic);
			printed += symbol__fprintf_symname_offs(al.sym, &al, fp);
			if (print_opts & PRINT_INSN_IMM_HEX)
				printed += fprintf(fp, " [%#" PRIx64 "]", op->imm);
			addr_location__exit(&al);
			return printed;
		}
		addr_location__exit(&al);
	}

	printed += fprintf(fp, "%s %s", insn[0].mnemonic, insn[0].op_str);
	return printed;
}

ssize_t capstone__fprintf_insn_asm(struct machine *machine, struct thread *thread, u8 cpumode,
				   bool is64bit, const uint8_t *code, size_t code_size, uint64_t ip,
				   int *lenp, int print_opts, FILE *fp)
{
	size_t printed;
	struct cs_insn *insn;
	csh cs_handle;
	size_t count;
	bool is_big_endian = false;
	uint16_t e_machine = thread__e_machine_endian(thread, machine,
						      /*e_flags=*/NULL, &is_big_endian);
	int ret;

	/* TODO: Try to initiate capstone only once but need a proper place. */
	ret = capstone_init(e_machine, &cs_handle, is64bit, is_big_endian,
			    /*disassembler_style=*/true);
	if (ret < 0)
		return ret;

	count = perf_cs_disasm(cs_handle, code, code_size, ip, 1, &insn);
	if (count > 0) {
		if (e_machine == EM_X86_64 || e_machine == EM_386)
			printed = print_insn_x86(thread, cpumode, &insn[0], print_opts, fp);
		else
			printed = fprintf(fp, "%s %s", insn[0].mnemonic, insn[0].op_str);
		if (lenp)
			*lenp = insn->size;
		perf_cs_free(insn, count);
	} else {
		printed = -1;
	}

	perf_cs_close(&cs_handle);
	return printed;
}

static void print_capstone_detail(struct cs_insn *insn, char *buf, size_t len,
				  struct annotate_args *args, u64 addr)
{
	int i;
	struct map *map = args->ms->map;
	struct symbol *sym;

	/* TODO: support more architectures */
	if (!arch__is_x86(args->arch))
		return;

	if (insn->detail == NULL)
		return;

	for (i = 0; i < insn->detail->x86.op_count; i++) {
		struct cs_x86_op *op = &insn->detail->x86.operands[i];
		u64 orig_addr;
		struct map *found_map = NULL;

		if (op->type != X86_OP_MEM)
			continue;

		/* only print RIP-based global symbols for now */
		if (op->mem.base != X86_REG_RIP)
			continue;

		/* get the target address */
		orig_addr = addr + insn->size + op->mem.disp;
		addr = map__objdump_2mem(map, orig_addr);

		if (dso__kernel(map__dso(map))) {
			/*
			 * The kernel maps can be split into sections, let's
			 * find the map first and then search the symbol.
			 */
			found_map = maps__find(map__kmaps(map), addr);
			if (found_map == NULL)
				continue;
			map = found_map;
		}

		/* convert it to map-relative address for search */
		addr = map__map_ip(map, addr);

		sym = map__find_symbol(map, addr);
		if (sym == NULL) {
			map__put(found_map);
			continue;
		}

		if (addr == sym->start) {
			scnprintf(buf, len, "\t# %"PRIx64" <%s>",
				  orig_addr, sym->name);
		} else {
			scnprintf(buf, len, "\t# %"PRIx64" <%s+%#"PRIx64">",
				  orig_addr, sym->name, addr - sym->start);
		}
		map__put(found_map);
		break;
	}
}

struct find_file_offset_data {
	u64 ip;
	u64 offset;
};

/* This will be called for each PHDR in an ELF binary */
static int find_file_offset(u64 start, u64 len, u64 pgoff, void *arg)
{
	struct find_file_offset_data *data = arg;

	if (start <= data->ip && data->ip < start + len) {
		data->offset = pgoff + data->ip - start;
		return 1;
	}
	return 0;
}

int symbol__disassemble_capstone(const char *filename, struct symbol *sym,
				 struct annotate_args *args)
{
	struct annotation *notes = symbol__annotation(sym);
	struct map *map = args->ms->map;
	struct dso *dso = map__dso(map);
	u64 start = map__rip_2objdump(map, sym->start);
	u64 offset;
	int i, count, free_count;
	bool is_64bit = false;
	bool needs_cs_close = false;
	/* Malloc-ed buffer containing instructions read from disk. */
	u8 *code_buf = NULL;
	/* Pointer to code to be disassembled. */
	const u8 *buf;
	u64 buf_len;
	csh handle;
	struct cs_insn *insn = NULL;
	char disasm_buf[512];
	struct disasm_line *dl;
	bool disassembler_style = false;
	uint16_t e_machine = EM_NONE;
	bool is_big_endian = false;

	if (args->options->objdump_path)
		return -1;

	buf = dso__read_symbol(dso, filename, map, sym,
			       &code_buf, &buf_len, &is_64bit);
	if (buf == NULL)
		return errno;

	/* add the function address and name */
	scnprintf(disasm_buf, sizeof(disasm_buf), "%#"PRIx64" <%s>:",
		  start, sym->name);

	args->offset = -1;
	args->line = disasm_buf;
	args->line_nr = 0;
	args->fileloc = NULL;
	args->ms->sym = sym;

	dl = disasm_line__new(args);
	if (dl == NULL)
		goto err;

	annotation_line__add(&dl->al, &notes->src->source);

	if (!args->options->disassembler_style ||
	    !strcmp(args->options->disassembler_style, "att"))
		disassembler_style = true;

	if (args->ms->thread) {
		e_machine = thread__e_machine_endian(args->ms->thread,
						     /*machine=*/NULL,
						     /*e_flags=*/NULL, &is_big_endian);
	} else if (dso) {
		struct maps *kmaps = (map && dso__kernel(dso)) ? map__kmaps(map) : NULL;
		struct machine *kmap_machine = kmaps ? maps__machine(kmaps) : NULL;

		e_machine = dso__e_machine_endian(dso, kmap_machine, /*e_flags=*/NULL,
						  &is_big_endian);
	}
	if (!e_machine || e_machine == EM_NONE) {
		e_machine = thread__e_machine_endian(NULL,
						     /*machine=*/NULL,
						     /*e_flags=*/NULL, &is_big_endian);
	}
	if (capstone_init(e_machine, &handle, is_64bit, is_big_endian, disassembler_style) < 0)
		goto err;

	needs_cs_close = true;

	free_count = count = perf_cs_disasm(handle, buf, buf_len, start, buf_len, &insn);
	for (i = 0, offset = 0; i < count; i++) {
		int printed;

		printed = scnprintf(disasm_buf, sizeof(disasm_buf),
				    "       %-7s %s",
				    insn[i].mnemonic, insn[i].op_str);
		print_capstone_detail(&insn[i], disasm_buf + printed,
				      sizeof(disasm_buf) - printed, args,
				      start + offset);

		args->offset = offset;
		args->line = disasm_buf;

		dl = disasm_line__new(args);
		if (dl == NULL)
			goto err;

		annotation_line__add(&dl->al, &notes->src->source);

		offset += insn[i].size;
	}

	/* It failed in the middle: probably due to unknown instructions */
	if (offset != buf_len) {
		struct list_head *list = &notes->src->source;

		/* Discard all lines and fallback to objdump */
		while (!list_empty(list)) {
			dl = list_first_entry(list, struct disasm_line, al.node);

			list_del_init(&dl->al.node);
			disasm_line__free(dl);
		}
		count = -1;
	}

out:
	if (needs_cs_close) {
		perf_cs_close(&handle);
		if (free_count > 0)
			perf_cs_free(insn, free_count);
	}
	free(code_buf);
	return count < 0 ? count : 0;

err:
	if (needs_cs_close) {
		struct disasm_line *tmp;

		/*
		 * It probably failed in the middle of the above loop.
		 * Release any resources it might add.
		 */
		list_for_each_entry_safe(dl, tmp, &notes->src->source, al.node) {
			list_del(&dl->al.node);
			disasm_line__free(dl);
		}
	}
	count = -1;
	goto out;
}

int symbol__disassemble_capstone_powerpc(const char *filename __maybe_unused,
					 struct symbol *sym __maybe_unused,
					 struct annotate_args *args __maybe_unused)
{
	struct annotation *notes = symbol__annotation(sym);
	struct map *map = args->ms->map;
	struct dso *dso = map__dso(map);
	struct nscookie nsc;
	u64 start = map__rip_2objdump(map, sym->start);
	u64 end = map__rip_2objdump(map, sym->end);
	u64 len = end - start;
	u64 offset;
	int i, fd, count;
	bool is_64bit = false;
	bool needs_cs_close = false;
	u8 *buf = NULL;
	struct find_file_offset_data data = {
		.ip = start,
	};
	csh handle;
	char disasm_buf[512];
	struct disasm_line *dl;
	u32 *line;
	bool disassembler_style = false;
	uint16_t e_machine = EM_NONE;
	bool is_big_endian = false;

	if (args->options->objdump_path)
		return -1;

	nsinfo__mountns_enter(dso__nsinfo(dso), &nsc);
	fd = open(filename, O_RDONLY);
	nsinfo__mountns_exit(&nsc);
	if (fd < 0)
		return -1;

	if (file__read_maps(fd, /*exe=*/true, find_file_offset, &data,
			    &is_64bit) == 0)
		goto err;

	if (!args->options->disassembler_style ||
	    !strcmp(args->options->disassembler_style, "att"))
		disassembler_style = true;

	if (args->ms->thread) {
		e_machine = thread__e_machine_endian(args->ms->thread,
						     /*machine=*/NULL,
						     /*e_flags=*/NULL, &is_big_endian);
	} else if (dso) {
		struct maps *kmaps = (map && dso__kernel(dso)) ? map__kmaps(map) : NULL;
		struct machine *kmap_machine = kmaps ? maps__machine(kmaps) : NULL;

		e_machine = dso__e_machine_endian(dso, kmap_machine, /*e_flags=*/NULL,
						  &is_big_endian);
	}
	if (!e_machine || e_machine == EM_NONE) {
		e_machine = thread__e_machine_endian(NULL,
						     /*machine=*/NULL,
						     /*e_flags=*/NULL, &is_big_endian);
	}
	if (capstone_init(e_machine, &handle, is_64bit, is_big_endian, disassembler_style) < 0)
		goto err;

	needs_cs_close = true;

	buf = malloc(len);
	if (buf == NULL)
		goto err;

	count = pread(fd, buf, len, data.offset);
	close(fd);
	fd = -1;

	if ((u64)count != len)
		goto err;

	line = (u32 *)buf;

	/* add the function address and name */
	scnprintf(disasm_buf, sizeof(disasm_buf), "%#"PRIx64" <%s>:",
		  start, sym->name);

	args->offset = -1;
	args->line = disasm_buf;
	args->line_nr = 0;
	args->fileloc = NULL;
	args->ms->sym = sym;

	dl = disasm_line__new(args);
	if (dl == NULL)
		goto err;

	annotation_line__add(&dl->al, &notes->src->source);

	/*
	 * TODO: enable disassm for powerpc
	 * count = cs_disasm(handle, buf, len, start, len, &insn);
	 *
	 * For now, only binary code is saved in disassembled line
	 * to be used in "type" and "typeoff" sort keys. Each raw code
	 * is 32 bit instruction. So use "len/4" to get the number of
	 * entries.
	 */
	count = len/4;

	for (i = 0, offset = 0; i < count; i++) {
		args->offset = offset;
		sprintf(args->line, "%x", line[i]);

		dl = disasm_line__new(args);
		if (dl == NULL)
			break;

		annotation_line__add(&dl->al, &notes->src->source);

		offset += 4;
	}

	/* It failed in the middle */
	if (offset != len) {
		struct list_head *list = &notes->src->source;

		/* Discard all lines and fallback to objdump */
		while (!list_empty(list)) {
			dl = list_first_entry(list, struct disasm_line, al.node);

			list_del_init(&dl->al.node);
			disasm_line__free(dl);
		}
		count = -1;
	}

out:
	if (needs_cs_close)
		perf_cs_close(&handle);
	free(buf);
	return count < 0 ? count : 0;

err:
	if (fd >= 0)
		close(fd);
	count = -1;
	goto out;
}