// SPDX-License-Identifier: GPL-2.0-only
/*
 * Copyright (C) 2025 Ctrl IQ, Inc.
 * Author: Sultan Alsawaf <sultan@ciq.com>
 */

#include <linux/hash.h>
#include <linux/rbtree.h>
#include <linux/seqlock.h>
#include <linux/slab.h>

#include "../../p_lkrg_main.h"
#include "ed_task_tree.h"

/*
 * RCU-safe rbtree lookups were introduced in 4.2 in upstream commit
 * d72da4a4d973 ("rbtree: Make lockless searches non-fatal"). For kernels <4.2,
 * a rwlock is used instead. This check could be improved to cover older kernels
 * with commit d72da4a4d973 backported by checking for the existence of
 * rb_link_node_rcu() in an autoconf-like way, via a stub C file that the build
 * system attempts to compile to determine if rb_link_node_rcu() exists.
 */
#define HAS_RBTREE_RCU (LINUX_VERSION_CODE >= KERNEL_VERSION(4, 2, 0))

#define RB_HASH_BITS 9
#define RB_HASH_SIZE (1U << RB_HASH_BITS)

struct ed_rb_tree {
#if HAS_RBTREE_RCU
	seqcount_t seqcount;
	raw_spinlock_t lock;
#else
	rwlock_t lock;
#endif
	struct rb_root root;
};

/*
 * Align the hash table to a page boundary to ensure the entire hash table spans
 * the fewest TLB entries possible for performance. Since `struct ed_rb_tree` is
 * 16 bytes on common 64-bit architectures and the hash table contains 512
 * entries (RB_HASH_SIZE), this means the hash table occupies exactly two pages'
 * worth of memory when the page size is 4 KiB. As a result, aligning the hash
 * table to a page boundary typically doesn't increase its memory footprint.
 */
static struct ed_rb_tree ed_hash_tbl[RB_HASH_SIZE] __aligned(PAGE_SIZE) = {
	[0 ... RB_HASH_SIZE - 1] = {
#if HAS_RBTREE_RCU
		.seqcount = SEQCNT_ZERO(ed_hash_tbl.seqcount),
		.lock = __RAW_SPIN_LOCK_UNLOCKED(ed_hash_tbl.lock),
#else
		.lock = __RW_LOCK_UNLOCKED(ed_hash_tbl.lock),
#endif
		.root = RB_ROOT
	}
};
static struct kmem_cache *ed_task_cache;
unsigned long cookie_kill_task_on_unlock;

static __always_inline struct ed_rb_tree *
task_rb_tree_bucket(const struct task_struct *tsk)
{
	return &ed_hash_tbl[hash_ptr(tsk, RB_HASH_BITS)];
}

static __always_inline void
__rb_add_or_del(struct p_ed_process *edp, struct ed_rb_tree *edtree,
		struct rb_node *parent, struct rb_node **link)
{
#if HAS_RBTREE_RCU
	write_seqcount_begin(&edtree->seqcount);
#endif
	if (link) {
#if HAS_RBTREE_RCU
		rb_link_node_rcu(&edp->node, parent, link);
#else
		rb_link_node(&edp->node, parent, link);
#endif
		rb_insert_color(&edp->node, &edtree->root);
	} else {
		rb_erase(&edp->node, &edtree->root);
	}
#if HAS_RBTREE_RCU
	write_seqcount_end(&edtree->seqcount);
#endif
}

