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 
11 DEFINE_SPINLOCK(rbtree_lock);
12 
find_process_jit_space(struct rb_root *root, int pid)13 struct 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 
update_process_jit_space(struct rb_root *root, int pid, unsigned long cookie, unsigned long *err)41 struct 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 
delete_process_jit_space(struct rb_root *root, int pid)77 struct 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 }