/* ----------------------------------------------------------------------- *
 *
 *   Copyright 2008 rPath, Inc. - All Rights Reserved
 *   Copyright 2010 Intel Corporation; author: H. Peter Anvin
 *
 *   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.
 *
 * ----------------------------------------------------------------------- */

/*
 * wrap.c
 *
 * Actually wrap the kernel image...
 */

#include <ctype.h>
#include <stdio.h>
#include <unistd.h>
#include <getopt.h>
#include <errno.h>
#include <stdlib.h>
#include <string.h>
#include <fcntl.h>
#include <sysexits.h>
#include <sys/types.h>
#include "segment.h"
#include "setup.h"
#include "elf32.h"
#include "le.h"
#include "wraplinux.h"

extern char reloc[];
extern uint32_t reloc_size;
extern char highmove[];
extern uint32_t highmove_size;

/* Find the last instance of a particular command line argument
   (which should include the final =; do not use for boolean arguments) */
static const char *find_argument(const char *cmdline, const char *argument)
{
	int la = strlen(argument);
	const char *p = cmdline;
	int wasspace = 1;

	while (*p) {
		if (wasspace && !memcmp(p, argument, la))
			return p + la;

		wasspace = isspace(*p++);
	}

	return NULL;
}

/* Truncate to 32 bits, with saturate */
static inline uint32_t saturate32(unsigned long long v)
{
	return (v > 0xffffffff) ? 0xffffffff : (uint32_t) v;
}

/* Get a value with a potential suffix (k/m/g/t/p/e) */
static unsigned long long suffix_number(const char *str)
{
	char *ep;
	unsigned long long v;
	int shift;

	v = strtoull(str, &ep, 0);
	switch (*ep | 0x20) {
	case 'k':
		shift = 10;
		break;
	case 'm':
		shift = 20;
		break;
	case 'g':
		shift = 30;
		break;
	case 't':
		shift = 40;
		break;
	case 'p':
		shift = 50;
		break;
	case 'e':
		shift = 60;
		break;
	default:
		shift = 0;
		break;
	}
	v <<= shift;

	return v;
}