/* Add @edp to its bucket's red-black tree. Returns 0 on success. */
int ed_task_add(struct p_ed_process *edp)
{
	struct task_struct *tsk = edp->p_ed_task.p_task;
	struct ed_rb_tree *edtree = task_rb_tree_bucket(tsk);
	struct rb_node **link = &edtree->root.rb_node;
	struct rb_node *parent = NULL;
	struct p_ed_process *entry;
	int ret = 0;

#if HAS_RBTREE_RCU
	raw_spin_lock(&edtree->lock);
#else
	write_lock(&edtree->lock);
#endif
	while (*link) {
		parent = *link;
		entry = rb_entry(parent, typeof(*entry), node);
		if (tsk < entry->p_ed_task.p_task) {
			link = &(*link)->rb_left;
		} else if (tsk > entry->p_ed_task.p_task) {
			link = &(*link)->rb_right;
		} else {
			/* Task is already in the tree, bail out */
			ret = -EEXIST;
			goto unlock;
		}
	}
	__rb_add_or_del(edp, edtree, parent, link);
unlock:
#if HAS_RBTREE_RCU
	raw_spin_unlock(&edtree->lock);
#else
	write_unlock(&edtree->lock);
#endif

	return ret;
}

/*
 * Must be called under RCU read lock. The lockless rbtree traverse can return a
 * false negative (i.e., task not found even though it's in the rbtree), which
 * is mitigated using a seqlock. The seqlock is used to determine if the rbtree
 * was modified during traverse, which is a signal to retry traverse only if the
 * task isn't found. There's no risk of a false positive, so when a task is
 * found it means it's legitimate and the search is done.
 *
 * Note that the seqlock will break if a reader nests into a write-side critical
 * section. Generally, this can occur when reading from IRQ context for an IRQ
 * that fires in the middle of a write-side critical section. This isn't a
 * problem for LKRG though because there's no need to ever call this function
 * from IRQ context, since LKRG only monitors tasks.
 *
 * If readers nested into write-side critical sections is ever needed, a latched
 * rbtree (include/linux/rbtree_latch.h) should be used instead.
 *
 * Do not call this from IRQ context!
 */
static struct p_ed_process *__ed_tree_find_rcu(struct ed_rb_tree *edtree,
					       const struct task_struct *tsk)
{
	struct p_ed_process *edp;
	struct rb_node *node;
#if HAS_RBTREE_RCU
	unsigned seq;

	do {
		seq = read_seqcount_begin(&edtree->seqcount);
		node = rcu_dereference_raw(edtree->root.rb_node);
		while (node) {
			edp = rb_entry(node, typeof(*edp), node);
			if (tsk < edp->p_ed_task.p_task)
				node = rcu_dereference_raw(node->rb_left);
			else if (tsk > edp->p_ed_task.p_task)
				node = rcu_dereference_raw(node->rb_right);
			else
				return edp;
		}
	} while (read_seqcount_retry(&edtree->seqcount, seq));
#else
	read_lock(&edtree->lock);
	node = edtree->root.rb_node;
	while (node) {
		edp = rb_entry(node, typeof(*edp), node);
		if (tsk < edp->p_ed_task.p_task) {
			node = node->rb_left;
		} else if (tsk > edp->p_ed_task.p_task) {
			node = node->rb_right;
		} else {
			read_unlock(&edtree->lock);
			return edp;
		}
	}
	read_unlock(&edtree->lock);
#endif

	return NULL;
}

struct p_ed_process *__ed_task_find_rcu(const struct task_struct *tsk)
{
	struct ed_rb_tree *edtree = task_rb_tree_bucket(tsk);

	return __ed_tree_find_rcu(edtree, tsk);
}

/*
 * Get the ed struct for the current task. An RCU read lock doesn't need to be
 * held when calling this because a task's ed struct is only freed by that task
 * itself upon reaching do_exit().
 */
struct p_ed_process *__ed_task_current(void)
{
	const struct task_struct *tsk = current;
	struct ed_rb_tree *edtree = task_rb_tree_bucket(tsk);
	struct p_ed_process *edp;

	rcu_read_lock();
	edp = __ed_tree_find_rcu(edtree, tsk);
	rcu_read_unlock();

	return edp;
}

static void free_ed_task_rcu(struct rcu_head *rcu)
{
	struct p_ed_process *edp = container_of(rcu, typeof(*edp), rcu);

	free_ed_task(edp);
}

