/**
@file     procfs_manager.c
@brief    Manager for procfs
@details  Copyright (c) 2025 Acronis International GmbH
@author   Denis Kopyrin (denis.kopyrin@acronis.com)
@since    $Id: $
*/

#include "procfs_manager.h"

#include "hashtable_compat.h"
#include "hash_fast.h"
#include "path_tools.h"
#include "memory.h"

#include <linux/jiffies.h>
#include <linux/list.h>

#define TABLE_SIZE_BITS 15
#define TABLE_SIZE (1 << (TABLE_SIZE_BITS - 1)) // 16384
#define TABLE_CLEAN_SIZE 256

#ifndef list_first_entry_or_null
#define list_first_entry_or_null(ptr, type, member) (list_empty(ptr) ? NULL : list_first_entry(ptr, type, member))
#endif

#ifdef KERNEL_MOCK
#include "mock/mock.h"
#endif

// 30 seconds
#define TTL msecs_to_jiffies(30000)

typedef struct
{
	atomic_t refcount;

	struct hlist_node hash_node;

	struct list_head lru_list_node;
	bool lru_list_inserted;
	unsigned long lru_deadline;

	struct rcu_head rcu;

	hashtable_procfs_key_t key;
} hashtable_procfs_node_t;

typedef struct procfs_manager
{
	struct mutex table_writer_lock;
	bool active;
	ssize_t entries_count;
	// Entry per caller + callee as a key.
	DECLARE_HASHTABLE(seen_entries_hashtable, TABLE_SIZE_BITS);
	struct list_head seen_entries_lru_list;
} procfs_manager_t;

static procfs_manager_t *global_procfs_manager;

// MARK: Procfs node
static hashtable_procfs_node_t *node_alloc(const hashtable_procfs_key_t *key)
{
	hashtable_procfs_node_t *node = mem_alloc(sizeof(hashtable_procfs_node_t));
	if (!node)
		return NULL;

	atomic_set(&node->refcount, 1);
	node->lru_list_inserted = false;
	INIT_HLIST_NODE(&node->hash_node);
	INIT_LIST_HEAD(&node->lru_list_node);

	node->key = *key;
	return node;
}

static void node_rcu_free(struct rcu_head *rcu)
{
	hashtable_procfs_node_t *node = container_of(rcu, hashtable_procfs_node_t, rcu);
	mem_free(node);
}

static void node_put(hashtable_procfs_node_t *node)
{
	if (atomic_dec_and_test(&node->refcount))
		call_rcu(&node->rcu, node_rcu_free);
}

static bool key_equal(const hashtable_procfs_key_t *k1, const hashtable_procfs_key_t *k2)
{
	return k1->access_type == k2->access_type && k1->caller_pid_version == k2->caller_pid_version && k1->target_pid_version == k2->target_pid_version;
}

static int key_hash(const hashtable_procfs_key_t *key)
{
	return murmur_hash(&key->caller_pid_version, sizeof(key->caller_pid_version) + sizeof(key->target_pid_version) + sizeof(key->access_type)) >> (64 - TABLE_SIZE_BITS);
}

static hashtable_procfs_node_t *find_ref_rcu(int hash, const hashtable_procfs_key_t *key)
{
	hashtable_procfs_node_t *search_node;
	hlist_for_each_entry_rcu(search_node, &global_procfs_manager->seen_entries_hashtable[hash], hash_node)
	{
		if (!key_equal(&search_node->key, key))
			continue;

		if (atomic_inc_not_zero(&search_node->refcount))
			return search_node;
		else
			return NULL;
	}

	return NULL;
}

static hashtable_procfs_node_t *find(int hash, const hashtable_procfs_key_t *key)
{
	hashtable_procfs_node_t *search_node;
	hlist_for_each_entry(search_node, &global_procfs_manager->seen_entries_hashtable[hash], hash_node)
	{
		if (key_equal(&search_node->key, key))
			return search_node;
	}

	return NULL;
}

static void erase_impl(hashtable_procfs_node_t *node)
{
	hash_del_rcu(&node->hash_node);
	list_del(&node->lru_list_node);
	node->lru_list_inserted = false;
	global_procfs_manager->entries_count--;
	node_put(node);
}

static void refresh_impl(hashtable_procfs_node_t *node)
{
	if (node->lru_list_inserted)
	{
		node->lru_deadline = jiffies + TTL;
		list_del(&node->lru_list_node);
		list_add_tail(&node->lru_list_node, &global_procfs_manager->seen_entries_lru_list);
	}
}

// MARK: Procfs manager
int procfs_manager_init(void)
{
	global_procfs_manager = vmem_alloc(sizeof(procfs_manager_t));
	if (!global_procfs_manager)
		return -ENOMEM;

	mutex_init(&global_procfs_manager->table_writer_lock);
	global_procfs_manager->active = false;
	global_procfs_manager->entries_count = 0;
	hash_init(global_procfs_manager->seen_entries_hashtable);
	INIT_LIST_HEAD(&global_procfs_manager->seen_entries_lru_list);

	return 0;
}

void procfs_manager_deinit(void)
{
	if (!global_procfs_manager)
		return;

	vmem_free(global_procfs_manager);
}

void procfs_manager_activate(void)
{
	mutex_lock(&global_procfs_manager->table_writer_lock);
	global_procfs_manager->active = true;
	mutex_unlock(&global_procfs_manager->table_writer_lock);
}

