/*
 * SPDX-License-Identifier: GPL-2.0-only
 *
 * Copyright 2009 Red Hat, Inc.
 *
 * Author: Peter Jones <pjones@redhat.com>
 */
#include <dlfcn.h>
#include <gelf.h>
#include <stdio.h>
#include <strings.h>
#include <string.h>
#include <fcntl.h>
#include <stdlib.h>
#include <unistd.h>
#include <assert.h>

#include "elfcreator.h"

struct elf_creator {
	const char *path;
	int fd;

	Elf *elf;
	GElf_Ehdr *ehdr, ehdr_mem;

	Elf *oldelf;
	/* just because we have to look this up /so/ often... */
	Elf_Scn *dynscn;
	GElf_Shdr *dynshdr, dynshdr_mem;
	Elf_Data *dyndata;
};

static void clear(ElfCreator *ctor, int do_unlink)
{
	if (do_unlink) {
		if (ctor->elf)
			elf_end(ctor->elf);
		if (ctor->fd >= 0)
			close(ctor->fd);
		if (ctor->path)
			unlink(ctor->path);
	} else {
		if (ctor->elf) {
			elf_update(ctor->elf, ELF_C_WRITE_MMAP);
			elf_end(ctor->elf);
		}
		if (ctor->fd >= 0)
			close(ctor->fd);
	}
	memset(ctor, '\0', sizeof(*ctor));
}

ElfCreator *elfcreator_begin(char *path, Elf *elf) {
	ElfCreator *ctor = NULL;
	GElf_Ehdr ehdr_mem, *ehdr;
	GElf_Half machine;

	if (!(ctor = calloc(1, sizeof(*ctor))))
		return NULL;

	clear(ctor, 0);

	ctor->path = path;
	ctor->oldelf = elf;

	ehdr = gelf_getehdr(elf, &ehdr_mem);
	machine = ehdr->e_machine;

	if ((ctor->fd = open(path, O_RDWR|O_CREAT|O_TRUNC, 0755)) < 0) {
err:
		clear(ctor, 1);
		free(ctor);
		return NULL;
	}

	if (!(ctor->elf = elf_begin(ctor->fd, ELF_C_WRITE_MMAP, elf)))
		goto err;

	gelf_newehdr(ctor->elf, gelf_getclass(elf));
	gelf_update_ehdr(ctor->elf, ehdr);

	if (!(ctor->ehdr = gelf_getehdr(ctor->elf, &ctor->ehdr_mem)))
		goto err;

	return ctor;
}

static Elf_Scn *get_scn_by_type(ElfCreator *ctor, Elf64_Word sh_type)
{
	Elf_Scn *scn = NULL;

	while ((scn = elf_nextscn(ctor->elf, scn)) != NULL) {
		GElf_Shdr *shdr, shdr_mem;

		shdr = gelf_getshdr(scn, &shdr_mem);
		if (shdr->sh_type == sh_type)
			return scn;
	}
	return NULL;
}

static void update_dyn_cache(ElfCreator *ctor)
{
	ctor->dynscn = get_scn_by_type(ctor, SHT_DYNAMIC);
	if (ctor->dynscn == NULL)
		return;

	ctor->dynshdr = gelf_getshdr(ctor->dynscn, &ctor->dynshdr_mem);
	ctor->dyndata = elf_getdata(ctor->dynscn, NULL);
}

void elfcreator_copy_scn(ElfCreator *ctor, Elf *src, Elf_Scn *scn)
{
	Elf_Scn *newscn;
	Elf_Data *indata, *outdata;
	GElf_Shdr *oldshdr, oldshdr_mem;
	GElf_Shdr *newshdr, newshdr_mem;

	newscn = elf_newscn(ctor->elf);
	newshdr = gelf_getshdr(newscn, &newshdr_mem);

	oldshdr = gelf_getshdr(scn, &oldshdr_mem);

	memmove(newshdr, oldshdr, sizeof(*newshdr));
	gelf_update_shdr(newscn, newshdr);

	indata = NULL;
	while ((indata = elf_getdata(scn, indata)) != NULL) {
		outdata = elf_newdata(newscn);
		*outdata = *indata;
	}
	if (newshdr->sh_type == SHT_DYNAMIC)
		update_dyn_cache(ctor);
}

static GElf_Dyn *get_dyn_by_tag(ElfCreator *ctor, Elf64_Sxword d_tag,
				GElf_Dyn *mem, size_t *idx)
{
	size_t cnt;

	if (!ctor->dyndata)
		return NULL;

	for (cnt = 1; cnt < ctor->dynshdr->sh_size / ctor->dynshdr->sh_entsize;
			cnt++) {
		GElf_Dyn *dyn;

		if ((dyn = gelf_getdyn(ctor->dyndata, cnt, mem)) == NULL)
			break;

		if (dyn->d_tag == d_tag) {
			*idx = cnt;
			return dyn;
		}
	}
	return NULL;
}

