text: /* $OpenBSD: rdsetroot.c,v 1.3 2021/10/24 21:24:19 deraadt 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>
  
  /*
   * NOTE: Original from:
   *	https://raw.githubusercontent.com/openbsd/src/master/usr.sbin/rdsetroot/rdsetroot.c
   */
  #define ELF_SYMTAB      ".symtab"       /* symbol table */
  #define ELF_STRTAB      ".strtab"       /* string table */
  #define __dead
  #define pledge(x,y) (1)
  
  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)) < 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);
  	} 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);
  }
  