/* $OpenBSD: rdsetroot.c,v 1.2 2019/04/16 06:09:52 sunil Exp $ */

/*
 * Copyright (c) 2019 Sunil Nimmagadda <sunil@openbsd.org>
 *
 * Permission to use, copy, modify, and 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 THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR 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.
 */

#include <sys/mman.h>
#include <sys/stat.h>

#include <err.h>
#include <fcntl.h>
#include <gelf.h>
#include <libelf.h>
#include <stdio.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

int find_rd_root_image(uint64_t *, uint64_t *, off_t *, size_t *);
int symbol_get_u64(const char *, uint64_t *);
__dead void usage(void);

Elf	*e;
Elf_Scn	*symtab;
size_t	 nsymb, strtabndx, strtabsz;

int
main(int argc, char **argv)
{
	GElf_Shdr	 shdr;
	Elf_Scn		*scn;
	char		*dataseg, *kernel = NULL, *fs = NULL, *name;
	off_t		 mmap_off, rd_root_size_val;
	size_t		 shstrndx, mmap_size;
	uint64_t	 rd_root_size_off, rd_root_image_off;
	uint32_t	*ip;
	int		 ch, debug = 0, fsfd, kfd, n, xflag = 0;

	while ((ch = getopt(argc, argv, "dx")) != -1) {
		switch (ch) {
		case 'd':
			debug = 1;
			break;
		case 'x':
			xflag = 1;
			break;
		default:
			usage();
		}
	}
	argc -= optind;
	argv += optind;

	if (argc == 1)
		kernel = argv[0];
	else if (argc == 2) {
		kernel = argv[0];
		fs = argv[1];
	} else
		usage();

	if ((kfd = open(kernel, xflag ? O_RDONLY : O_RDWR, 0644)) < 0)
		err(1, "%s", kernel);

	if (fs) {
		if (xflag)
			fsfd = open(fs, O_RDWR | O_CREAT | O_TRUNC, 0644);
		else
			fsfd = open(fs, O_RDONLY, 0644);
	} else {
		if (xflag)
			fsfd = dup(STDOUT_FILENO);
		else
			fsfd = dup(STDIN_FILENO);
	}
	if (fsfd < 0)
		err(1, "%s", fs);

	if (pledge("stdio", NULL) == -1)
		err(1, "pledge");

	if (elf_version(EV_CURRENT) == EV_NONE)
		errx(1, "elf_version: %s", elf_errmsg(-1));

	if ((e = elf_begin(kfd, xflag ? ELF_C_READ : ELF_C_RDWR, NULL)) == NULL)
		errx(1, "elf_begin: %s", elf_errmsg(-1));

	if (elf_kind(e) != ELF_K_ELF)
		errx(1, "%s: not an elf", kernel);

	if (gelf_getclass(e) == ELFCLASSNONE)
		errx(1, "%s: invalid elf, not 32 or 64 bit", kernel);

	/* Retrieve index of section name string table. */
	if (elf_getshdrstrndx(e, &shstrndx) != 0)
		errx(1, "elf_getshdrstrndx: %s", elf_errmsg(-1));

	/* Find symbol table, string table. */
	scn = symtab = NULL;
	while ((scn = elf_nextscn(e, scn)) != NULL) {
		if (gelf_getshdr(scn, &shdr) != &shdr)
			errx(1, "elf_getshdr: %s", elf_errmsg(-1));

		if ((name = elf_strptr(e, shstrndx, shdr.sh_name)) == NULL)
			errx(1, "elf_strptr: %s", elf_errmsg(-1));

		if (strcmp(name, ELF_SYMTAB) == 0 &&
		    shdr.sh_type == SHT_SYMTAB && shdr.sh_entsize != 0) {
			symtab = scn;
			nsymb = shdr.sh_size / shdr.sh_entsize;
		}

		if (strcmp(name, ELF_STRTAB) == 0 &&
		    shdr.sh_type == SHT_STRTAB) {
			strtabndx = elf_ndxscn(scn);
			strtabsz = shdr.sh_size;
		}
	}

	if (symtab == NULL)
		errx(1, "symbol table not found");

	if (strtabndx == 0)
		errx(1, "string table not found");

	if (find_rd_root_image(&rd_root_size_off, &rd_root_image_off,
	    &mmap_off, &mmap_size) != 0)
		errx(1, "can't locate space for rd_root_image!");

	if (debug) {
		fprintf(stderr, "rd_root_size_off: 0x%llx\n", rd_root_size_off);
		fprintf(stderr, "rd_root_image_off: 0x%llx\n",
		    rd_root_image_off);
	}

	/*
	 * Map in the whole data segment.
	 * The file offset needs to be page aligned.
	 */
	dataseg = mmap(NULL, mmap_size,
	    xflag ? PROT_READ : PROT_READ | PROT_WRITE,
	    MAP_SHARED, kfd, mmap_off);
	if (dataseg == MAP_FAILED)
		err(1, "%s: cannot map data seg", kernel);

	/*
	 * Find value in the location: rd_root_size
	 */
	ip = (uint32_t *) (dataseg + rd_root_size_off);
	rd_root_size_val = *ip;
	if (debug) {
		fprintf(stderr, "rd_root_size  val: 0x%llx (%lld blocks)\n",
		    (unsigned long long)rd_root_size_val,
		    (unsigned long long)rd_root_size_val >> 9);
		fprintf(stderr, "copying root image...\n");
	}

	if (xflag) {
		n = write(fsfd, dataseg + rd_root_image_off,
		    (size_t)rd_root_size_val);
		if (n != rd_root_size_val)
			err(1, "write");
	} else {
		struct stat sstat;

		if (fstat(fsfd, &sstat) == -1)
			err(1, "fstat");
		if (S_ISREG(sstat.st_mode) &&
		    sstat.st_size > rd_root_size_val) {
			fprintf(stderr, "ramdisk too small 0x%llx 0x%llx\n",
			    (unsigned long long)sstat.st_size,
			    (unsigned long long)rd_root_size_val);
			exit(1);
		}
		n = read(fsfd, dataseg + rd_root_image_off,
		    (size_t)rd_root_size_val);
		if (n < 0)
			err(1, "read");

		msync(dataseg, mmap_size, 0);
	}

	if (debug)
		fprintf(stderr, "...copied %d bytes\n", n);

	elf_end(e);
	return 0;
}

