rkit_helper.h

#ifndef RKIT_HELPER_H
#define RKIT_HELPER_H

#include <linux/kprobes.h>
#include <linux/ftrace.h>
#include <linux/linkage.h>
#include <linux/kprobes.h>
#include <linux/kernel.h>
#include <linux/types.h>
#include <linux/module.h>
#include <linux/set_memory.h>

static bool ibt_status = false;

#define HOOK(_func_name, _new_func, _orig_func) \
{ \
        .func_name = (_func_name), \
        .new_func = (_new_func), \
        .orig_func = (_orig_func), \
} 


typedef unsigned long (*kallsyms_lookup_name_t)(const char *symbol_name);
typedef struct fhook hook;

static kallsyms_lookup_name_t get_kallsyms_lookup_name(void);
static void notrace ftrace_callback(unsigned long ip, unsigned long parent_ip, struct ftrace_ops *op, struct ftrace_regs *regs);
static int hook_kfunction(hook *hk, unsigned int hk_cnt);
static int unhook_kfunction(struct fhook *hk, unsigned int hk_cnt);
static void ibt_toggle(void); 
static bool ibt_is_on(void); 

#define FLAG_TRACED (1 << 1)
#define FLAG_FILTERED (1 << 2)

#define MSR_IA_S_CET 0x6a2
#define IBT_BIT_POSITION 2

struct fhook {
    const char *func_name;
    void *new_func;
    void *orig_func;
    unsigned long mcount_inst_addr;
    int status_flag;
    struct ftrace_ops fops;
};

static kallsyms_lookup_name_t get_kallsyms_lookup_name(void)
{
	kallsyms_lookup_name_t kallsyms_lookup_name_fptr = NULL;
	struct kprobe kp = {
		.symbol_name = "kallsyms_lookup_name"
	};
		
	if(!register_kprobe(&kp)) {
		kallsyms_lookup_name_fptr = (kallsyms_lookup_name_t)kp.addr;
		unregister_kprobe(&kp);
	} 
	return kallsyms_lookup_name_fptr;
}

static void notrace ftrace_callback(unsigned long ip, unsigned long parent_ip, struct ftrace_ops *op, struct ftrace_regs *regs)
{
    struct fhook *hk = (struct fhook*) container_of(op, struct fhook, fops);
    
    struct pt_regs * _regs;
    _regs = (struct pt_regs*)regs;
    if(!within_module(parent_ip, THIS_MODULE))
        _regs->ip = (unsigned long)hk->new_func;
}

static int hook_kfunction(struct fhook *hk, unsigned int hk_cnt)
{
    int error;
    kallsyms_lookup_name_t kallsyms_lookup_name;
    
    ibt_status = ibt_is_on();
    if(ibt_status)
        pr_debug("Detected Indirect Branch Tracking...");
    if(ibt_status)
        ibt_toggle();

    kallsyms_lookup_name = get_kallsyms_lookup_name();
    
    for(int i = 0; i < hk_cnt; i++, hk++) {
        /*
         * Initially hk->mcount_inst_addr will hold the desired function to be hooked.
         * However after instrumentation with ftrace the address would point to the instrumentation code for ftrace.
         */
        hk->mcount_inst_addr = kallsyms_lookup_name(hk->func_name); 
        if(!hk->mcount_inst_addr) {
            pr_debug("kallsyms_lookup_name() failed when locating %s function\n", hk->func_name); 
            return -1;
        }
        
        hk->fops.func = ftrace_callback;
        hk->fops.flags = FTRACE_OPS_FL_SAVE_REGS | FTRACE_OPS_FL_RECURSION | FTRACE_OPS_FL_IPMODIFY; 
        
        if((error = ftrace_set_filter_ip(&hk->fops, hk->mcount_inst_addr, 0, 0))) {
            pr_debug("ftrace_set_filter_ip() failed with error code %d for function %s\n", error, hk->func_name);
            return -1;
        }

        hk->status_flag = FLAG_FILTERED; 

        if((error = register_ftrace_function(&hk->fops))) {
            pr_debug("register_ftrace_function failed with error code %d for function %s\n", error, hk->func_name);
            return -1;
        }

        hk->status_flag |= FLAG_TRACED;
        pr_debug("%s function hooked\n", hk->func_name); 
        *((unsigned long*)hk->orig_func) = hk->mcount_inst_addr;
    }
    return 0;
}

static int unhook_kfunction(struct fhook *hk, unsigned int hk_cnt)
{
    int error;
    pr_debug("unregistering hooks and removing any filters...");

    for(int i = 0; i < hk_cnt; i++, hk++) {
       if((hk->status_flag & FLAG_TRACED) && (error = unregister_ftrace_function(&hk->fops))) {
           pr_debug("unregister_ftrace_function() failed with error code %d for function %s\n", error, hk->func_name);
           return -1;  
       }

       if((hk->status_flag & FLAG_FILTERED) && (error = ftrace_set_filter_ip(&hk->fops, *((unsigned long*)hk->orig_func), 1, 0))) {
        pr_debug("ftrace_set_filter_ip() failed with error code %d for function %s\n", error, hk-> func_name);
        return -1; 
       }

       pr_debug("%s function unhooked\n", hk->func_name);
    }

    if(ibt_status)
        ibt_toggle();
    
    return 0;
}

static void ibt_toggle(void) 
{
    u64 msr_value;
    asm volatile (
        "rdmsr"
        : "=A" (msr_value)
        : "c" (MSR_IA_S_CET)
    );

    msr_value ^= (1ULL << IBT_BIT_POSITION);
    asm volatile (
        "wrmsr"
        :
        : "c" (MSR_IA_S_CET),
        "A" (msr_value)
    );    
}

static bool ibt_is_on(void) 
{
    u64 msr_value;
    if(boot_cpu_has(X86_FEATURE_IBT)) {
        asm volatile (
            "rdmsr"
            : "=A" (msr_value)
            : "c" (MSR_IA_S_CET)
        );
        return test_bit(IBT_BIT_POSITION, (unsigned long *)&msr_value);

    }
    return false;
}

#endif