1// SPDX-License-Identifier: GPL-2.0-or-later
2/*
3 * Copyright (c) 2023 Huawei Device Co., Ltd.
4 */
5
6#include <linux/sched.h>
7#include <linux/types.h>
8#include "avc.h"
9#include "jit_process.h"
10
11DEFINE_SPINLOCK(rbtree_lock);
12
13struct result_of_find_process find_process_jit_space(struct rb_root *root, int pid)
14{
15	struct rb_node **node = &(root->rb_node), *parent = NULL;
16	struct result_of_find_process result = {NULL, NULL, NULL};
17
18	spin_lock(&rbtree_lock);
19	while (*node) {
20		struct jit_process *now = container_of(*node, struct jit_process, node);
21
22		parent = *node;
23		if (now->pid == pid) {
24			result.head = &(now->head);
25			break;
26		}
27		else if (now->pid < pid) {
28			node = &((*node)->rb_left);
29		}
30		else if (now->pid > pid) {
31			node = &((*node)->rb_right);
32		}
33	}
34	spin_unlock(&rbtree_lock);
35
36	result.node = node;
37	result.parent = parent;
38	return result;
39}
40
41struct list_head *update_process_jit_space(struct rb_root *root,
42	int pid, unsigned long cookie, unsigned long *err)
43{
44	struct result_of_find_process result = find_process_jit_space(root, pid);
45
46	if (result.head != NULL) {
47		// find node which already exist
48		struct jit_process *now = container_of(result.head, struct jit_process, head);
49		if (now->cookie == cookie) {
50			return result.head;
51		} else {
52			*err = -EACCES;
53			return NULL;
54		}
55	} else {
56	// init node
57		struct jit_process *process = kmalloc(sizeof(struct jit_process), GFP_KERNEL);
58		if (process == NULL) {
59			jit_memory_log_error("malloc for rbTree node failed");
60			*err = -ENOMEM;
61			return NULL;
62		}
63		process->cookie = cookie;
64		process->pid = pid;
65		process->head.next = &(process->head);
66		process->head.prev = &(process->head);
67		/* Add new node and rebalance tree. */
68		spin_lock(&rbtree_lock);
69		rb_link_node(&(process->node), result.parent, result.node);
70		rb_insert_color(&(process->node), root);
71		spin_unlock(&rbtree_lock);
72
73		return &(process->head);
74	}
75}
76
77struct jit_process *delete_process_jit_space(struct rb_root *root, int pid)
78{
79	struct list_head *head = (find_process_jit_space(root, pid).head);
80	if (head == NULL)
81		return NULL;
82
83	struct jit_process *victim = container_of(head, struct jit_process, head);
84	if (victim == NULL)
85		return NULL;
86
87	spin_lock(&rbtree_lock);
88	rb_erase(&(victim->node), root);
89	spin_unlock(&rbtree_lock);
90
91	return victim;
92}