/*
 * Find, delete, and free the current task's entry if one is found. Returns 0 on
 * success, -ENOENT if no entry was found.
 */
int ed_task_del_current(void)
{
	const struct task_struct *tsk = current;
	struct ed_rb_tree *edtree = task_rb_tree_bucket(tsk);
	struct p_ed_process *edp;

	/*
	 * The RCU read lock isn't needed to keep @edp pinned in memory after
	 * finding it because this is the only place where the current task's
	 * @edp can be freed. The lockless lookup is an optimization since
	 * there's no need to try and find @edp under the tree's lock.
	 */
	rcu_read_lock();
	edp = __ed_tree_find_rcu(edtree, tsk);
	rcu_read_unlock();
	if (!edp)
		return -ENOENT;

	/* Remove @edp from the tree */
#if HAS_RBTREE_RCU
	raw_spin_lock(&edtree->lock);
#else
	write_lock(&edtree->lock);
#endif
	__rb_add_or_del(edp, edtree, NULL, NULL);
#if HAS_RBTREE_RCU
	raw_spin_unlock(&edtree->lock);
#else
	write_unlock(&edtree->lock);
#endif

	/*
	 * Free @edp after all RCU users are done using it. Note that @edp's
	 * rb_node must also remain intact during this time because there may be
	 * lockless tree traversals in-flight.
	 */
	call_rcu(&edp->rcu, free_ed_task_rcu);

	return 0;
}

/* Allocate an entry for tracking a specific task */
struct p_ed_process *alloc_ed_task(struct task_struct *tsk)
{
	struct p_ed_process *edp;

	edp = kmem_cache_zalloc(ed_task_cache, GFP_ATOMIC);
	if (!edp)
		return NULL;

	raw_spin_lock_init(&edp->lock);
	edp->kill_task_on_unlock = cookie_kill_task_on_unlock;
	edp->p_ed_task.p_task = tsk;
	return edp;
}

void free_ed_task(struct p_ed_process *edp)
{
	kmem_cache_free(ed_task_cache, edp);
}

int init_ed_task_cache(void)
{
	ed_task_cache = kmem_cache_create("ed_tasks",
					  sizeof(struct p_ed_process), 0,
					  P_LKRG_CACHE_FLAGS, NULL);
	if (!ed_task_cache)
		return -ENOMEM;

	/*
	 * Generate a random, nonzero cookie to use as the default value for
	 * kill_task_on_unlock for each task. This is more secure than just
	 * making kill_task_on_unlock a `bool` because it makes it harder for an
	 * attacker to exploit kill_task_on_unlock.
	 */
	while (!(cookie_kill_task_on_unlock = get_random_long()));

	return 0;
}

static void free_one_tree(struct ed_rb_tree *edtree)
{
#ifdef rbtree_postorder_for_each_entry_safe
	struct p_ed_process *entry, *n;

	/*
	 * Free all entries in this tree without calling rb_erase() which could
	 * rebalance the tree and cause the traverse to miss some nodes.
	 */
	rbtree_postorder_for_each_entry_safe(entry, n, &edtree->root, node)
		free_ed_task(entry);
#else
	struct p_ed_process *entry;
	struct rb_node *node;

	/* Free all entries in this tree for ancient kernels (this is slower) */
	while ((node = rb_first(&edtree->root))) {
		entry = rb_entry(node, typeof(*entry), node);
		rb_erase(&entry->node, &edtree->root);
		free_ed_task(entry);
	}
#endif
}

/* Called after all probes are unregistered and are no longer running */
void destroy_ed_task_cache(void)
{
	int i;

	if (!ed_task_cache)
		return;

	/* Flush all free_ed_task_rcu() callbacks */
	rcu_barrier();

	/* Directly free every remaining tree entry and destroy the cache */
	for (i = 0; i < RB_HASH_SIZE; i++)
		free_one_tree(&ed_hash_tbl[i]);
	kmem_cache_destroy(ed_task_cache);

	ed_task_cache = NULL;
}
