/*
 * Copyright (C) 2009  Internet Systems Consortium, Inc. ("ISC")
 *
 * Permission to use, copy, modify, and/or distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND ISC DISCLAIMS ALL WARRANTIES WITH
 * REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
 * AND FITNESS.  IN NO EVENT SHALL ISC BE LIABLE FOR ANY SPECIAL, DIRECT,
 * INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
 * LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE
 * OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
 * PERFORMANCE OF THIS SOFTWARE.
 */

/* $Id$ */
#ifdef __cplusplus
extern "C" {
#endif	/* __cplusplus */
#include <sys/types.h>

#include <elf.h>
#include <fcntl.h>
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>

static int show_backtrace(void **, int);
static int foo();
static int bar();
int main(int, char **);

#define MAX_OFFSET 256		/* ad-hoc */

#ifdef __ia64__
static void *trace0[3];
#else
static void *trace0[3] = { (void *)bar, (void *)foo, (void *)main };
#endif

extern void __register_frame (void *begin);
struct _dwarf_fde {
	uint32_t length;
	uint32_t CIE_delta;
};
struct _dwarf_cie {
	uint32_t length;
	uint32_t CIE_id;
};

static void
get_sectioninfo(int fd, Elf64_Ehdr *ehdr, unsigned int secid, 
		int *nidp, off_t *offsetp, size_t *sizep)
{
	off_t offset;
	size_t hdrsize;
	Elf64_Shdr shdr;

	if (ehdr->e_ident[EI_CLASS] == ELFCLASS32) {
		offset = (off_t)((Elf32_Ehdr *)ehdr)->e_shoff +
			secid * ((Elf32_Ehdr *)ehdr)->e_shentsize;
		hdrsize = (size_t)(((Elf32_Ehdr *)ehdr)->e_shentsize);
	} else {
		offset = (off_t)ehdr->e_shoff + secid * ehdr->e_shentsize;
		hdrsize = (size_t)ehdr->e_shentsize;
	}
	if (lseek(fd, offset, SEEK_SET) < 0) {
		perror("lseek for section header");
		exit(1);
	}
	if (read(fd, &shdr, hdrsize) < hdrsize) {
		fprintf(stderr, "failed to read ELF section header\n");
		exit(1);
	}
	if (ehdr->e_ident[EI_CLASS] == ELFCLASS32) {
		*nidp = (int)((Elf32_Shdr *)&shdr)->sh_name;
		*offsetp = (off_t)((Elf32_Shdr *)&shdr)->sh_offset;
		*sizep = (size_t)((Elf32_Shdr *)&shdr)->sh_size;
	} else {
		*nidp = (int)shdr.sh_name;
		*offsetp = (off_t)shdr.sh_offset;
		*sizep = (size_t)shdr.sh_size;
	}
}

static void
read_elfdata(int fd, off_t offset, size_t size, char *buf)
{
	if (lseek(fd, offset, SEEK_SET) < 0) {
		perror("lseek for ELF section");
		exit(1);
	}
	if (read(fd, buf, size) < size) {
		fprintf(stderr, "failed to read ELF section\n");
		exit(1);
	}
}

static void
elfhack(char *progname)
{
	char *cp, *ep;
	Elf64_Ehdr ehdr, *ehdr64;
	Elf32_Ehdr *ehdr32;
	size_t size;
	off_t offset;
	int i, fd, nid, e_shnum;
	unsigned int e_shstrndx;
	struct _dwarf_fde *fde;
	char *shstrndx, *finfo;

	fd = open(progname, O_RDONLY);
	if (fd < 0) {
		perror("open");
		exit(1);
	}
	/* XXX: assume it contains at least sizeof(Elf64_Ehdr) bytes */
	if (read(fd, &ehdr, sizeof(ehdr)) < sizeof(ehdr)) {
		fprintf(stderr, "failed to read ELF header\n");
		exit(1);
	}
	if (ehdr.e_ident[EI_CLASS] == ELFCLASS32) {
		ehdr32 = (Elf32_Ehdr *)&ehdr;
		e_shstrndx = (unsigned int)ehdr32->e_shstrndx;
		e_shnum = ehdr32->e_shnum;
	} else if (ehdr.e_ident[EI_CLASS] == ELFCLASS64) {
		ehdr64 = &ehdr;
		e_shstrndx = (unsigned int)ehdr64->e_shstrndx;
		e_shnum = ehdr64->e_shnum;
	} else {
		fprintf(stderr, "Unknown ELF class: %d\n",
			ehdr.e_ident[EI_CLASS]);
		exit(1);
	}

	/* read string table for section headers. */
	get_sectioninfo(fd, &ehdr, e_shstrndx, &nid, &offset, &size);
	if ((shstrndx = (char *)malloc(size)) == NULL) {
		perror("malloc failed for ELF string table");
		exit(1);
	}
	read_elfdata(fd, offset, size, shstrndx);

	/*
	 * identify the .debug_frame section and read it (XXX: skip validation
	 * for brevity).
	 */
	for (i = 0; i < e_shnum; i++) {
		get_sectioninfo(fd, &ehdr, i, &nid, &offset, &size);
		if (strcmp(&shstrndx[nid], ".debug_frame") == 0) {
			/*
			 * copy the section.  note we need additional 4 bytes
			 * to indicate the end of data.
			 */
			if ((finfo = (char *)malloc(size + sizeof(uint32_t)))
			    == NULL) {
				perror("malloc failed for .debug_frame");
				exit(1);
			}
			read_elfdata(fd, offset, size, finfo);
			break;
		}
	}
	if (i == e_shnum) {
		fprintf(stderr, ".debug_frame section not found\n");
		exit(1);
	}

	free(shstrndx);
	close(fd);

	/* Fix the .debug_frame section */
	cp = finfo;
	ep = cp + size;
	while (cp < ep) {
		/* XXX: length validation is omitted for simplicity */
		fde = (struct _dwarf_fde *)cp;
		if (fde->CIE_delta == 0xffffffff) {
			/* This is CIE */
			fde->CIE_delta = 0;
		} else {
			fde->CIE_delta = (char *)&fde->CIE_delta -
				(finfo + fde->CIE_delta); 
		}
		cp += fde->length + sizeof(fde->length);
	}
	if (cp != ep) {
		fprintf(stderr, "assumption failure: malformed frame info?\n");
		exit(1);
	}
	*(uint32_t *)cp = 0;
#ifndef __ia64__
	/*
	 * __register_frame() doesn't exist for IA64.  But it doesn't matter
	 * in the first place because .debug_frame wouldn't exist either.
	 */
	__register_frame(finfo);
#endif
	/* Don't free finfo.  It will be used in _Unwind_Backtrace() */
}

static void
usage()
{
	fprintf(stderr, "Usage: unwindtest [-e]\n");
	exit(1);
}

static void
setup(int argc, char *argv[])
{
	int ch;
	int do_elfhack = 0;

	while ((ch = getopt(argc, argv, "e")) != -1) {
		switch (ch) {
		case 'e':
			do_elfhack = 1;
			break;
		case '?':
		default:
			usage();
		}
	}

	if (do_elfhack)
		elfhack(argv[0]);

#ifdef __ia64__
	/*
	 * XXX: in IA64, function pointers do not point to the entrance point
	 * of the function code; instead it points to a data structure whose
	 * first quadword is the address of the function entrance.
	 */
	trace0[0] = ((void **)bar)[0];
	trace0[1] = ((void **)foo)[0];
	trace0[2] = ((void **)main)[0];
#endif
}

extern int _Unwind_Backtrace(int (*fn)(void *, void *), void* a);
extern void* _Unwind_GetIP(void* ctx);

typedef struct {
	void **result;
	int max_depth;
	int skip_count;
	int count;
} trace_arg_t;

static int
btcallback(void *uc, void *opq) {
	trace_arg_t *arg = (trace_arg_t *)opq;

	if (arg->skip_count > 0)
		arg->skip_count--;
	else
		arg->result[arg->count++] = (void *)_Unwind_GetIP(uc);
	if (arg->count == arg->max_depth)
		return (5); /* _URC_END_OF_STACK */

	return (0); /* _URC_NO_REASON */
}

static int
show_backtrace(void **tracebuf, int max_depth) {
	trace_arg_t arg;

	arg.skip_count = 1;
	arg.result = tracebuf;
	arg.max_depth = max_depth;
	arg.count = 0;
	_Unwind_Backtrace(btcallback, &arg);

	return (arg.count);
}

static int
bar() {
	int i, nframe;
	int ret = 0;
	void *trace[3];

	memset(trace, 0, sizeof(trace));
	nframe = show_backtrace(trace, 3);
	if (nframe < 3)
		ret = 1;
	for (i = 0; i < nframe; i++) {
		printf("frame[%d] %p: entrance pt=%p, diff=%ld\n", i,
		       trace[i], trace0[i],
		       (long)((char *)trace[i] - (char *)trace0[i]));
		if (trace[i] < trace0[i])
			ret = 2;
		if ((char *)trace[i] - (char *)trace0[i] > MAX_OFFSET)
			ret = 3;
	}

	return (ret);
}

static int
foo() {
	return (bar());
}

int
main(int argc, char *argv[])
{
	setup(argc, argv);
	return (foo());
}

#ifdef __cplusplus
}
#endif	/* __cplusplus */