/*
 * Kernel core dump.
 *
 * 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; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program 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 General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
 *
 * Copyright (C) Hui Zhu (teawater@gmail.com), 2009
 *
 * 2002-09-22	Demo release.  Just support X86_32.
 */

#include <linux/version.h>
#include <linux/kernel.h>
#include <linux/module.h>
#include <linux/kprobes.h>
#include <linux/elfcore.h>
#include <linux/init.h>
#include <linux/uaccess.h>
#include <linux/proc_fs.h>
#include <linux/seq_file.h>

#ifndef CONFIG_KPROBES
#error "Linux Kernel doesn't support KPROBES.  Please open it in 'General setup->Kprobes'."
#endif

#define KNAME_NOT_SET		"NOT SET"
#define ELFNOTELEN(datalen)	(sizeof(struct elf_note) + roundup(sizeof ("CORE"), 4) + roundup(datalen, 4))
#define STACKLEN		512
#define KBUFLEN			(sizeof(struct elfhdr) + sizeof(struct elf_phdr) * 2 + ELFNOTELEN(sizeof(struct elf_prstatus)) + ELFNOTELEN(sizeof(struct elf_prpsinfo)) + STACKLEN)

static char			*kname = KNAME_NOT_SET;
static unsigned int		koffset = 0;
static unsigned int		kaddr = 0;
static struct proc_dir_entry	*kcoredump_file = NULL;

static uint8_t			kbuf[KBUFLEN];
static struct elf_prstatus	*kpstat = NULL;
static struct elf_phdr		*kphdr = NULL;
static uint8_t			*kstack;

static int
handler_pre(struct kprobe *p, struct pt_regs *regs)
{
	static int dumped = 0;
	static DEFINE_SPINLOCK(dumped_lock);

	spin_lock(&dumped_lock);
	if (dumped) {
		spin_unlock(&dumped_lock);
		return 0;
	}
	dumped = 1;
	spin_unlock(&dumped_lock);

#ifdef CONFIG_X86_32
	/* Set regs to kpstat->pr_reg. */
#if (LINUX_VERSION_CODE > KERNEL_VERSION(2,6,24))
	kpstat->pr_reg[0] = regs->bx;
	kpstat->pr_reg[1] = regs->cx;
	kpstat->pr_reg[2] = regs->dx;
	kpstat->pr_reg[3] = regs->si;
	kpstat->pr_reg[4] = regs->di;
	kpstat->pr_reg[5] = regs->bp;
	kpstat->pr_reg[6] = regs->ax;
	kpstat->pr_reg[7] = regs->ds;
	kpstat->pr_reg[8] = regs->es;
	kpstat->pr_reg[9] = regs->fs;
#if (LINUX_VERSION_CODE > KERNEL_VERSION(2,6,29))
	kpstat->pr_reg[10] = regs->gs;
#else
	kpstat->pr_reg[10] = 0;
#endif
	kpstat->pr_reg[11] = regs->orig_ax;
	kpstat->pr_reg[12] = regs->ip - 1;
	kpstat->pr_reg[13] = regs->cs;
	kpstat->pr_reg[14] = regs->flags;
	kpstat->pr_reg[15] = (unsigned long)&regs->sp;
	kpstat->pr_reg[16] = 0;
#else
	kpstat->pr_reg[0] = regs->ebx;
	kpstat->pr_reg[1] = regs->ecx;
	kpstat->pr_reg[2] = regs->edx;
	kpstat->pr_reg[3] = regs->esi;
	kpstat->pr_reg[4] = regs->edi;
	kpstat->pr_reg[5] = regs->ebp;
	kpstat->pr_reg[6] = regs->eax;
	kpstat->pr_reg[7] = regs->xds;
	kpstat->pr_reg[8] = regs->xes;
	kpstat->pr_reg[9] = regs->xfs;
	/* kpstat->pr_reg[10] = regs->gs; */
	kpstat->pr_reg[11] = regs->orig_eax;
	kpstat->pr_reg[12] = regs->eip - 1;
	kpstat->pr_reg[13] = regs->xcs;
	kpstat->pr_reg[14] = regs->eflags;
	kpstat->pr_reg[15] = (unsigned long)&regs->esp;
	kpstat->pr_reg[16] = 0;
#endif
	/* Set kpstat->pr_reg[15] to kphdr->p_vaddr. */
	kphdr->p_vaddr = kpstat->pr_reg[15] & ~63;

#if (LINUX_VERSION_CODE > KERNEL_VERSION(2,6,25))
	if (probe_kernel_read (kstack, (void *)kphdr->p_vaddr, STACKLEN))
		memset (kstack, 0, STACKLEN);
#else
	memcpy (kstack, (void *)kphdr->p_vaddr, STACKLEN);
#endif

	return 0;
#else
#error "This architecture doesn't support kcoredump."
#endif
}

