/**
@file
@brief    Process (aka task) info storage
@details  Copyright (c) 2017-2021 Acronis International GmbH
@author   Ivan Matveev (ivan.matveev@acronis.com)
@since    $Id: $
*/

#include "task_info_map.h"
#include "transport.h"
#include "compat.h"
#include "debug.h"
#include "memory.h"

#include <linux/mutex.h>
#include <linux/mm.h>   // get_task_exe_file()

static struct mutex global_task_info_maps_mutex;
static task_info_map_t* global_task_info_maps[MAX_TASK_INFO_MAP_SIZE] = {0};

const char *task_status_to_string(task_status_t status)
{
	switch (status)
	{
#define CASE_TS_RETURN(t) case TS_##t: return #t
	CASE_TS_RETURN(UNKNOWN);
	CASE_TS_RETURN(IGNORE);
	CASE_TS_RETURN(WHITE);
	CASE_TS_RETURN(BLACK);
	CASE_TS_RETURN(GREY);
#undef CASE_TS_RETURN
	default: return "?";
	}
}

static task_info_t *task_info_init(task_info_t *task_info, pid_t pid)
{
	DPRINTF("task_info=%p pid=%i", task_info, pid);
	RB_CLEAR_NODE(&task_info->rb_node);
	task_info->pid = pid;

	atomic_set(&task_info->ref_cnt, 1);
	spin_lock_init(&task_info->spinlock);

	task_info->status = TS_UNKNOWN;
	/*
		struct path {
			struct vfsmount *mnt;
			struct dentry *dentry;
		};
	*/
	task_info->exe_key = (file_key_t){};
    task_info->is_exe_path_changed = true;

	return task_info;
}

static task_info_t *task_info_new(pid_t pid)
{
	task_info_t *task_info = mem_alloc0(sizeof(task_info_t));
	if (task_info) {
		task_info_init(task_info, pid);
	}
	return task_info;
}

static task_info_t *task_info_ref(task_info_t *task_info)
{
	atomic_inc(&task_info->ref_cnt);
	return task_info;
}

static void task_info_free(task_info_t *task_info)
{
	DPRINTF("task_info=%p", task_info);
	mem_free(task_info);
}

static void task_info_unref(task_info_t *task_info)
{
	DPRINTF("pid=%d ref_cnt=%d", task_info->pid, atomic_read(&task_info->ref_cnt));
	if (atomic_dec_and_test(&task_info->ref_cnt)) {
		task_info_free(task_info);
	}
}

static int task_info_map_init(task_info_map_t *map)
{
	spin_lock_init(&map->spinlock);
	map->root = RB_ROOT;

    atomic_set(&map->refcount, 1);
	return 0;
}

static task_info_map_t* task_info_map_alloc(void)
{
    task_info_map_t *map = mem_alloc(sizeof(task_info_map_t));
    if (map) {
        task_info_map_init(map);
    }
    return map;
}

static void task_info_map_clear(task_info_map_t *map);
static void task_info_map_free_deferred(struct rcu_head* _map)
{
    task_info_map_t *map = container_of(_map, task_info_map_t, rcu);
    task_info_map_clear(map);
    mem_free(map);
}

static task_info_map_t* task_info_map_get(int idx)
{
    task_info_map_t *map = NULL;
    if (idx >= MAX_TASK_INFO_MAP_SIZE)
    {
        return NULL;
    }
    rcu_read_lock();
    map = rcu_dereference(global_task_info_maps[idx]);
    if (map)
    {
        if (atomic_inc_not_zero(&map->refcount))
        {
            DPRINTF("task_info_map_get: %d", idx);
        }
        else
        {
            map = NULL;
        }
    }
    rcu_read_unlock();

    return map;
}

static void task_info_map_put(task_info_map_t *map)
{
    if (atomic_dec_and_test(&map->refcount))
    {
        call_rcu(&map->rcu, task_info_map_free_deferred);
    }
}

void task_info_maps_init(void)
{
    unsigned char i = 0;

    mutex_init(&global_task_info_maps_mutex);
    DPRINTF("task_info_maps_init");
    for (i = 0; i < MAX_TASK_INFO_MAP_SIZE; i++) {
        global_task_info_maps[i] = NULL;
    }
}