static void remove_dyn(ElfCreator *ctor, size_t idx)
{
	size_t cnt;

	for (cnt = idx; cnt < ctor->dynshdr->sh_size/ctor->dynshdr->sh_entsize;
			cnt++) {
		GElf_Dyn *dyn, dyn_mem;

		if (cnt+1 == ctor->dynshdr->sh_size/ctor->dynshdr->sh_entsize) {
			memset(&dyn_mem, '\0', sizeof(dyn_mem));
			gelf_update_dyn(ctor->dyndata, cnt, &dyn_mem);
			break;
		}

		dyn = gelf_getdyn(ctor->dyndata, cnt+1, &dyn_mem);
		gelf_update_dyn(ctor->dyndata, cnt, dyn);
	}
	ctor->dynshdr->sh_size--;
	gelf_update_shdr(ctor->dynscn, ctor->dynshdr);
	update_dyn_cache(ctor);
}

typedef void (*dyn_fixup_fn)(ElfCreator *ctor, Elf64_Sxword d_tag, Elf_Scn *scn);

static void generic_dyn_fixup_fn(ElfCreator *ctor, Elf64_Sxword d_tag, Elf_Scn *scn)
{
	GElf_Shdr *shdr, shdr_mem;
	GElf_Dyn *dyn, dyn_mem;
	size_t idx;

	dyn = get_dyn_by_tag(ctor, d_tag, &dyn_mem, &idx);
	shdr = gelf_getshdr(scn, &shdr_mem);
	if (shdr) {
		dyn->d_un.d_ptr = shdr->sh_addr;
		gelf_update_dyn(ctor->dyndata, idx, dyn);
	} else {
		remove_dyn(ctor, idx);
	}
}

static void rela_dyn_fixup_fn(ElfCreator *ctor, Elf64_Sxword d_tag, Elf_Scn *scn)
{
	GElf_Shdr *shdr, shdr_mem;
	GElf_Dyn *dyn, dyn_mem;
	size_t idx;

	dyn = get_dyn_by_tag(ctor, d_tag, &dyn_mem, &idx);
	shdr = gelf_getshdr(scn, &shdr_mem);
	if (shdr) {
		dyn->d_un.d_ptr = shdr->sh_addr;
		gelf_update_dyn(ctor->dyndata, idx, dyn);
	} else {
		remove_dyn(ctor, idx);
		dyn = get_dyn_by_tag(ctor, DT_RELASZ, &dyn_mem, &idx);
		if (dyn) {
			dyn->d_un.d_val = 0;
			gelf_update_dyn(ctor->dyndata, idx, dyn);
		}
	}
}

static void rel_dyn_fixup_fn(ElfCreator *ctor, Elf64_Sxword d_tag, Elf_Scn *scn)
{
	GElf_Shdr *shdr, shdr_mem;
	GElf_Dyn *dyn, dyn_mem;
	size_t idx;

	dyn = get_dyn_by_tag(ctor, d_tag, &dyn_mem, &idx);
	shdr = gelf_getshdr(scn, &shdr_mem);
	if (shdr) {
		dyn->d_un.d_ptr = shdr->sh_addr;
		gelf_update_dyn(ctor->dyndata, idx, dyn);
	} else {
		remove_dyn(ctor, idx);
		dyn = get_dyn_by_tag(ctor, DT_RELSZ, &dyn_mem, &idx);
		if (dyn) {
			dyn->d_un.d_val = 0;
			gelf_update_dyn(ctor->dyndata, idx, dyn);
		}
	}
}

static void fixup_dynamic(ElfCreator *ctor)
{
	struct {
		Elf64_Sxword d_tag;
		Elf64_Word sh_type;
		dyn_fixup_fn fn;
	} fixups[] = {
		{ DT_HASH, SHT_HASH, NULL },
		{ DT_STRTAB, SHT_STRTAB, NULL },
		{ DT_SYMTAB, SHT_SYMTAB, NULL },
		{ DT_RELA, SHT_RELA, rela_dyn_fixup_fn},
		{ DT_REL, SHT_REL, rel_dyn_fixup_fn},
		{ DT_GNU_HASH, SHT_GNU_HASH, NULL },
		{ DT_NULL, SHT_NULL, NULL }
	};
	int i;

	for (i = 0; fixups[i].d_tag != DT_NULL; i++) {
		Elf_Scn *scn;

		scn = get_scn_by_type(ctor, fixups[i].sh_type);
		if (fixups[i].fn)
			fixups[i].fn(ctor, fixups[i].d_tag, scn);
		else
			generic_dyn_fixup_fn(ctor, fixups[i].d_tag, scn);
	}
}

void elfcreator_end(ElfCreator *ctor)
{
	GElf_Phdr phdr_mem, *phdr;
	int m,n;

	for (m = 0; (phdr = gelf_getphdr(ctor->oldelf, m, &phdr_mem)) != NULL; m++)
		/* XXX this should check if an entry is needed */;

	gelf_newphdr(ctor->elf, m);
	elf_update(ctor->elf, ELF_C_NULL);
	update_dyn_cache(ctor);

	for (n = 0; n < m; n++) {
		/* XXX this should check if an entry is needed */
		phdr = gelf_getphdr(ctor->oldelf, n, &phdr_mem);
		if (ctor->dynshdr && phdr->p_type == PT_DYNAMIC)
			phdr->p_offset = ctor->dynshdr->sh_offset;

		gelf_update_phdr(ctor->elf, n, phdr);
	}

	fixup_dynamic(ctor);

	clear(ctor, 0);
	free(ctor);
}