/*
 * corefilter: truncate a core dump file to a certain size while preserving
 *             data useful to the debugger
 *
 * Copyright (C) Jean-Marc Saffroy <saffroy@gmail.com> 2008
 * This program is free software, distributed under the terms of the
 * GNU General Public License version 2.
 *
 * The core dump is read from stdin. The program is meant to be invoked by
 * the kernel on a core dump, after being installed like this:
 *
 * # echo '|/bin/corefilter -c %c -p %p -g %g -u %u' > \
 *       /proc/sys/kernel/core_pattern
 */

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <errno.h>
#include <string.h>
#include <fcntl.h>
#include <elf.h>
#include <sys/fsuid.h>
#include <limits.h>

#define min(x,y) (x < y ? x : y)
#define BUFFER_SIZE (1 << 20)

#ifdef _LP64

typedef Elf64_Ehdr	elf_ehdr_t;
typedef Elf64_Phdr	elf_phdr_t;

#define EXPECTED_ELF_CLASS ELFCLASS64

#else /* !USE_ELF64 */

typedef Elf32_Ehdr	elf_ehdr_t;
typedef Elf32_Phdr	elf_phdr_t;

#define EXPECTED_ELF_CLASS ELFCLASS32

#endif /* !USE_ELF64 */

void usage(char *progname)
{
	printf("usage: %s <options>\n"
	       "	-p <pid>\n"
	       "	-u <uid>\n"
	       "	-g <gid>\n"
	       "	-c <core file limit (in kB)>\n", progname);
	exit(1);
}

void err(char *msg)
{
	fprintf(stderr, "error: %s\n", msg);
	exit(1);
}

void errno_exit(char *msg)
{
	fprintf(stderr, "error: %s: %s\n", msg, strerror(errno));
	exit(1);
}

void read_pipe(void *buf, int bytes, off_t *offset)
{
	int num_read = 0;

	while (num_read < bytes) {
		int rc = read(0, buf, bytes - num_read);

		if ((rc < 0) && (errno != EINTR))
			errno_exit("read");
		if (rc > 0)
			num_read += rc;
	}

	*offset += bytes;
}

void seek_pipe(off_t new_offset, off_t *offset)
{
	if (new_offset < *offset)
		errno_exit("try to seek before current pipe offset");

	while (*offset < new_offset) {
		char buf[BUFFER_SIZE];

		read_pipe(buf, min(sizeof(buf), (new_offset - *offset)),
			  offset);
	}
}

void seek_file(int fd, off_t new_offset)
{
	if (lseek(fd, new_offset, SEEK_SET) < 0)
		errno_exit("lseek");
}

void write_file(int fd, void *buf, int bytes)
{
	int num_written = 0;

	while (num_written < bytes) {
		int rc = write(fd, buf, bytes - num_written);

		if ((rc < 0) && (errno != EINTR))
			errno_exit("write");
		if (rc > 0)
			num_written += rc;
	}
}

struct segs {
	int seg_num;		/* segment index in array of program headers */
	long seg_sz;		/* segment file size */
};

int cmp_segs(const void *a, const void *b)
{
	const struct segs *sa = a;
	const struct segs *sb = b;

	return sa->seg_sz - sb->seg_sz;
}

void trim_core(elf_ehdr_t *hdrin, elf_ehdr_t *hdrout,
	       elf_phdr_t *phdrin, elf_phdr_t *phdrout,
	       unsigned long limit)
{
	int num_load_segs, num_other_segs, i;
	struct segs load_segs[hdrin->e_phnum];
	unsigned long offset;

	/* copy ELF header */
	memcpy(hdrout, hdrin, sizeof(*hdrin));
	/* copy program headers */
	memcpy(phdrout, phdrin, hdrin->e_phentsize * hdrin->e_phnum);

#if 1
	/*
	 * Real memory data is in PT_LOAD segments.
	 * These segments are sorted by file size, and
	 * trimmed according to limit.
	 */
	num_load_segs = num_other_segs = 0;
	for (i = 0; i < hdrin->e_phnum; i++) {
		if (phdrin[i].p_type == PT_LOAD) {
			load_segs[num_load_segs].seg_num = i;
			load_segs[num_load_segs].seg_sz =
			    phdrin[i].p_filesz;
			num_load_segs++;
		}
	}

	/* sort segments by file size */
	qsort(load_segs, num_load_segs, sizeof(load_segs[0]), cmp_segs);

	/* trim segments  */
	for (i = 0; i < num_load_segs; i++) {
		unsigned long avg_size, seg_size;
		elf_phdr_t *phdr;

		phdr = &phdrout[load_segs[i].seg_num];
		avg_size = limit / (num_load_segs - i);
		seg_size = phdr->p_filesz;
		if (seg_size > avg_size)
			seg_size = avg_size;
		if (phdr->p_align)
			seg_size &= ~(phdr->p_align - 1);
		phdr->p_filesz = seg_size;
		limit -= seg_size;
	}

	/* fix file offsets */
	offset = phdrin[0].p_offset;
	for (i = 0; i < hdrout->e_phnum; i++) {
		if (phdrout[i].p_align) {
			offset += phdrout[i].p_align - 1;
			offset &= ~(phdrout[i].p_align - 1);
		}
		phdrout[i].p_offset = offset;
		offset += phdrout[i].p_filesz;
	}
#endif
}