static void task_info_map_clear(task_info_map_t *map)
{
	struct rb_node *node;

	node = map->root.rb_node;
	while (node) {
        task_info_t *task_info;
		task_info = rb_entry(node, task_info_t, rb_node);
		rb_erase(&task_info->rb_node, &map->root);
        node = map->root.rb_node;
        if (atomic_read(&task_info->ref_cnt) != 1)
        {
            DPRINTF("task info [%d] ref_cnf[%d] is not equal to 1 when clearing", task_info->pid, atomic_read(&task_info->ref_cnt));
        }
        task_info_unref(task_info);
	}
}

// Called on module shutdown
void task_info_maps_clear(void)
{
    unsigned char i = 0;

    DPRINTF("task_info_maps_clear");
    for (i = 0; i < MAX_TASK_INFO_MAP_SIZE; i++)
    {
        task_info_map_t *map;
        mutex_lock(&global_task_info_maps_mutex);
        map = global_task_info_maps[i];
        rcu_assign_pointer(global_task_info_maps[i], NULL);
        mutex_unlock(&global_task_info_maps_mutex);

        if (map) {
            task_info_map_put(map);
        }
    }

    synchronize_rcu();
	rcu_barrier();
}

// Assume that each transport only use it's unique valid map_id after acquiring
unsigned char acquire_task_info_map(void)
{
    unsigned char map_id = MAX_TASK_INFO_MAP_SIZE + 1;
    unsigned char i = 0;
    mutex_lock(&global_task_info_maps_mutex);
    for (i = 0; i < MAX_TASK_INFO_MAP_SIZE; i++)
    {
        task_info_map_t* map;
        if (global_task_info_maps[i])
            continue;

        map = task_info_map_alloc();
        if (!map) {
            EPRINTF("Failed to allocate task_info_map");
            mutex_unlock(&global_task_info_maps_mutex);
            return 0;
        }

        rcu_assign_pointer(global_task_info_maps[i], map);
        map_id = i;
        break;
    }
    mutex_unlock(&global_task_info_maps_mutex);

    if (map_id < MAX_TASK_INFO_MAP_SIZE)
    {
        DPRINTF_LEVEL(LOG_LEVEL_DEBUG1, "Acquire task_info_map ok: %u", map_id);
    }
    else
    {
        EPRINTF("Failed to acquire task_info_map");
    }
    return map_id;
}

void release_task_info_map(unsigned char map_id)
{
    task_info_map_t *map;
    if (map_id >= MAX_TASK_INFO_MAP_SIZE)
    {
        return;
    }

    mutex_lock(&global_task_info_maps_mutex);
    map = global_task_info_maps[map_id];
    rcu_assign_pointer(global_task_info_maps[map_id], NULL);
    mutex_unlock(&global_task_info_maps_mutex);

    if (map)
        task_info_map_put(map);
}

static task_info_t *task_info_lookup(task_info_map_t *map, pid_t pid)
{
	struct rb_node *node;
	task_info_t *task_info = NULL;

	DPRINTF("pid=%d", pid);
	spin_lock(&map->spinlock);
    node = map->root.rb_node;
	while (node) {
		task_info_t *node_task_info = rb_entry(node, task_info_t, rb_node);
		pid_t node_pid = node_task_info->pid;

		if (pid < node_pid) {
			node = node->rb_left;
		} else if (pid > node_pid) {
			node = node->rb_right;
		} else {
			task_info = task_info_ref(node_task_info);
			break;
		}
	}

	spin_unlock(&map->spinlock);
	DPRINTF_RATELIMITED("task_info=%p pid=%d", task_info, pid);
	return task_info;
}

