/*
 * radix-tree-test.c
 *
 * This module inserts a kprobe into the radix_tree_node_alloc function that
 * simulates a failed kmem_cache_alloc.  It then performs an insert at an
 * index that will require allocating the entire path to the maximum height
 * of the tree.  This ensures that the radix_tree_preload array is large
 * enough to satisfy this pathological case.
 */

#include <linux/kernel.h>
#include <linux/kallsyms.h>
#include <linux/kprobes.h>
#include <linux/module.h>
#include <linux/radix-tree.h>

/*
 * Change this to the location of the test for (ret == NULL in
 * radix_tree_node_alloc.
 */
#define RET_GETS_NULL_CHECK  0xffffffff811131ee
/*
 * Change this to the address of the radix_tree_node_cachep for your kernel.
 */
#define RADIX_TREE_CACHEP    0xffff810037c20080

RADIX_TREE(test_tree, GFP_ATOMIC);
unsigned long tree_index = ~0UL - 1;

static int probe_radix_tree_node_alloc(struct kprobe *p, struct pt_regs *regs)
{
	/*
	 * regs->rax points to the return value from kmem_cache_alloc.
	 */
	kmem_cache_free((struct kmem_cache *)RADIX_TREE_CACHEP,
			(struct radix_tree_node *)regs->rax);
	regs->rax = 0;
	return 0;
}

struct kprobe rtna_probe = {
	.addr = (kprobe_opcode_t *)RET_GETS_NULL_CHECK,
	.pre_handler = probe_radix_tree_node_alloc,
};

int init_module(void)
{
	int ret;

	ret = register_kprobe(&rtna_probe);
	if (ret != 0) {
		printk("Failed to register kprobe!\n");
		return ret;
	}

	ret = radix_tree_preload(GFP_KERNEL);
	if (ret < 0)
		goto out;

	preempt_disable();
	ret = radix_tree_insert(&test_tree, tree_index, (void *)0xdead0000);
	if (ret != 0)
		printk("Failed to insert element into tree!\n");
	preempt_enable();

	printk("Successfully inserted a radix tree member at maximum height"
	       " with preemption disabled and a failed kmalloc.\n");
out:
	unregister_kprobe(&rtna_probe);
	return ret;
}

void cleanup_module(void)
{
	radix_tree_delete(&test_tree, tree_index);
	return;
}

MODULE_LICENSE("GPL");