#include <linux/module.h>
#include <linux/init.h>
#include <linux/printk.h>
#include <linux/sched.h>
#include <linux/mmu_context.h>
#include <linux/mm_types.h>
#include <linux/mm.h>
#include <linux/kthread.h>
#include <linux/oom.h>

struct task_struct *th = NULL;

static int read_vma(struct vm_area_struct *vma)
{
	unsigned long addr;
	for (addr = vma->vm_start; addr < vma->vm_end; addr += PAGE_SIZE) {
		char __user *ptr = (char __user *)addr;
		int ret;
		char c;

		if ((ret = get_user(c, ptr)) < 0 && test_bit(MMF_UNSTABLE, &vma->vm_mm->flags)) {
			pr_info("Found a dead vma: ret:%d", ret);
			dump_vma(vma);
			return 1;
		}
	}
	return 0;
}

static int scan_memory(void *data)
{
	struct task_struct *victim = data;
	struct mm_struct *mm;
	int reported = 0;

	pr_info("Starting with pid:%d %s\n", victim->pid, victim->comm);
	mm = get_task_mm(victim);
	if (!mm) {
		pr_info("Failed to get mm\n");
		return 1;
	}
	use_mm(mm);
	mmput(mm);

	while (!kthread_should_stop()) {
		struct vm_area_struct *vma = mm->mmap;
		
		if (!reported && test_bit(MMF_UNSTABLE, &mm->flags)) {
			pr_info("mm is unstable\n");
			reported = 1;
		}
		for (; vma; vma = vma->vm_next) {
			if (!(vma->vm_flags & VM_MAYREAD))
				continue;
			if (read_vma(vma))
				goto out;
		}
		schedule_timeout_idle(HZ);
	}
out:
	unuse_mm(mm);
	while (!kthread_should_stop()) {
		set_current_state(TASK_UNINTERRUPTIBLE);
		if (kthread_should_stop())
			break;
		schedule();
	}
	return 0;
}
static int __init mymodule_init(void)
{
	struct task_struct *p, *victim = NULL;

	rcu_read_lock();
	for_each_process(p) {
		if (p->signal->oom_score_adj == OOM_SCORE_ADJ_MAX) {
			get_task_struct(p);
			victim = p;
			break;
		}
	}
	rcu_read_unlock();
	if (!victim) {
		pr_info("No potential victim found\n");
		return 1;
	}

	th = kthread_run(scan_memory, victim, "scan_memory");
	return 0;
}

static void __exit mymodule_exit(void)
{
	if (!th)
		return;

	pr_info("Waiting for kthread to stop\n");
	kthread_stop(th);
	put_task_struct(th);
	pr_info("Done\n");
}

module_init(mymodule_init);
module_exit(mymodule_exit);

MODULE_LICENSE("GPL");