static task_info_t *task_info_map_insert(task_info_map_t *map, task_info_t *new_task_info)
{
	pid_t pid = new_task_info->pid;
	struct rb_node *parent = NULL;
	struct rb_node **link;

	DPRINTF_RATELIMITED("new_task_info=%p pid=%i", new_task_info, pid);
	spin_lock(&map->spinlock);
    link = &(map->root.rb_node);
	while (*link) {
		task_info_t *node_task_info;
		pid_t node_pid;

		parent = *link;
		node_task_info = rb_entry(parent, task_info_t, rb_node);
		node_pid = node_task_info->pid;

		if (pid < node_pid) {
			link = &parent->rb_left;
		} else if (pid > node_pid) {
			link = &parent->rb_right;
		} else {
			// collision
			DPRINTF_RATELIMITED("collision");
			task_info_ref(node_task_info);
			spin_unlock(&map->spinlock);
			return node_task_info;
		}
	}
	// do 'inc' for 'map->root'
	task_info_ref(new_task_info);
	rb_link_node(&new_task_info->rb_node, parent, link);
	rb_insert_color(&new_task_info->rb_node, &map->root);
	DPRINTF_RATELIMITED("inserted");

	spin_unlock(&map->spinlock);
	return new_task_info;
}

// Warning: May block!
static task_info_t *task_info_add(task_info_map_t *map, pid_t pid)
{
	task_info_t *task_info;
	task_info_t *new_task_info;

	DPRINTF_RATELIMITED("pid=%d", pid);
	task_info = task_info_lookup(map, pid);
	if (task_info) {
		DPRINTF_RATELIMITED("pid=%i is already in the map (task_info=%p)", pid, task_info);
		return task_info;
	}

	new_task_info = task_info_new(pid);
	if (!new_task_info) {
		DPRINTF_RATELIMITED("out of memory");
		return NULL;
	}

	task_info = task_info_map_insert(map, new_task_info);
	if (task_info != new_task_info) {
		// collision
		DPRINTF_RATELIMITED("collision");
		task_info_unref(new_task_info);
	}
	return task_info;
}

static int task_info_map_del(task_info_map_t *map, pid_t pid)
{
	task_info_t *task_info;

	DPRINTF("pid=%d", pid);
	task_info = task_info_lookup(map, pid);
	if (!task_info) {
        // task info may be deleted in task_info_map_clear_exited_list()
        // so here no need to print warning
		DPRINTF("pid=%d is missing in the map", pid);
		return -ENOENT;
	}

	spin_lock(&map->spinlock);
    if (map->root.rb_node != NULL)
    {
        rb_erase(&task_info->rb_node, &map->root);
    }
    spin_unlock(&map->spinlock);

	// undo 'inc' done for 'map->root'
	task_info_unref(task_info);

	// undo 'inc' done in 'task_info_lookup()'
	task_info_unref(task_info);
	DPRINTF("pid=%d", pid);
	return 0;
}

static void task_info_map_remove(task_info_map_t *map, task_info_t *task_info)
{
    spin_lock(&map->spinlock);
    if (map->root.rb_node != NULL)
    {
        rb_erase(&task_info->rb_node, &map->root);
    }
    if (atomic_read(&task_info->ref_cnt) >= 2)
    {
        // undo 'inc' done for 'map->root'
        task_info_unref(task_info);
    }
    spin_unlock(&map->spinlock);
}

int task_info_map_del_by_map_id(unsigned char map_id, pid_t pid)
{
    task_info_map_t *map;
    int ret = 0;
    if (map_id >= MAX_TASK_INFO_MAP_SIZE)
    {
        return ENOENT;
    }

    map = task_info_map_get(map_id);
    if (!map)
        return ENOENT;

    ret = task_info_map_del(map, pid);
    task_info_map_put(map);
    return ret;
}

static task_info_t *get_task_info_by_map_id(unsigned char map_id, pid_t pid)
{
    task_info_map_t *map = NULL;
    task_info_t *task_info = NULL;
    if (map_id >= MAX_TASK_INFO_MAP_SIZE)
    {
        return NULL;
    }

    map = task_info_map_get(map_id);
    if (!map)
        return NULL;

    task_info = task_info_lookup(map, pid);
    task_info_map_put(map);

    return task_info;
}

