#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <inttypes.h>
#include <dwarf.h>
#include <elfutils/libdwfl.h>

#define INTUSE(x)	x

/* From libdwfl/argp-std.c */

static char *debuginfo_path;

static const Dwfl_Callbacks offline_callbacks =
  {
    .find_debuginfo = INTUSE(dwfl_standard_find_debuginfo),
    .debuginfo_path = &debuginfo_path,
    .section_address = INTUSE(dwfl_offline_section_address),
  };


void error(char *msg)
{
	printf("Error %s: %s\n", msg, dwarf_errmsg(-1));
}

/*
 * @ra_regno contains the return address. We need to find
 * the expression for that register
 */
static int
check_return_reg(int ra_regno, Dwarf_Frame *frame)
{
	Dwarf_Op ops_mem[2];
	Dwarf_Op dummy;
	Dwarf_Op *ops = &dummy;
  	size_t nops;
	int i, result;

	result = dwarf_frame_register(frame, ra_regno, ops_mem, &ops, &nops);
	if (result < 0) {
		error("dwarf_frame_register()");
		return -1;
	}

	if (nops != 0 || ops != NULL)
		return 0;	/* return address on the stack */

	/*
	 * Return address in LR. Check if a frame was allocated
	 * but not yet used.
	 */
	result = dwarf_frame_cfa(frame, &ops, &nops);
	if (result < 0) {
		error("dwarf_frame_cfa()");
		return -1;
	}

	if (nops == 1 && ops[0].atom == DW_OP_bregx &&
				ops[0].number == 1 && ops[0].number2 == 0)

		return 1;	/* no new frame allocated */

	return 2;		/* new frame allocated, but not yet used */
}

/*
 * Return 1 if the address @pc is in LR. O otherwise.
 * Return -1 in case of errors
 */
int check_ret_addr_in_lr(const char *exec_file, Dwarf_Addr pc)
{
	int		result;
	Dwfl		*dwfl;
	Dwfl_Module	*mod;
	Dwarf_CFI	*cfi;
	Dwarf_Frame	*frame;
	Dwarf_Addr	bias;
	int		ra_regno;
	Dwarf_Addr 	start = pc;
	Dwarf_Addr 	end = pc;
	bool 		signalp;

	dwfl = dwfl_begin(&offline_callbacks);
	if (dwfl == NULL) {
		error("dwfl_begin() failed\n");
		return -1;
	}

	if (dwfl_report_offline(dwfl, "",  exec_file, -1) == NULL) {
		error("dwfl_report_offline()");
		return -1;
	}

	mod = dwfl_addrmodule(dwfl, pc);
	if (!mod) {
		error("dwfl_addrmodule");
		return -1;
	}

	cfi = dwfl_module_dwarf_cfi(mod, &bias);
	if (!cfi) {
		printf("No dwarf cfi, trying eh cfi\n");
		cfi = dwfl_module_eh_cfi(mod, &bias);	// CHECK
		if (!cfi) {
			error("dwfl_module_dwarf_cfi(): no CFI -");
			return -1;
		}
	}

	result = dwarf_cfi_addrframe(cfi, pc - bias, &frame);
	if (result != 0) {
		error("dwarf_cfi_addrframe(): ");
		return -1;
	}

	ra_regno = dwarf_frame_info(frame, &start, &end, &signalp);

	if (ra_regno < 0) {
		printf("\treturn address register unavailable (%s)\n",
				dwarf_errmsg (0));
		return -1;
	}

	return check_return_reg(ra_regno, frame);
}

int
main(int argc, const char *argv[])
{
	int 		rc;
	const char 	*exec_file;
	Dwarf_Addr	pc = 0x100006c4;
	Dwarf_Addr	end;

	if (argc > 1)
		pc = strtol(argv[1], NULL, 16);

	end = pc;
	if (argc > 2)
		end = strtol(argv[2], NULL, 16);

	exec_file = "a.out";
	if (argc > 3)
		exec_file = argv[3];

	while (pc <= end) {
		rc = check_ret_addr_in_lr(exec_file, pc);
		if (rc < 0)
			printf("Error with addr 0x%" PRIx64 "\n", pc);
		else if (rc == 1) 
			printf("Ret Address 0x%" PRIx64 " is in LR\n", pc);
		else if (rc == 2) 
			printf("Ret Address 0x%" PRIx64 " is in LR, slot1 invalid\n", pc);
		else
			printf("Ret Address 0x%" PRIx64 " is NOT in LR\n", pc);
		pc += 1;
	}
}