int main(int argc, char **argv)
{
	pid_t pid = 0;
	uid_t uid = 0;
	gid_t gid = 0;
	unsigned long limit = -1;
	char path[PATH_MAX];
	int c, fd, i;
	off_t offset = 0;
	elf_ehdr_t hdrin, hdrout;
	elf_phdr_t *phdrin, *phdrout;
	int debug = 0;

	/* parse args */
	while ((c = getopt(argc, argv, "u:g:p:c:d")) != -1) {
		switch (c) {
		case 'u':
			uid = atoi(optarg);
			break;
		case 'g':
			gid = atoi(optarg);
			break;
		case 'p':
			pid = atoi(optarg);
			break;
		case 'c':
			limit = atol(optarg);
			break;
		case 'd':
			debug = 1;
			break;
		default:
			fprintf(stderr, "error: invalid flag -%c\n", c);
			usage(argv[0]);
		}
	}

	/* anything to do? */
	if (limit == 0)
		exit(0);

	/* setup stderr for debugging */
	if (debug && !isatty(2)) {
		fflush(stderr);
		fclose(stderr);
		if (pid > 0) {
			snprintf(path, sizeof(path),
				 "/tmp/corefilter.debug.%d", pid);
			stderr = fopen(path, "w");
		} else {
			snprintf(path, sizeof(path),
				 "/tmp/corefilter.debug");
			stderr = fopen(path, "a");
		}
	}
	if (debug) {
		for (i = 0; i < argc; i++)
			fprintf(stderr, "argv[%d] = '%s'\n", i, argv[i]);
	}

	/* set cwd */
	if (pid > 0) {
		snprintf(path, sizeof(path), "/proc/%d/cwd", pid);
		if (debug) {
			char buf[PATH_MAX];

			if (readlink(path, buf, sizeof(buf)) < 0)
				errno_exit("readlink");
			fprintf(stderr, "cwd: %s -> %s\n", path, buf);
		}
		if (chdir(path))
			errno_exit("chdir");
	}

	/* set credentials */
	if (setgid(gid))
		errno_exit("setgid");
	if (setuid(uid))
		errno_exit("setuid");

	/* open output file for writing */
	if (pid > 0)
		snprintf(path, sizeof(path), "core.%d", pid);
	else
		snprintf(path, sizeof(path), "core");
	fd = creat(path, 0600);
	if (fd < 0)
		errno_exit("creat");

	/* read ELF header, perform basic sanity checks */
	read_pipe(&hdrin, sizeof(hdrin), &offset);
	if(memcmp(hdrin.e_ident, ELFMAG, SELFMAG) != 0)
		err("bad ELF magic");
	if(hdrin.e_ident[EI_CLASS] != EXPECTED_ELF_CLASS)
		err("bad ELF class");
	if(hdrin.e_type != ET_CORE)
		err("bad ELF type");
	

	/* read program headers */
	phdrin = malloc(hdrin.e_phentsize * hdrin.e_phnum);
	if (!phdrin)
		errno_exit("malloc");
	seek_pipe(hdrin.e_phoff, &offset);
	read_pipe(phdrin, hdrin.e_phentsize * hdrin.e_phnum, &offset);

	/* recompute layout */
	phdrout = malloc(hdrin.e_phentsize * hdrin.e_phnum);
	if (!phdrout)
		errno_exit("malloc");
	trim_core(&hdrin, &hdrout, phdrin, phdrout, limit);
	if (debug && stderr) {
		fprintf(stderr, "%-16s %9s %9s\n",
			"segment", "old size", "new size");
		for (i = 0; i < hdrin.e_phnum; i++)
			fprintf(stderr, "%016lx % 9ld % 9ld\n",
				phdrin[i].p_vaddr,
				phdrin[i].p_filesz, phdrout[i].p_filesz);
	}

	/* write ELF header */
	write_file(fd, &hdrout, sizeof(hdrout));

	/* write modified program headers */
	seek_file(fd, hdrout.e_phoff);
	write_file(fd, phdrout, hdrout.e_phentsize * hdrout.e_phnum);

	/* read each segment, write none/some/all of its content */
	for (i = 0; i < hdrin.e_phnum; i++) {
		int towrite;

		seek_pipe(phdrin[i].p_offset, &offset);
		seek_file(fd, phdrout[i].p_offset);
		towrite = phdrout[i].p_filesz;

		while (towrite > 0) {
			char buf[BUFFER_SIZE];
			int bytes;

			bytes = min(sizeof(buf), towrite);
			read_pipe(buf, bytes, &offset);
			write_file(fd, buf, bytes);
			towrite -= bytes;
		}
	}

	/* close output file */
	close(fd);

	/* all done! */
	return 0;
}