int
find_rd_root_image(uint64_t *rd_root_size_off, uint64_t *rd_root_image_off,
    off_t *pmmap_off, size_t *pmmap_size)
{
	GElf_Phdr	phdr;
	size_t		i, phdrnum;
	unsigned long	kernel_start, kernel_size;
	uint64_t	adiff, rd_root_size, rd_root_image, size_off, image_off;
	int		error = 1;

	if (symbol_get_u64("rd_root_size", &rd_root_size) != 0)
		errx(1, "no rd_root_image symbols?");

	if (symbol_get_u64("rd_root_image", &rd_root_image) != 0)
		errx(1, "no rd_root_image symbols?");

	/* Retrieve number of program headers. */
	if (elf_getphdrnum(e, &phdrnum) != 0)
		errx(1, "elf_getphdrnum: %s", elf_errmsg(-1));

	/* Locate the data segment. */
	for (i = 0; i < phdrnum; i++) {
		if (gelf_getphdr(e, i, &phdr) != &phdr)
			errx(1, "gelf_getphdr: %s", elf_errmsg(-1));

		if (phdr.p_type != PT_LOAD)
			continue;

		kernel_start = phdr.p_paddr;
		kernel_size = phdr.p_filesz;
		adiff = phdr.p_vaddr - phdr.p_paddr;

		size_off = rd_root_size - kernel_start;
		image_off = rd_root_image - kernel_start;
		if (size_off < adiff || image_off < adiff)
			continue;

		size_off -= adiff;
		image_off -= adiff;
		if (image_off >= kernel_size)
			continue;
		if (size_off >= kernel_size)
			errx(1, "rd_root_size not in data segment");

		*pmmap_off = phdr.p_offset;
		*pmmap_size = kernel_size;
		*rd_root_size_off = size_off;
		*rd_root_image_off = image_off;
		error = 0;
		break;
	}

	return error;
}

int
symbol_get_u64(const char *symbol, uint64_t *result)
{
	GElf_Sym	 sym;
	Elf_Data	*data;
	const char	*name;
	size_t		 i;
	int		 error = 1;

	data = NULL;
	while ((data = elf_rawdata(symtab, data)) != NULL) {
		for (i = 0; i < nsymb; i++) {
			if (gelf_getsym(data, i, &sym) != &sym)
				continue;

			if (sym.st_name >= strtabsz)
				break;

			if ((name = elf_strptr(e, strtabndx,
			    sym.st_name)) == NULL)
				continue;

			if (strcmp(name, symbol) == 0) {
				if (result)
					*result = sym.st_value;
				error = 0;
				break;
			}
		}
	}

	return error;
}

__dead void
usage(void)
{
	extern char *__progname;

	fprintf(stderr, "usage: %s [-dx] bsd [fs]\n", __progname);
	exit(1);
}