// Warning: May block!
int task_info_status_set_by_map_id(unsigned char map_id, pid_t pid, task_status_t status)
{
    task_info_t *task_info = NULL;
    task_info_map_t *map = NULL;
    DPRINTF_LEVEL(LOG_LEVEL_DEBUG1, "map_id:%u pid=%i status=%i/%s", map_id, pid, status, task_status_to_string(status));
    if (map_id >= MAX_TASK_INFO_MAP_SIZE)
    {
        return ENOENT;
    }

    map = task_info_map_get(map_id);
    if (!map)
        return ENOENT;

    task_info = task_info_add(map, pid);
    if (!task_info) {
		WPRINTF("'%s(pid=%i)' failure", "task_info_add", pid);
        task_info_map_put(map);
		return -ENOMEM;
	}
    spin_lock(&task_info->spinlock);
	task_info->status = status;
	spin_unlock(&task_info->spinlock);
	task_info_unref(task_info);
    task_info_map_put(map);
	return 0;
}

void task_info_map_on_exit_event(pid_t tgid, pid_t pid)
{
    unsigned char i = 0;
    task_info_t *task_info = NULL;

    (void) tgid;
    for (i = 0; i < MAX_TASK_INFO_MAP_SIZE; i++)
    {
        task_info_map_t *map = task_info_map_get(i);
        if (!map)
            continue;

        task_info = task_info_lookup(map, pid);
        if (task_info) {
            task_info_map_remove(map, task_info);
            task_info_unref(task_info);
        }
        task_info_map_put(map);
    }
}

static inline bool file_key_equal(const file_key_t* k1, const file_key_t* k2)
{
    return k1->ptr == k2->ptr && k1->ino == k2->ino && k1->dev == k2->dev && k1->gen == k2->gen;
}

static inline void make_key_from_inode(file_key_t* key, const struct inode *inode) {
	key->ptr = (uint64_t) inode;
	key->ino = (uint64_t) inode->i_ino;
	key->gen = (uint64_t) inode->i_generation;
	key->dev = (uint64_t) inode->i_sb ? inode->i_sb->s_dev : 0;
}

static inline void make_key(file_key_t* key, const struct path *path) {
	return make_key_from_inode(key, path->dentry->d_inode);
}

int check_exec_with_task_info(pid_t tgid, struct path exe_path)
{
    unsigned char i = 0;
    task_info_t *task_info = NULL;
    bool is_exe_path_changed = false;
    file_key_t exe_key = (file_key_t){};

    if (exe_path.dentry && exe_path.dentry->d_inode)
        make_key(&exe_key, &exe_path);

    for (i = 0; i < MAX_TASK_INFO_MAP_SIZE; i++)
    {
        task_info_map_t *map = task_info_map_get(i);
        if (!map)
            continue;

        task_info = NULL;
        task_info = task_info_add(map, tgid);
        if (task_info)
        {
            spin_lock(&task_info->spinlock);
            if (file_key_equal(&task_info->exe_key, &exe_key))
            {
            }
            else
            {
                task_info->is_exe_path_changed = true;
                task_info->exe_key = exe_key;
            }
            if(task_info->is_exe_path_changed){
                is_exe_path_changed = true;
            }
            spin_unlock(&task_info->spinlock);
            task_info_unref(task_info);
        }
        else
        {
            EPRINTF("Failed to add task info");
        }

        task_info_map_put(map);
    }

    if (is_exe_path_changed)
    {
        return 1;
    }
    return 0;
}

static bool msg_should_skip(task_info_t *task_info, msg_type_img_t msg_type)
{
    bool ret = false;
    task_status_t status;

    if (msg_type == MT_EXEC)
    {
    }
    else if (msg_type == MT_EXIT || msg_type == MT_FORK ||
             msg_type == MT_FILE_PRE_CREATE || msg_type == MT_FILE_CREATE_EX ||
             msg_type == MT_PRE_RENAME_EX || msg_type == MT_RENAME_EX || msg_type == MT_RENAME ||
             msg_type == MT_PRE_UNLINK_EX || msg_type == MT_UNLINK_EX || msg_type == MT_UNLINK)
    {
        // for event in EXIT, FORK, CREATE, RENAME and UNLINK, always send
    }
    else
    {
        spin_lock(&task_info->spinlock);
        status = task_info->status;
        spin_unlock(&task_info->spinlock);

        // For event not in EXIT, FORK, CREATE, RENAME and UNLINK, check status is TS_WHITE or TS_IGNORE
        if (status == TS_IGNORE || status == TS_WHITE)
        {
            ret = true;
        }
    }

    return ret;
}

