/* ----------------------------------------------------------------------- *
 *
 *   Copyright 2008 rPath, Inc. - All Rights Reserved
 *
 *   This program is free software; you can redistribute it and/or modify
 *   it under the terms of the GNU General Public License as published by
 *   the Free Software Foundation, Inc., 51 Franklin St, Fifth Floor,
 *   Boston MA 02110-1301, USA; either version 2 of the License, or
 *   (at your option) any later version; incorporated herein by reference.
 *
 * ----------------------------------------------------------------------- */

/*
 * elf.c
 *
 * Take a linked list of segments and output it as an ELF32 image
 */

#include "wraplinux.h"

#include <assert.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <inttypes.h>
#include <stdbool.h>
#include "elf32.h"
#include "segment.h"
#include "multiboot.h"
#include "le.h"

static inline uint32_t p_type(const struct segment *s)
{
	switch (s->sh_type) {
	case SHT_PROGBITS:
	case SHT_NOBITS:
		return (s->sh_flags & SHF_ALLOC) ? PT_LOAD : PT_NULL;
	case SHT_NOTE:
		return PT_NOTE;
	default:
		return PT_NULL; /* No other types supported */
	}
}

/*
 * Return true if we should use a single PHDR for s and s->next
 */
static bool merge_phdrs(const struct segment *s)
{
	const struct segment *n = s->next;

	if (!n)
		return false;

	if (p_type(s) != PT_LOAD || p_type(n) != PT_LOAD)
		return false;

	if (s->sh_type != n->sh_type)
		return false;

	if (s->sh_flags != n->sh_flags)
		return false;

	if (s->align < n->align)
		return false;

	if (n->address != align_up(s->address+s->length, n->align))
		return false;

	return true;
}

static int gen_elf(struct segment *segs, Elf32_Ehdr *ehdr, FILE *out)
{
	Elf32_Shdr shdr;
	Elf32_Phdr phdr;
	int nsections = 0, npheader = 0;
	int namebytes = 11;	/* Null entry plus .shstrtab */
	struct segment *s;
	uint32_t data_start, data_offset, file_offset;
	uint32_t name_offset;

	for (s = segs; s; s = s->next) {
		nsections++;
		if (s->name)
			namebytes += strlen(s->name)+1;
	}

	for (s = segs; s; s = s->next) {
		uint32_t memsize;
		addr_t align = (addr_t)1 << s->align;
		addr_t address, pad;

		if (p_type(s) == PT_NULL)
			continue;

		npheader++;

		address = s->address;
		memsize = s->length;

		while (merge_phdrs(s)) {
			s = s->next;
			align = UINT32_C(1) << s->align;
			pad = (s->address - (address+memsize)) & (align-1);
			memsize += pad + s->length;
		}
	}

	wrle16(sizeof phdr, &ehdr->e_phentsize);
	wrle16(npheader, &ehdr->e_phnum);
	wrle16(sizeof shdr, &ehdr->e_shentsize);
	wrle16(nsections+2, &ehdr->e_shnum);
	wrle16(nsections+1, &ehdr->e_shstrndx);
	file_offset = c_fwrite(ehdr, sizeof *ehdr, out);

	/* First actual data byte */
	data_start = (sizeof *ehdr) + npheader*(sizeof phdr);

	/* Program header table */
	data_offset = data_start;
	for (s = segs; s; s = s->next) {
		uint32_t filesize, memsize;
		uint32_t ptype = p_type(s);
		uint32_t p_flags;
		uint32_t align = UINT32_C(1) << s->align;
		addr_t address, offset, pad;

		filesize = (s->sh_type == SHT_NOBITS) ? 0 : s->length;
		memsize = (s->sh_flags & SHF_ALLOC) ? s->length : 0;

		if (s->sh_type == SHT_NOBITS)
			pad = 0;
		else
			pad = (s->address - data_offset) & (align-1);

		data_offset += pad;

		if (ptype != PT_NULL) {
			p_flags = 0;
			if (s->sh_flags & SHF_ALLOC) /* Somewhat cheesy... */
				p_flags |= PF_R;
			if (s->sh_flags & SHF_EXECINSTR)
				p_flags |= PF_X;
			if (s->sh_flags & SHF_WRITE)
				p_flags |= PF_W;

			memset(&phdr, 0, sizeof phdr);
			wrle32(ptype, &phdr.p_type);
			wrle32(data_offset, &phdr.p_offset);
			wrle32(s->address, &phdr.p_vaddr);
			wrle32(s->address, &phdr.p_paddr);
			wrle32(p_flags, &phdr.p_flags);
			wrle32(align, &phdr.p_align);

			address = s->address;
			offset  = data_offset + filesize;

			while (merge_phdrs(s)) {
				s = s->next;
				align = UINT32_C(1) << s->align;
				pad = (s->address - (address+memsize))
					& (align-1);
				if (s->sh_type != SHT_NOBITS) {
					assert(pad == ((s->address - offset)
						       & (align-1)));
					filesize += s->length + pad;
					offset   += s->length + pad;
				}
				memsize += s->length + pad;
			}

			wrle32(filesize, &phdr.p_filesz);
			wrle32(memsize, &phdr.p_memsz);

			file_offset += c_fwrite(&phdr, sizeof phdr, out);
		}
		data_offset += filesize;
	}

	/* Actual file data */
	assert(file_offset == data_start);
	for (s = segs; s; s = s->next) {
		if (s->sh_type != SHT_NOBITS) {
			uint32_t align = UINT32_C(1) << s->align;
			uint32_t pad = (s->address - file_offset) & (align-1);
			file_offset += c_writezero(pad, out);
			file_offset += c_fwrite(s->data, s->length, out);
		}
	}

	/* String table */
	file_offset += c_fwrite("", 1, out);	/* Null string table entry */
	for (s = segs; s; s = s->next) {
		uint32_t len = strlen(s->name)+1;
		file_offset += c_fwrite(s->name, len, out);
	}
	file_offset += c_fwrite(".shstrtab", 10, out);

	/* Align the section header table */
	file_offset += c_writezero(-file_offset & 3, out);

	/* Section header table */
	wrle32(file_offset, &ehdr->e_shoff); /* Valid on the next interation */

	data_offset = data_start;

	memset(&shdr, 0, sizeof shdr);
	c_fwrite(&shdr, sizeof shdr, out); /* Null section header */

	name_offset = 1;
	for (s = segs; s; s = s->next) {
		uint32_t align = UINT32_C(1) << s->align;

		if (s->sh_type != SHT_NOBITS)
			data_offset += (s->address - data_offset) & (align-1);

		memset(&shdr, 0, sizeof shdr);
		wrle32(s->name ? name_offset : 0, &shdr.sh_name);
		wrle32(s->sh_type, &shdr.sh_type);
		wrle32(s->sh_flags, &shdr.sh_flags);
		wrle32(s->address, &shdr.sh_addr);
		wrle32(data_offset, &shdr.sh_offset); /* Even for SHT_NOBITS */
		wrle32(s->length, &shdr.sh_size);
		wrle32(align, &shdr.sh_addralign);

		file_offset += c_fwrite(&shdr, sizeof shdr, out);

		if (s->name)
			name_offset += strlen(s->name)+1;

		if (s->sh_type != SHT_NOBITS)
			data_offset += s->length;
	}

	/* String table section header */
	memset(&shdr, 0, sizeof shdr);
	wrle32(name_offset, &shdr.sh_name);
	wrle32(SHT_STRTAB, &shdr.sh_type);
	wrle32(data_offset, &shdr.sh_offset);
	wrle32(namebytes, &shdr.sh_size);
	wrle32(1, &shdr.sh_addralign);

	c_fwrite(&shdr, sizeof shdr, out);

	return 0;
}


