#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <sys/types.h>

enum {
	EI_MAG0=0, EI_MAG1, EI_MAG2, EI_MAG3, EI_CLASS, EI_DATA, EI_VERSION,
	EI_OSABI, EI_ABIVERSION, EI_PAD, EI_NIDENT=16
};
enum {
	ELFCLASSNONE=0, ELFCLASS32, ELFCLASS64
};
enum {
	ELFDATANONE=0, ELFDATA2LSB, ELFDATA2MSB
};
enum {
	ET_NONE=0, ET_REL, ET_EXEC, ET_DYN, ET_CORE
};
enum {
	EV_NONE=0, EV_CURRENT
};
enum {
	PT_NULL=0, PT_LOAD, PT_DYNAMIC, PT_INTERP, PT_NOTE, PT_SHLIB, PT_PHDR
};
enum {
	PF_X=1<<0, PF_W=1<<1, PF_R=1<<2
};
enum {
	SHT_PROGBITS=1, SHT_SYMTAB, SHT_STRTAB, SHT_RELA, SHT_HASH,
	SHT_DYNAMIC, SHT_NOTE, SHT_NOBITS
};

struct elf32_ehdr {
	uint8_t  e_ident[EI_NIDENT];
	uint16_t e_type;
	uint16_t e_machine;
	uint32_t e_version;
	uint32_t e_entry;
	uint32_t e_phoff;
	uint32_t e_shoff;
	uint32_t e_flags;
	uint16_t e_ehsize;
	uint16_t e_phentsize;
	uint16_t e_phnum;
	uint16_t e_shentsize;
	uint16_t e_shnum;
	uint16_t e_shstrndx;
};

#define EM_PPC			20

struct elf32_phdr {
	uint32_t p_type;
	uint32_t p_offset;
	uint32_t p_vaddr;
	uint32_t p_paddr;
	uint32_t p_filesz;
	uint32_t p_memsz;
	uint32_t p_flags;
	uint32_t p_align;
};

struct elf32_shdr {
	uint32_t sh_name;
	uint32_t sh_type;
	uint32_t sh_flags;
	uint32_t sh_addr;
	uint32_t sh_offset;
	uint32_t sh_size;
	uint32_t sh_link;
	uint32_t sh_info;
	uint32_t sh_addralign;
	uint32_t sh_entsize;
};

#define SHF_ALLOC		(1 << 1)

struct elf32_sym {
	uint32_t st_name;
	uint32_t st_value;
	uint32_t st_size;
	uint8_t  st_info;
	uint8_t  st_other;
	uint16_t st_shndx;
};

#define SHN_UNDEF		0
#define SHN_ABS			0xfff1
#define SHN_COMMON		0xfff2

struct elf32_rela {
	uint32_t r_offset;
	uint32_t r_info;
	int32_t  r_addend;
};

#define ELF32_R_SYM(i)		((i)>>8)
#define ELF32_R_TYPE(x)		((unsigned char)(x))

#define R_NONE			0
#define R_PPC_ADDR32		1
#define R_PPC_ADDR16_LO		4
#define R_PPC_ADDR16_HI		5
#define R_PPC_ADDR16_HA		6
#define R_PPC_REL24		10
#define R_PPC_REL14		11
#define R_PPC_REL14_BRTAKEN	12
#define R_PPC_REL14_BRNTAKEN	13
#define R_PPC_REL32		26

/* we have to relocate these symbols, although they are absolute */
#define LNKSYMCNT 5
static const char *linkersyms[LNKSYMCNT] = {
	"_edata", "__bss_start", "end", "_end", "__end"
};

static int relocate(uint8_t *, uint32_t);
static int fixreloclist(uint8_t *, struct elf32_rela *, int relacnt,
    struct elf32_sym *, char *, uint32_t, uint32_t);
static int islnksym(char *, struct elf32_sym *);


int main(int argc, char *argv[])
{
	void *buf;
	FILE *f;
	long len;
	int err;

	err = 1;

	if (argc != 2) {
		fprintf(stderr, "usage: %s filename\n", argv[0]);
		return 1;
	}

	if (f = fopen(argv[1], "r+")) {
		(void)fseek(f, 0, SEEK_END);
		len = ftell(f);
		rewind(f);
		buf = malloc(len);
		if (buf != NULL) {
			if (fread(buf, 1, len, f) == len) {
				if (relocate(buf, 0x78000000)) {
					rewind(f);
					if (fwrite(buf, 1, len, f) == len)
						err = 0;
					else
						fprintf(stderr,
						    "Write error!\n");
				}
			}
			else
				fprintf(stderr, "Read error!\n");
			free(buf);
		} else
			fprintf(stderr, "Cannot allocate memory!\n");
		fclose(f);
	} else
		perror(argv[1]);

	return err;
}

/*
 * Relocate a PPC ELF kernel to a new base address.
 */