bool is_exe_path_changed(task_info_t *task_info)
{
    struct file *exe_file;
    file_key_t exe_key = (file_key_t){};
    bool ret = false;

    exe_file = get_task_exe_file_compat(current);
    if (!exe_file)
    {
        return false;
    }
    if (exe_file->f_path.dentry && exe_file->f_path.dentry->d_inode)
        make_key(&exe_key, &exe_file->f_path);

    fput(exe_file);

    spin_lock(&task_info->spinlock);
    if (!file_key_equal(&task_info->exe_key, &exe_key))
    {
        ret = true;
    }
    spin_unlock(&task_info->spinlock);

    return ret;
}

bool sys_call_fs_events_should_skip(pid_t pid, uint64_t events_mask)
{
    unsigned char i = 0;
    task_info_t *task_info = NULL;
    bool ret = true;

    for (i = 0; i < MAX_TASK_INFO_MAP_SIZE; i++)
    {
        task_info_map_t *map;
        msg_type_img_t msg_type;

        map = task_info_map_get(i);
        if (!map)
            continue;

        task_info = task_info_lookup(map, pid);
        if (!task_info)
        {
            task_info_map_put(map);
            return false;
        }

        if (is_exe_path_changed(task_info))
        {
            task_info_unref(task_info);
            task_info_map_put(map);
            return false;
        }

        for (msg_type = MT_EXEC; msg_type < MT_MAX; msg_type++)
        {
            uint64_t event_mask = MSG_TYPE_TO_EVENT_MASK(msg_type);
            if (event_mask & events_mask)
            {
                if (!msg_should_skip(task_info, msg_type))
                {
                    ret = false;
                    break;
                }
            }
        }

        task_info_unref(task_info);
        task_info_map_put(map);

        if (!ret)
        {
            break;
        }
    }
    return ret;
}

static bool should_send(unsigned char map_id, msg_type_img_t msg_type)
{
    bool ret = true;
    task_info_t *task_info = NULL;
    task_status_t status;

    task_info = get_task_info_by_map_id(map_id, current->tgid);
    if (task_info == NULL)
    {
        return false;
    }

    if (msg_type == MT_EXEC)
    {
        spin_lock(&task_info->spinlock);
        if (task_info->is_exe_path_changed)
        {
            task_info->is_exe_path_changed = false;
        }
        else
        {
            ret = false;
        }
        spin_unlock(&task_info->spinlock);
    }
    else if (msg_type == MT_EXIT || msg_type == MT_FORK ||
             msg_type == MT_FILE_PRE_CREATE || msg_type == MT_FILE_CREATE_EX ||
             msg_type == MT_PRE_RENAME_EX || msg_type == MT_RENAME_EX || msg_type == MT_RENAME ||
             msg_type == MT_PRE_UNLINK_EX || msg_type == MT_UNLINK_EX || msg_type == MT_UNLINK)
    {
        // for event in EXIT, FORK, CREATE, RENAME and UNLINK, always send
    }
    else
    {
        spin_lock(&task_info->spinlock);
        status = task_info->status;
        spin_unlock(&task_info->spinlock);

        // For event not in EXIT, FORK, CREATE, RENAME and UNLINK, check status is TS_WHITE or TS_IGNORE
        if (status == TS_IGNORE || status == TS_WHITE)
        {
            ret = false;
        }
    }

    task_info_unref(task_info);
    return ret;
}

void get_should_send_task_info_map_ids(bool *should_send_map_ids, msg_type_img_t msg_type)
{
    int i;
    for (i = 0; i < MAX_TASK_INFO_MAP_SIZE; i++)
    {
        should_send_map_ids[i] = false;
        if (should_send(i, msg_type))
        {
            should_send_map_ids[i] = true;
        }
    }
}