/*

 * libhugetlbfs - Easy use of Linux hugepages

 * Copyright (C) 2008 Adam Litke, IBM Corporation.

 *

 * This program is free software; you can redistribute it and/or

 * modify it under the terms of the GNU Lesser General Public License

 * as published by the Free Software Foundation; either version 2.1 of

 * the License, or (at your option) any later version.

 *

 * This library is distributed in the hope that it will be useful, but

 * WITHOUT ANY WARRANTY; without even the implied warranty of

 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU

 * Lesser General Public License for more details.

 *

 * You should have received a copy of the GNU Lesser General Public

 * License along with this library; if not, write to the Free Software

 * Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA

 */

#define _GNU_SOURCE

#include <stdio.h>

#include <stdlib.h>

#include <string.h>

#include <unistd.h>

#include <fcntl.h>

#include <elf.h>

#include <link.h>

#include <getopt.h>

#include <errno.h>

#include <sys/stat.h>

#include <sys/mman.h>

/*

 * Eventually we plan to use the libhugetlbfs reporting facility,

 * but until that is possible, redefine a simpler version here.

 */

#define REPORT(level, prefix, format, ...) \

	do { \

		fprintf(stderr, "hugeedit: " prefix ": " format, \

			##__VA_ARGS__); \

	} while (0)

#include "libhugetlbfs_internal.h"

/*

 * All MAP_* options are tagged with MAP_BASE to differentiate them as options

 * in the options parser.  This must be removed before they are compared.

 */

#define MAP_BASE	0x1000

#define MAP_DISABLE	0x0001

#define MAP_TEXT	0x0002

#define MAP_DATA	0x0004

#define PF_LINUX_HUGETLB	0x100000

extern int optind;

extern char *optarg;

#define OPTION(opts, text)	fprintf(stderr, " %-25s  %s\n", opts, text)

#define CONT(text)		fprintf(stderr, " %-25s  %s\n", "", text)

void print_usage()

{

	fprintf(stderr, "hugeedit [options] target\n");

	fprintf(stderr, "options:\n");

	OPTION("--text", "Remap program text into huge pages by default");

	OPTION("--data", "Remap program data into huge pages by default");

	OPTION("--disable", "Remap no segments into huge pages by default");

	OPTION("--help, -h", "Print this usage information");

}

int check_elf_wordsize(void *ehdr)

{

	char *e_ident = (char *) ehdr;

	if (strncmp(e_ident, ELFMAG, SELFMAG)) {

		ERROR("Not a valid ELF executable\n");

		exit(EXIT_FAILURE);

	}

	switch (e_ident[EI_CLASS]) {

		case ELFCLASS32:

		case ELFCLASS64:

			return e_ident[EI_CLASS];

		default:

			ERROR("Can not determine word size\n");

			exit(EXIT_FAILURE);

	}

}

/*

 * We need to map enough of the binary so that we can access the ELF header and

 * all of the program headers.  This function takes a pointer to the first page

 * of ELF headers which is guaranteed to be enough data to determine if we need

 * to map more of the binary.  Use mremap to enlarge the mapping if needed.

 *

 * void **elf - may be updated with a new address if mremap moved it

 * unsigned long *size - may be updated with the new mapping size

 */

#define elf_ph_end_offset(e) ((e)->e_phoff + (e)->e_phentsize * (e)->e_phnum)

void check_remap_elf(void **elf, unsigned long *size, int wordsize)

{

	unsigned long newsize;

	int pagesize = getpagesize();

	if (wordsize == ELFCLASS32) {

		Elf32_Ehdr *ehdr = *(Elf32_Ehdr **) elf;

		newsize = elf_ph_end_offset(ehdr);

	} else {

		Elf64_Ehdr *ehdr = *(Elf64_Ehdr **) elf;

		newsize = elf_ph_end_offset(ehdr);

	}

	newsize = ALIGN_UP(newsize, pagesize);

	if (newsize > *size) {

		*size = newsize;

		*elf = mremap(*elf, *size, newsize, MREMAP_MAYMOVE);

		if (*elf == MAP_FAILED) {

			ERROR("Remapping failed: %s\n", strerror(errno));

			exit(EXIT_FAILURE);

		}

	}

}

#define is_text(p) ((((p)->p_flags & (PF_R|PF_W|PF_X)) == (PF_R|PF_X)) && \

			((p)->p_memsz == (p)->p_filesz))

#if defined(__powerpc__) && !defined(__powerpc64__)

#define is_data(p) (((p)->p_flags & (PF_R|PF_W|PF_X)) == (PF_R|PF_W|PF_X))