void procfs_manager_deactivate(void)
{
	mutex_lock(&global_procfs_manager->table_writer_lock);
	if (global_procfs_manager->active)
	{
		while (1)
		{
			hashtable_procfs_node_t *node = list_first_entry_or_null(&global_procfs_manager->seen_entries_lru_list, hashtable_procfs_node_t, lru_list_node);
			if (!node)
				break;

			erase_impl(node);
		}

		global_procfs_manager->active = false;
	}
	mutex_unlock(&global_procfs_manager->table_writer_lock);
}

static void sweep_impl(void)
{
	while (1)
	{
		hashtable_procfs_node_t *node = list_first_entry_or_null(&global_procfs_manager->seen_entries_lru_list, hashtable_procfs_node_t, lru_list_node);
		if (!node)
			break;

		if (time_after(jiffies, node->lru_deadline))
			erase_impl(node);
		else
			break;
	}

	while (global_procfs_manager->entries_count >= TABLE_SIZE - TABLE_CLEAN_SIZE)
	{
		hashtable_procfs_node_t *node = list_first_entry_or_null(&global_procfs_manager->seen_entries_lru_list, hashtable_procfs_node_t, lru_list_node);
		if (!node)
			break;

		erase_impl(node);
	}
}

static bool procfs_manager_key_exist(const hashtable_procfs_key_t *key)
{
	int hash;
	hashtable_procfs_node_t *node;

	if (!global_procfs_manager->active)
	{
		return false;
	}

	hash = key_hash(key);

	rcu_read_lock();
	node = find_ref_rcu(hash, key);
	rcu_read_unlock();
	if (node)
	{
		mutex_lock(&global_procfs_manager->table_writer_lock);
		refresh_impl(node);
		mutex_unlock(&global_procfs_manager->table_writer_lock);
		node_put(node);
		return true;
	}

	mutex_lock(&global_procfs_manager->table_writer_lock);
	if(!READ_ONCE(global_procfs_manager->active))
	{
		mutex_unlock(&global_procfs_manager->table_writer_lock);
		return false;
	}

	node = find(hash, key);
	if (node)
	{
		refresh_impl(node);
		mutex_unlock(&global_procfs_manager->table_writer_lock);
		return true;
	}

	sweep_impl();

	node = node_alloc(key);
	if (!node)
	{
		mutex_unlock(&global_procfs_manager->table_writer_lock);
		return false;
	}

	hlist_add_head_rcu(&node->hash_node, &global_procfs_manager->seen_entries_hashtable[hash]);

	node->lru_deadline = jiffies + TTL;
	list_add_tail(&node->lru_list_node, &global_procfs_manager->seen_entries_lru_list);
	node->lru_list_inserted = true;

	global_procfs_manager->entries_count++;
	mutex_unlock(&global_procfs_manager->table_writer_lock);

	return false;
}

static inline bool dname_eq(const struct dentry *d, const char *name, size_t name_len)
{
	return d->d_name.len == name_len && !memcmp(d->d_name.name, name, name_len);
}

bool procfs_should_send(task_info_t *caller_task_info, const struct path *target_path, hashtable_procfs_key_t *out_key, task_info_t **target_task_info)
{
	SiProcfsAccessType access_type;
	uint64_t caller_pid_version = 0;
	uint64_t target_pid_version = 0;
	pid_t target_pid = 0;

	if(!READ_ONCE(global_procfs_manager->active))
		return false;
		
	if (!caller_task_info || !target_path || !out_key)
		return false;
	caller_pid_version = READ_ONCE(caller_task_info->pid_version);

	target_pid = pid_under_proc(target_path);
	if (target_pid == -1)
		return false;

	// target_path here is guaranteed to be under /proc/<target_pid>/...
	{
		struct dentry *root_dentry = target_path->mnt->mnt_root;
		struct dentry *dentry = target_path->dentry;
		int i = 0;
		rcu_read_lock();
		// target_path must be like /proc/<pid>/...
		for (; i < 2; i++)
		{
			if (!dentry || dentry == root_dentry)
			{
				rcu_read_unlock();
				return false;
			}
			dentry = rcu_dereference(dentry->d_parent);
		}
		if (!dentry || dentry != root_dentry)
		{
			rcu_read_unlock();
			return false;
		}
		rcu_read_unlock();

		dentry = target_path->dentry;
		if (dname_eq(dentry, "mem", 3))
		{
			access_type = SI_PROCFS_AT_MEM;
		}
		else if (dname_eq(dentry, "maps", 4))
		{
			access_type = SI_PROCFS_AT_MAPS;
		}
		else if (dname_eq(dentry, "exe", 3))
		{
			access_type = SI_PROCFS_AT_EXE;
		}
		else if (dname_eq(dentry, "comm", 4))
		{
			access_type = SI_PROCFS_AT_COMM;
		}
		else
		{
			// not interested
			return false;
		}
	}

	{
		*target_task_info = task_info_map_get_by_pid(target_pid, 0);
		if (!*target_task_info)
			return false;
		target_pid_version = READ_ONCE((*target_task_info)->pid_version);
	}

	out_key->access_type = access_type;
	out_key->caller_pid_version = caller_pid_version;
	out_key->target_pid_version = target_pid_version;

	if (procfs_manager_key_exist(out_key))
	{
		return false;
	}

	return true;
}