static void
handler_post(struct kprobe *p, struct pt_regs *regs,
				unsigned long flags)
{
}

static int
handler_fault(struct kprobe *p, struct pt_regs *regs, int trapnr)
{
	printk(KERN_ERR "kcoredump: get fault addr = 0x%p, trap #%dn\n",
		p->addr, trapnr);
	return 0;
}

static struct kprobe kp = {
	.pre_handler	= handler_pre,
	.post_handler	= handler_post,
	.fault_handler	= handler_fault,
};

static int
fill_elf_note (uint8_t *p, const char *name, int type, void **data, int data_len)
{
	int		ret = 0;
	struct elf_note	*en = (struct elf_note *) p;

	en = (struct elf_note *) p;
	p += sizeof(struct elf_note);
	ret += sizeof(struct elf_note);
	en->n_namesz = strlen(name) + 1;
	en->n_type = type;
	en->n_descsz = data_len;

	memcpy (p, name, en->n_namesz);
	p += en->n_namesz;
	p = (uint8_t *) roundup((unsigned long) p, 4);
	ret += roundup((en->n_namesz), 4);

	*data = p;
	p += data_len;
	p = (uint8_t *) roundup((unsigned long) p, 4);
	ret += roundup(data_len, 4);

	return ret;
}

static int
init_kbuf (void)
{
	uint8_t			*wp = kbuf;
	struct elfhdr		*elf;
	struct elf_phdr		*nhdr;
	struct elf_prpsinfo	*pinfo;

	/* elf head */
	elf = (struct elfhdr *) wp;
	wp += sizeof(struct elfhdr);
	memset(elf, 0, sizeof(*elf));
	memcpy(elf->e_ident, ELFMAG, SELFMAG);
	elf->e_ident[EI_CLASS] = ELF_CLASS;
	elf->e_ident[EI_DATA] = ELF_DATA;
	elf->e_ident[EI_VERSION] = EV_CURRENT;
	elf->e_ident[EI_OSABI] = ELF_OSABI;
	elf->e_type = ET_CORE;
	elf->e_machine = ELF_ARCH;
	elf->e_version = EV_CURRENT;
	elf->e_phoff = sizeof(struct elfhdr);
	elf->e_flags = 0;
	elf->e_ehsize = sizeof(struct elfhdr);
	elf->e_phentsize = sizeof(struct elf_phdr);
	/* sections number (vma + 1) */
	elf->e_phnum = 1 + 1;

	/* nhdr */
	nhdr = (struct elf_phdr *) wp;
	wp += sizeof(struct elf_phdr);
	nhdr->p_type	= PT_NOTE;
	nhdr->p_offset	= sizeof(struct elfhdr) + sizeof(struct elf_phdr) * 2;
	nhdr->p_vaddr	= 0;
	nhdr->p_paddr	= 0;
	nhdr->p_filesz	= 0;
	nhdr->p_memsz	= 0;
	nhdr->p_flags	= 0;
	nhdr->p_align	= 0;

	/* kphdr */
	kphdr = (struct elf_phdr *) wp;
	wp += sizeof(struct elf_phdr);
	kphdr->p_type = PT_LOAD;
	kphdr->p_flags = PF_R;
	kphdr->p_offset = nhdr->p_offset;
	/* kphdr->p_vaddr will be set in handler_pre.  */
	kphdr->p_vaddr = 0;
	kphdr->p_paddr = 0;
	kphdr->p_filesz = kphdr->p_memsz = STACKLEN;
	kphdr->p_align = 1;

	/* note0 kpstat */
	if (fill_elf_note (wp, "CORE", NT_PRSTATUS, ((void **)&kpstat),
				sizeof(struct elf_prstatus)) != ELFNOTELEN(sizeof(struct elf_prstatus)))
		return -1;
	wp += ELFNOTELEN(sizeof(struct elf_prstatus));
	nhdr->p_filesz += ELFNOTELEN(sizeof(struct elf_prstatus));
	/* kpstat->pr_reg will be set in handler_pre.  */
	memset (kpstat, 0, sizeof(struct elf_prstatus));

	/* note1 pinfo */
	if (fill_elf_note (wp, "CORE", NT_PRPSINFO, ((void **)&pinfo),
				sizeof(struct elf_prpsinfo)) != ELFNOTELEN(sizeof(struct elf_prpsinfo)))
		return -1;
	wp += ELFNOTELEN(sizeof(struct elf_prpsinfo));
	nhdr->p_filesz += ELFNOTELEN(sizeof(struct elf_prpsinfo));
	memset (pinfo, 0, sizeof(struct elf_prpsinfo));
	pinfo->pr_state = 0;
	pinfo->pr_sname = 'R';
	pinfo->pr_zomb = 0;
	snprintf(pinfo->pr_fname, 16, "%s", "vmlinux");
	//snprintf(pinfo->pr_psargs, ELF_PRARGSZ, "%s", saved_command_line);
	snprintf(pinfo->pr_psargs, ELF_PRARGSZ, "%s", "");

	/* kstack */
	kphdr->p_offset += nhdr->p_filesz;
	kstack = wp;
	memset (kstack, 0, STACKLEN);
	wp += STACKLEN;

	/* Check size.  */
	if (wp - kbuf != KBUFLEN)
		return -1;

	return 0;
}

