#include <linux/sched.h>
#include <linux/pid.h>
#include <linux/utrace.h>
#include <linux/err.h>
#include <linux/module.h>
#include <linux/errno.h>

MODULE_DESCRIPTION("automatic suspend on crash");
MODULE_LICENSE("GPL");

static int target_pid;
static int verbose;

module_param_named(pid, target_pid, int, 0);
module_param(verbose, bool, 0);

#define MY_EVENTS (UTRACE_EVENT(CLONE) | UTRACE_EVENT(DEATH) |		\
		   UTRACE_EVENT(SIGNAL_CORE) | UTRACE_EVENT(JCTL) |	\
		   UTRACE_EVENT(QUIESCE))

/*
 * This is the interesting hook.
 */
static u32
crash_suspend_signal(u32 action, struct utrace_engine *engine,
		     struct pt_regs *regs, siginfo_t *info,
		     const struct k_sigaction *orig_ka,
		     struct k_sigaction *return_ka)
{
	enum utrace_signal_action what = utrace_signal_action(action);

	/*
	 * If another engine is doing something, just get out of the way.
	 * If we return to user, clear the record of a previous crash.
	 */
	if (what != UTRACE_SIGNAL_CORE) {
		if (!signal_pending(current))
			engine->data = NULL;
		return what | UTRACE_RESUME;
	}

	/*
	 * If we've been here before and not been back to user mode yet,
	 * then let it go on to the normal core dump.
	 */
	if (engine->data)
		return what | UTRACE_RESUME;

	engine->data = (void *) 1L;
	return UTRACE_RESUME | UTRACE_SIGNAL_TSTP | UTRACE_SIGNAL_HOLD;
}

static u32
crash_suspend_quiesce(u32 action, struct utrace_engine *engine,
		      unsigned long event)
{
	/*
	 * If we return to user, clear the record of a previous crash.
	 */
	if (!event)
		engine->data = NULL;
	return UTRACE_RESUME;
}

static u32
crash_suspend_jctl(u32 action, struct utrace_engine *engine,
		   int type, int notify)
{
	if (type == CLD_STOPPED) {
		int signr = current->exit_code;
		pid_t pgid = task_pgrp_nr(current);
		if (verbose)
			printk("crash-suspend stopped"
			       " pgrp %d for pid %d signal %d\n",
			       pgid, current->pid, signr);
		if (signr != SIGSTOP && signr != SIGTSTP
		    && signr != SIGTTOU && signr != SIGTTIN)
			/*
			 * This is an unnatural stop induced by us, above.
			 * Now that we have ourselves stopped with the
			 * proper weirdo status, stop the rest of the
			 * process group too in normal job control fashion.
			 */
			(void) kill_pgrp(find_vpid(-pgid), SIGTTOU, 1);
	} else if (engine->data) {
		/*
		 * We've been resumed after a crash.
		 * Make sure we get a report before returning to user mode.
		 */
		return UTRACE_REPORT;
	}

	return UTRACE_RESUME;
}

/*
 * On clone, attach to the child.
 */
static u32
crash_suspend_clone(u32 action, struct utrace_engine *engine,
		    unsigned long clone_flags,
		    struct task_struct *child)
{
	struct utrace_engine *child_engine;

	child_engine = utrace_attach_task(child, UTRACE_ATTACH_CREATE,
					  engine->ops, 0);
	if (IS_ERR(child_engine)) {
		printk("attach to clone child %d (%lx) from 0x%p => %ld\n",
		       child->pid, clone_flags, engine, PTR_ERR(child_engine));
	} else {
		int err = utrace_set_events(child, child_engine, MY_EVENTS);
		WARN_ON(err);
		utrace_engine_put(child_engine);
	}

	return UTRACE_RESUME;
}

/*
 * If we are still attached at task death, it didn't die by core dump signal.
 * Just detach and let it go.
 */
static u32
crash_suspend_death(struct utrace_engine *engine,
		    bool group_dead, int signal)
{
	return UTRACE_DETACH;
}

static const struct utrace_engine_ops crash_suspend_ops =
{
	.report_clone = crash_suspend_clone,
	.report_death = crash_suspend_death,
	.report_signal = crash_suspend_signal,
	.report_jctl = crash_suspend_jctl,
	.report_quiesce = crash_suspend_quiesce,
};

static int __init init_crash_suspend(void)
{
	struct pid *pid;
	struct utrace_engine *engine;
	int ret;

	pid = find_get_pid(target_pid);
	if (pid == NULL) {
		printk("cannot find PID %d\n", target_pid);
		return -ECHILD;	/* ESRCH is translated strangely by insmod */
	}

	engine = utrace_attach_pid(pid, UTRACE_ATTACH_CREATE,
				   &crash_suspend_ops, 0);
	if (IS_ERR(engine))
		printk("utrace_attach: %ld\n", PTR_ERR(engine));
	else if (engine == NULL)
		printk("utrace_attach => null!\n");
	else
		printk("attached to %d => 0x%p\n", pid_vnr(pid), engine);

	ret = utrace_set_events_pid(pid, engine, MY_EVENTS);
	if (ret == -ESRCH)
		printk("pid %d died during setup\n", pid_vnr(pid));
	else
		WARN_ON(ret);

	put_pid(pid);
	if (engine && !IS_ERR(engine))
		utrace_engine_put(engine);

	return 0;
}

static void __exit exit_crash_suspend(void)
{
	struct task_struct *t;
	struct utrace_engine *engine;
	int n = 0;
	int ret;

restart:
	rcu_read_lock();
	for_each_process(t) {
		engine = utrace_attach_task(t, UTRACE_ATTACH_MATCH_OPS,
					    &crash_suspend_ops, 0);
		if (IS_ERR(engine)) {
			int error = -PTR_ERR(engine);
			if (error != ENOENT)
				printk("!!! utrace_attach returned %d on %d\n",
				       error, t->pid);
			continue;
		}

		ret = utrace_control(t, engine, UTRACE_DETACH);
		if (ret == -EINPROGRESS) {
			/*
			 * It's running our callback, so we have to
			 * synchronize.  We can't keep rcu_read_lock,
			 * so the task pointer might die.  But it's
			 * safe to call utrace_barrier() even with
			 * a stale task pointer, if we have an engine ref.
			 */
			rcu_read_unlock();
			cond_resched();

			ret = utrace_barrier(t, engine);
			WARN_ON(ret && ret != -ESRCH);

			utrace_engine_put(engine);

			++n;
			goto restart;
		}

		WARN_ON(ret && ret != -ESRCH);
		utrace_engine_put(engine);
		++n;
	}
	rcu_read_unlock();

	printk("detached from %d threads, unloading\n", n);
}

module_init(init_crash_suspend);
module_exit(exit_crash_suspend);