static int do_elf(struct segment *segs, addr_t entry, FILE *out)
{
	Elf32_Ehdr ehdr;
	int rv;

	memset(&ehdr, 0, sizeof ehdr);

	ehdr.e_ident[EI_MAG0]    = ELFMAG0;
	ehdr.e_ident[EI_MAG1]    = ELFMAG1;
	ehdr.e_ident[EI_MAG2]    = ELFMAG2;
	ehdr.e_ident[EI_MAG3]    = ELFMAG3;
	ehdr.e_ident[EI_CLASS]   = ELFCLASS32;
	ehdr.e_ident[EI_DATA]    = ELFDATA2LSB;
	ehdr.e_ident[EI_VERSION] = EV_CURRENT;
	ehdr.e_ident[EI_OSABI]   = ELFOSABI_STANDALONE;

	wrle16(ET_EXEC, &ehdr.e_type);
	wrle16(EM_386, &ehdr.e_machine);
	wrle32(EV_CURRENT, &ehdr.e_version);
	wrle32(sizeof ehdr, &ehdr.e_phoff);
	wrle32(entry, &ehdr.e_entry);
	wrle16(sizeof ehdr, &ehdr.e_ehsize);

	/* Dummy run which produces some additional data */
	rv = gen_elf(segs, &ehdr, NULL);
	if (rv)
		return rv;

	/* The real run */
	return gen_elf(segs, &ehdr, out);
}

int output_elf(struct segment *segs, addr_t entry, FILE *out)
{
	segs = sort_segments(segs);
	return do_elf(segs, entry, out);
}

int output_multiboot(struct segment *segs, addr_t entry, FILE *out)
{
	struct segment mb_seg;
	struct multiboot_header mb_hdr;
	uint32_t magic, flags, csum;

	segs = sort_segments(segs);

	/* Stick a Multiboot header in a segment at
	   the beginning of the file. */
	mb_seg.next	= segs;
	mb_seg.length	= sizeof mb_hdr;
	mb_seg.align	= 2;
	mb_seg.address	= 0;
	mb_seg.sh_type	= SHT_PROGBITS;
	mb_seg.sh_flags	= 0;	/* No ALLOC, EXECINSTR, or WRITE */
	mb_seg.name	= ".multiboot";
	mb_seg.data	= &mb_hdr;
	segs = &mb_seg;

	/* Multiboot header */
	magic	= MULTIBOOT_MAGIC;
	flags	= MB_FLAG_PAGE;

	memset(&mb_hdr, 0, sizeof mb_hdr);
	wrle32(magic, &mb_hdr.magic);
	wrle32(flags, &mb_hdr.flags);

	csum    = magic + flags;
	wrle32(-csum, &mb_hdr.checksum);

	/* Generate the ELF image */
	return do_elf(segs, entry, out);
}