#else

#define is_data(p) (((p)->p_flags & (PF_R|PF_W|PF_X)) == (PF_R|PF_W))

#endif

#define update_phdrs(_BITS_) 						\

void update_phdrs##_BITS_(Elf##_BITS_##_Ehdr *ehdr, int remap_opts)	\

{									\

	int i;								\

	Elf##_BITS_##_Phdr *phdr;					\

	unsigned long long start, end;					\

									\

	phdr = (Elf##_BITS_##_Phdr *)((char *)ehdr + ehdr->e_phoff);	\

	for (i = 0; i < ehdr->e_phnum; i++) {				\

		if (phdr[i].p_type != PT_LOAD)				\

			continue;					\

		if (remap_opts)						\

			phdr[i].p_flags &= ~PF_LINUX_HUGETLB;		\

		if ((remap_opts & MAP_TEXT) && is_text(&phdr[i]))	\

			phdr[i].p_flags |= PF_LINUX_HUGETLB;		\

		if ((remap_opts & MAP_DATA) && is_data(&phdr[i]))	\

			phdr[i].p_flags |= PF_LINUX_HUGETLB;		\

		start = (unsigned long long) phdr[i].p_vaddr;		\

		end = start + phdr[i].p_memsz;				\

		printf("Segment %i 0x%llx - 0x%llx (%s%s) default is "	\

			"%s pages\n", i, start, end,			\

			is_text(&phdr[i]) ? "TEXT" : "",		\

			is_data(&phdr[i]) ? "DATA" : "",		\

			(phdr[i].p_flags & PF_LINUX_HUGETLB) ?		\

			"HUGE" : "BASE");				\

	}								\

}

update_phdrs(32)

update_phdrs(64)

int main(int argc, char ** argv)

{

	char opts[] = "+h";

	struct option long_opts[] = {

		{"help",	no_argument, NULL, 'h'},

		{"disable",	no_argument, NULL, MAP_BASE|MAP_DISABLE},

		{"text",	no_argument, NULL, MAP_BASE|MAP_TEXT},

		{"data",	no_argument, NULL, MAP_BASE|MAP_DATA},

		{0},

	};

	int ret = 0, index = 0, remap_opts = 0;

	int fd;

	const char *target;

	void *ehdr;

	unsigned long mapsize = getpagesize();

	int target_wordsize;

	while (ret != -1) {

		ret = getopt_long(argc, argv, opts, long_opts, &index);

		if (ret > 0 && (ret & MAP_BASE)) {

			remap_opts |= ret;

			continue;

		}

		switch (ret) {

		case '?':

			print_usage();

			exit(EXIT_FAILURE);

		case 'h':

			print_usage();

			exit(EXIT_SUCCESS);

		default:

			ret = -1;

			break;

		}

	}

	index = optind;

	remap_opts &= ~MAP_BASE;

	if (remap_opts & MAP_DISABLE && remap_opts != MAP_DISABLE) {

		ERROR("--disable is not compatible with --text or --data\n");

		exit(EXIT_FAILURE);

	}

	if ((argc - index) != 1) {

		print_usage();

		exit(EXIT_FAILURE);

	}

	target = argv[index];

	/* We don't need write access unless we plan to alter the binary */

	fd = open(target, (remap_opts ? O_RDWR : O_RDONLY));

	if (fd < 0) {

		ERROR("Opening %s failed: %s\n", target, strerror(errno));

		exit(EXIT_FAILURE);

	}

	ehdr = mmap(NULL, mapsize, PROT_READ | (remap_opts ? PROT_WRITE : 0),

						MAP_SHARED, fd, 0);

	if (ehdr == MAP_FAILED) {

		ERROR("Mapping %s failed: %s\n", target, strerror(errno));

		exit(EXIT_FAILURE);

	}

	target_wordsize = check_elf_wordsize(ehdr);

	check_remap_elf(&ehdr, &mapsize, target_wordsize);

	if (target_wordsize == ELFCLASS64)

		update_phdrs64((Elf64_Ehdr *) ehdr, remap_opts);

	else

		update_phdrs32((Elf32_Ehdr *) ehdr, remap_opts);

	if (munmap(ehdr, mapsize) != 0) {

		ERROR("Unmapping %s failed: %s\n", target, strerror(errno));

		exit(EXIT_FAILURE);

	}

	if (close(fd) != 0) {

		ERROR("Final close of %s failed: %s -- possible data loss!\n",

			target, strerror(errno));

		exit(EXIT_FAILURE);

	}

	exit(EXIT_SUCCESS);

}