#if (LINUX_VERSION_CODE < KERNEL_VERSION(2,6,31))
static int
seq_write(struct seq_file *seq, const void *data, size_t len)
{
	if (seq->count + len < seq->size) {
		memcpy(seq->buf + seq->count, data, len);
		seq->count += len;
		return 0;
	}
	seq->count = seq->size;
	return -1;
}
#endif

static int
kcoredump_seq_show(struct seq_file *s, void *v)
{
	return seq_write (s, kbuf, KBUFLEN);
}

static int
kcoredump_proc_open(struct inode *inode, struct file *file)
{
	return single_open(file, kcoredump_seq_show, NULL);
}

static const struct file_operations kcoredump_proc_ops = {
	.owner		= THIS_MODULE,
	.open		= kcoredump_proc_open,
	.read		= seq_read,
	.llseek		= seq_lseek,
	.release	= seq_release,
};

#if (LINUX_VERSION_CODE < KERNEL_VERSION(2,6,25))
static inline struct proc_dir_entry *
proc_create(const char *name, mode_t mode, struct proc_dir_entry *parent,
		const struct file_operations *proc_fops)
{
	struct proc_dir_entry *ret = create_proc_entry(name, mode, parent);
	if (!ret)
		return NULL;

	ret->proc_fops = proc_fops;

	return ret;
}
#endif

int init_module(void)
{
	int	ret = 0;

	/* Check argument.  */
	if (strcmp (kname, KNAME_NOT_SET)) {
		kp.symbol_name = kname;
		kp.offset = koffset;
	}
	else {
		if (kaddr) {
			kp.addr = (kprobe_opcode_t *)kaddr;
		}
		else {
			printk (KERN_ERR "kcoredump: must set the name or addr.\n");
			return -EINVAL;
		}
	}

	if (init_kbuf ()) {
		printk (KERN_ERR "kcoredump: init buf error.\n");
		return -EIO;
	}

	kcoredump_file = proc_create ("kcoredump", 0444, NULL, &kcoredump_proc_ops);
	if (!kcoredump_file) {
		printk (KERN_ERR "kcoredump: create proc file error.\n");
		return -ENOMEM;
	}
	/* kcoredump_file->owner = THIS_MODULE; */
	kcoredump_file->size = KBUFLEN;

	ret = register_kprobe(&kp);
	if (ret < 0) {
		printk (KERN_ERR "kcoredump: register kprobe error.\n");
		remove_proc_entry("kcoredump", NULL);
		return ret;
	}

	printk(KERN_NOTICE "kcoredump: dump in 0x%p.\n", kp.addr);

	return 0;
}

void cleanup_module(void)
{
	unregister_kprobe(&kp);
	remove_proc_entry("kcoredump", NULL);
}

module_param_named(name, kname, charp, 0644);
module_param_named(offset, koffset, uint, 0644);
module_param_named(addr, kaddr, uint, 0644);

MODULE_AUTHOR("Hui Zhu");
MODULE_LICENSE("GPL");