int wrap_kernel(const char *kernel_file, const char *cmdline,
		const struct string_list *initrd_list, FILE *out)
{
	struct initrd_info {
		int fd;
		struct segment seg;
	} *ird = NULL;
	int kernel_fd = -1;
	char *kernel = NULL;
	size_t kernel_len = 0;
	struct segment srel, ssup, scmd, skrn, shmv;
	struct startup_info *info = (void *)reloc;
	struct setup_header *hdr;
	struct highmove_info *hmv = (void *)highmove;
	size_t setup_len;
	size_t initrd_len;
	size_t initrd_addr;
	int rv = EX_SOFTWARE;	/* Should never be returned... */
	int ninitrd = 0;
	int i;
	const struct string_list *ip;
	int setup_ver;		/* Setup protocol version */
	int setup_space;	/* How much space for the setup */
	const char *cmd;
	uint32_t initrd_max;
	uint32_t entry;

	/* Process the kernel file */

	kernel_fd = open(kernel_file, O_RDONLY);
	if (kernel_fd < 0) {
		fprintf(stderr, "%s: %s: %s\n", program, kernel_file,
			strerror(errno));
		rv = EX_NOINPUT;
		goto err;
	}

	kernel = mapfile(kernel_fd, &kernel_len, 1);
	if (!kernel) {
		fprintf(stderr, "%s: %s: %s\n", program, kernel_file,
			strerror(errno));
		rv = EX_NOINPUT;
		goto err;
	}

	if (kernel_len < 1024) {
		fprintf(stderr, "%s: %s: kernel file too small\n", program,
			kernel_file);
		rv = EX_NOINPUT;
		errno = EINVAL;
		goto err;
	}

	/* Pick apart the header... */

	hdr = (struct setup_header *)(kernel + 0x1f1);
	setup_len = (hdr->setup_sects + 1) << 9;
	if (setup_len == 512)
		setup_len += 2048;	/* Really old kernel */

	if (rdle32(&hdr->header) != LINUX_MAGIC)
		setup_ver = 0;	/* Ancient kernel */
	else
		setup_ver = rdle16(&hdr->version);

	if (setup_ver >= 0x200 && (hdr->loadflags & LOADED_HIGH)) {
		skrn.address = 0x100000;
		ssup.address = 0x10000;
		setup_space = setup_ver >= 0x202 ? 0x10000 : 0xa000;
	} else {
		skrn.address = 0x10000;
		ssup.address = 0x90000;
		setup_space = 0xa000;
	}

	if (setup_ver <= 0x200)
		initrd_list = NULL;	/* No initrd for ancient kernel */

	if (setup_ver >= 0x203)
		initrd_max = rdle32(&hdr->initrd_addr_max);
	else
		initrd_max = 0x37ffffff;

	if ((cmd = find_argument(cmdline, "mem="))) {
		uint32_t mem = saturate32(suffix_number(cmd));
		if (initrd_max >= mem)
			initrd_max = mem - 1;
	}

	if ((cmd = find_argument(cmdline, "vga="))) {
		uint16_t vga;

		switch (cmd[0] | 0x20) {
		case 'a':	/* "ask" */
			vga = 0xfffd;
			break;
		case 'e':	/* "ext" */
			vga = 0xfffe;
			break;
		case 'n':	/* "normal" */
			vga = 0xffff;
			break;
		default:
			vga = strtoul(cmd, NULL, 0);
			break;
		}
		wrle16(vga, &hdr->vid_mode);
	}

	/* Process the initrd file(s) */

	ninitrd = 0;
	for (ip = initrd_list; ip; ip = ip->next)
		ninitrd++;

	if (ninitrd) {
		ird = xcalloc(ninitrd, sizeof *ird);
		for (i = 0; i < ninitrd; i++)
			ird[i].fd = -1;
	}

	initrd_len = 0;
	for (ip = initrd_list, i = 0; ip; ip = ip->next, i++) {
		/* Each sub-initrd is aligned to a 4-byte boundary */
		initrd_len = align_up(initrd_len, 2);

		ird[i].fd = open(ip->str, O_RDONLY);
		if (ird[i].fd < 0) {
			fprintf(stderr, "%s: %s: %s", program, ip->str,
				strerror(errno));
			rv = EX_NOINPUT;
			goto err;
		}

		ird[i].seg.data = mapfile(ird[i].fd, &ird[i].seg.length, 0);
		if (!ird[i].seg.data) {
			fprintf(stderr, "%s: %s: %s", program, ip->str,
				strerror(errno));
			rv = EX_NOINPUT;
			goto err;
		}

		initrd_len += ird[i].seg.length;
	}

	/* Segment: relocation code */
	/* We put this immediately after the kernel setup, in the memory
	   which will be reclaimed for setup. */
	srel.next = &ssup;
	srel.address = (ssup.address + setup_len + 15) & ~15;
	srel.align = 4;		/* 2**4 = 16 bytes */
	srel.length = reloc_size;
	srel.sh_type = SHT_PROGBITS;
	srel.sh_flags = SHF_ALLOC | SHF_WRITE | SHF_EXECINSTR;
	srel.name = "reloc";
	srel.data = reloc;

	/* Segment: Linux kernel setup */
	ssup.next = &scmd;
	ssup.align = 4;		/* 2**4 = 16 bytes */
	ssup.length = setup_len;
	ssup.sh_type = SHT_PROGBITS;
	ssup.sh_flags = SHF_ALLOC | SHF_WRITE | SHF_EXECINSTR;
	ssup.name = "setup";
	ssup.data = kernel;
	if (setup_ver >= 0x200)
		hdr->type_of_loader = 0xff;	/* "Other modern loader" */

	/* Segment: kernel command line */
	scmd.next = &skrn;
	scmd.length = strlen(cmdline) + 1;
	scmd.address = (ssup.address + setup_space - scmd.length) & ~15;
	scmd.align = 4;		/* 2**4 = 16 bytes */
	if (srel.address + reloc_size > scmd.address) {
		/* Uh-oh, we're short on space... push the command line
		   higher. */
		scmd.address = (srel.address + reloc_size + 15) & ~15;
	}
	scmd.sh_type = SHT_PROGBITS;
	scmd.sh_flags = SHF_ALLOC;
	scmd.name = "cmdline";
	scmd.data = cmdline;
	if (setup_ver >= 0x202) {
		wrle32(scmd.address, &hdr->cmd_line_ptr);
	} else {
		/* Old-style command line protocol */
		wrle16(OLD_CMDLINE_MAGIC, (uint16_t *) (kernel + 0x20));
		wrle16(scmd.address - ssup.address,
		       (uint16_t *) (kernel + 0x22));
	}
	if (setup_ver >= 0x201) {
		wrle16(scmd.address - ssup.address - 0x200, &hdr->heap_end_ptr);
		hdr->loadflags |= CAN_USE_HEAP;
	}


	/* Setup information */
	wrle32(ssup.address, &info->setup_addr);
	wrle32(scmd.address, &info->cmdline_addr);
	entry = srel.address + sizeof *info;

	/* Segment: Linux kernel proper */
	skrn.next = ninitrd ? &ird[0].seg : NULL;
	skrn.align = 4;		/* 2**4 = 16 bytes */
	skrn.length = kernel_len - setup_len;
	skrn.sh_type = SHT_PROGBITS;
	skrn.sh_flags = SHF_ALLOC;
	skrn.name = "kernel";
	skrn.data = kernel + setup_len;

	if (skrn.address < 0x100000)
		initrd_addr = 0x100000;
	else
		initrd_addr = skrn.address + skrn.length;

	/* Loadhigh modifications */
	if (opt.loadhigh) {
		uint32_t rm_base, rm_len, delta;

		if (skrn.address < 0x100000) {
			/* zImage, need to move the kernel as well */
			wrle32(skrn.address, &hmv->mv[1].dst);
			wrle32(skrn.length, &hmv->mv[1].len);
			skrn.address = 0x100000;
			wrle32(skrn.address, &hmv->mv[1].src);
			initrd_addr = skrn.address + skrn.length;
		}

		shmv.next = skrn.next;
		skrn.next = &shmv;
		shmv.address = align_up(initrd_addr, 4);
		shmv.align = 4;
		shmv.length = highmove_size;
		shmv.sh_type = SHT_PROGBITS;
		shmv.sh_flags = SHF_ALLOC | SHF_EXECINSTR;
		shmv.name = "highmove";
		shmv.data = highmove;

		rm_base = ssup.address;
		rm_len = scmd.address + scmd.length - rm_base;
		wrle32(rm_base, &hmv->mv[0].dst);
		wrle32(rm_len, &hmv->mv[0].len);
		wrle32(entry, &hmv->mv_entry);

		entry = shmv.address + sizeof *hmv;

		initrd_addr = shmv.address + shmv.length;
		initrd_addr = align_up(initrd_addr, 4);

		wrle32(initrd_addr, &hmv->mv[0].src);

		delta = initrd_addr - ssup.address;

		ssup.address += delta;
		srel.address += delta;
		scmd.address += delta;

		initrd_addr += rm_len;
	}

	/* Additional segments: initrd */
	for (i = 0; i < ninitrd; i++) {
		char *name;

		initrd_addr = align_up(initrd_addr, 2);

		ird[i].seg.next = (i < ninitrd - 1) ? &ird[i + 1].seg : 0;
		ird[i].seg.address = initrd_addr;
		ird[i].seg.align = 2;	/* 2**2 = 4 bytes */
		ird[i].seg.sh_type = SHT_PROGBITS;
		ird[i].seg.sh_flags = SHF_ALLOC;
		xasprintf(&name, "initrd.%d", i);
		ird[i].seg.name = name;

		initrd_addr += ird[i].seg.length;
	}
	if (setup_ver >= 0x200)
		wrle32(initrd_len, &hdr->ramdisk_size);

	/* Initrd information in the startup info */
	wrle32(ninitrd ? ird[0].seg.address : 0, &info->rd_addr);
	wrle32(initrd_len, &info->rd_len);
	wrle32(initrd_max, &info->rd_maxaddr);

	rv = opt.output(&srel, entry, out);

 err:
	if (ird) {
		for (i = 0; i < ninitrd; i++)
			unmapfile(ird[i].fd, (void *)ird[i].seg.data,
				  ird[i].seg.length);
		free(ird);
	}

	unmapfile(kernel_fd, kernel, kernel_len);
	return rv;
}