static int relocate(uint8_t *kern, uint32_t relocaddr)
{
	struct elf32_ehdr *hdr;
	struct elf32_phdr *phdr;
	struct elf32_shdr *shdr;
	struct elf32_sym *symbols;
	char *strtab;
	uint32_t origaddr;
	int i, str, sym, symcnt;

	hdr = (struct elf32_ehdr *)kern;
	if (strncmp(hdr->e_ident, "\177ELF", 4) ||
	    hdr->e_ident[EI_CLASS] != ELFCLASS32 ||
	    hdr->e_ident[EI_DATA] != ELFDATA2MSB ||
	    hdr->e_ident[EI_VERSION] != EV_CURRENT ||
	    hdr->e_phnum != 1 || hdr->e_type != ET_EXEC ||
	    hdr->e_machine != EM_PPC || hdr->e_version != EV_CURRENT ||
	    hdr->e_ehsize != sizeof(struct elf32_ehdr) ||
	    hdr->e_phentsize != sizeof(struct elf32_phdr) ||
	    hdr->e_shentsize != sizeof(struct elf32_shdr)) {
		fprintf(stderr,
		    "Not a single-segment PowerPC ELF executable!\n");
		return 0;
	}
	if (hdr->e_shnum * sizeof(struct elf32_shdr) == 0) {
		fprintf(stderr, "Missing section headers!\n");
		return 0;
	}

	/* relocate program entry and vaddr of first segment */
	phdr = (struct elf32_phdr *)(kern + hdr->e_phoff);
	origaddr = phdr[0].p_vaddr;
	phdr[0].p_vaddr = phdr[0].p_paddr = relocaddr;
	hdr->e_entry = hdr->e_entry - origaddr + relocaddr;

	/* locate symbol table and string table */
	shdr = (struct elf32_shdr *)(kern + hdr->e_shoff);
	for (i = 0, sym = -1; i < hdr->e_shnum; i++)
		if (shdr[i].sh_type == SHT_SYMTAB) {
			sym = i;
			str = shdr[i].sh_link;
			strtab = kern + shdr[str].sh_offset;
			break;
		}
	if (sym == -1) {
		fprintf(stderr, "Missing symbol table!\n");
		return 0;
	}
	if (shdr[sym].sh_size % shdr[sym].sh_entsize != 0) {
		fprintf(stderr, "Corrupted symbol table!\n");
		return 0;
	}
	symcnt = shdr[sym].sh_size / shdr[sym].sh_entsize;

	/*
	 * Relocate all symbols. Take care of some linker symbols,
	 * which also have to be relocated.
	 */
	symbols = (struct elf32_sym *)(kern + shdr[sym].sh_offset);
	for (i = 0; i < symcnt; i++) {
		if (symbols[i].st_shndx == SHN_UNDEF ||
		    symbols[i].st_shndx >= SHN_ABS)
			if (!islnksym(strtab, &symbols[i]))
				continue;
		symbols[i].st_value =
		    symbols[i].st_value - origaddr + relocaddr;
	}

	/* fix all relocation entries and the VA of allocated sections */
	for (i = 0; i < hdr->e_shnum; i++) {
		if ((shdr[i].sh_type == SHT_PROGBITS ||
		    shdr[i].sh_type == SHT_NOBITS ||
		    shdr[i].sh_type == SHT_NOTE) &&
		    (shdr[i].sh_flags & SHF_ALLOC) != 0)
			shdr[i].sh_addr =
			    shdr[i].sh_addr - origaddr + relocaddr;

		else if (shdr[i].sh_type == SHT_RELA) {
			if (shdr[i].sh_size % shdr[i].sh_entsize != 0) {
				fprintf(stderr, "Corrupted reloc section!\n");
				return 0;
			}
			if (fixreloclist(kern + phdr[0].p_offset,
			    (struct elf32_rela *)(kern + shdr[i].sh_offset),
			    shdr[i].sh_size / shdr[i].sh_entsize,
			    symbols, strtab, origaddr, relocaddr) == 0)
				return 0;
		}
	}

	return 1;
}

/* fix all relocation of a section */
static int fixreloclist(uint8_t *text, struct elf32_rela *rela, int relacnt,
    struct elf32_sym *syms, char *strtab, uint32_t origaddr, uint32_t relocaddr)
{
	uint32_t o, v;
	int i, j;

	for (i = 0; i < relacnt; i++) {
		j = ELF32_R_SYM(rela[i].r_info);
		if (syms[j].st_shndx == SHN_UNDEF ||
		    syms[j].st_shndx >= SHN_ABS)
			if (!islnksym(strtab, &syms[j]))
				continue;
		v = syms[j].st_value + rela[i].r_addend;
		o = rela[i].r_offset - origaddr;
		rela[i].r_offset = relocaddr + o;

		switch (ELF32_R_TYPE(rela[i].r_info)) {
		case R_PPC_ADDR32:
			*(uint32_t *)(text + o) = v;
			break;
		case R_PPC_ADDR16_HI:
			*(uint16_t *)(text + o) = v >> 16;
			break;
		case R_PPC_ADDR16_HA:
			*(uint16_t *)(text + o) =
			    (v >> 16) + (v & 0x8000 ? 1 : 0);
			break;
		case R_PPC_ADDR16_LO:
			*(uint16_t *)(text + o) = v & 0xffff;
			break;
		case R_NONE:
		case R_PPC_REL32:
		case R_PPC_REL24:
		case R_PPC_REL14:
		case R_PPC_REL14_BRTAKEN:
		case R_PPC_REL14_BRNTAKEN:
			break;
		default:
			fprintf(stderr, "Unknown relocation type!\n");
			return 0;
		}
	}
	return 1;
}

/*
 * Check whether this is a linker symbol, which needs to be relocated
 * despite being absolute.
 * XXX Did we get all of them? What about new ones?
 */
static int islnksym(char *str, struct elf32_sym *sym)
{
	int i;
	char *name;

	name = str + sym->st_name;
	if (sym->st_shndx == SHN_ABS) {
		if (!strncmp(name, "__start_link_set", 16) ||
		    !strncmp(name, "__stop_link_set", 15))
			return 1;

		for (i = 0; i < LNKSYMCNT; i++) {
			if (!strcmp(name, linkersyms[i]))
				return 1;
		}
	}
	return 0;
}