/*
 * pi3's Linux kernel Runtime Guard
 *
 * Component:
 *  - Red-black tree for keeping track of usermode processes
 *
 * Notes:
 *  - Makes sense with own kmem_cache_* allocation
 *
 * Timeline:
 *  - Created: 07.IX.2017
 *
 * Author:
 *  - Adam 'pi3' Zabrocki (http://pi3.com.pl)
 *
 */

#ifndef P_LKRG_EXPLOIT_DETECTION_RB_TREE_H
#define P_LKRG_EXPLOIT_DETECTION_RB_TREE_H

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

#define RB_HASH_MULT (2 * L1_CACHE_BYTES / sizeof(p_rb_hash[0]) + 3)
#define RB_HASH_FUNC(pid) ((uint32_t)(pid) * (uint32_t)RB_HASH_MULT) & RB_HASH_MASK

#define p_alloc_ed_pids()      kmem_cache_alloc(p_ed_pids_cache, GFP_ATOMIC)
#define p_free_ed_pids(name)   kmem_cache_free(p_ed_pids_cache, (void *)(name))

//#ifdef P_LKRG_DEBUG
#define P_DUMP_RB_ED_PIDS_TREE                                                                        \
do {                                                                                                  \
   struct rb_node *p_node;                                                                            \
   unsigned int i;                                                                                    \
                                                                                                      \
   for (i=0; i<RB_HASH_SIZE; i++) {                                                                   \
      for (p_node = rb_first(&p_rb_hash[i].p_tree.tree); p_node; p_node = rb_next(p_node)) {          \
        p_debug_log(P_LOG_DEBUG, "<%s> pid[%d] [uid[%d] gid[%d] ruid[%d] rgid[%d]",                   \
                 rb_entry(p_node, struct p_ed_process, p_rb)->p_ed_task.p_comm,                       \
                 rb_entry(p_node, struct p_ed_process, p_rb)->p_ed_task.p_pid,                        \
                 p_get_uid(&rb_entry(p_node, struct p_ed_process, p_rb)->p_ed_task.p_cred.uid),       \
                 p_get_gid(&rb_entry(p_node, struct p_ed_process, p_rb)->p_ed_task.p_cred.gid),       \
                 p_get_uid(&rb_entry(p_node, struct p_ed_process, p_rb)->p_ed_task.p_real_cred.uid),  \
                 p_get_gid(&rb_entry(p_node, struct p_ed_process, p_rb)->p_ed_task.p_real_cred.gid)   \
                 );                                                                                   \
      }                                                                                               \
   }                                                                                                  \
} while(0);
//#endif

struct p_tasks_root {

   union {
      rwlock_t lock;
      struct rb_root dummy;
   } p_lock;

   union {
      struct rb_root tree;
      rwlock_t dummy;
   } p_tree;

};

struct p_ed_process {

   struct rb_node p_rb;
   struct p_ed_process_task p_ed_task;
   /* ... add other driver-specific fields */

};

extern struct kmem_cache *p_ed_pids_cache;
extern struct p_tasks_root p_rb_hash[RB_HASH_SIZE];

struct p_ed_process *p_rb_find_ed_pid(struct rb_root *p_root, pid_t p_arg);
struct p_ed_process *p_rb_add_ed_pid(struct rb_root *p_root, pid_t p_arg, struct p_ed_process *p_source);
void p_rb_del_ed_pid(struct rb_root *p_root, struct p_ed_process *p_source);
int p_init_rb_ed_pids(void);
void p_delete_rb_ed_pids(void);


static inline struct rb_root *p_rb_hash_tree_lookup(pid_t p_pid) {
   return &p_rb_hash[RB_HASH_FUNC(p_pid)].p_tree.tree;
}

static inline rwlock_t *p_rb_hash_lock_lookup(pid_t p_pid) {
   return &p_rb_hash[RB_HASH_FUNC(p_pid)].p_lock.lock;
}

static inline struct p_ed_process *p_find_ed_by_pid(pid_t p_arg) {
   return p_rb_find_ed_pid(p_rb_hash_tree_lookup(p_arg), p_arg);
}

static inline void p_rb_init_ed_pid_node(struct rb_node *rb) {

   rb->__rb_parent_color = 0;
   rb->rb_right = NULL;
   rb->rb_left = NULL;
   RB_CLEAR_NODE(rb);
}

#endif
