[WIP] Fix multiple issues, make it work with the latest kernel and performance improvements.

Makefile

@@ -1,13 +1,22 @@
 
 .PHONY: all
 
+test: all
+	sudo rmmod afl_snapshot || echo "Not loaded anyways..."
+	sudo insmod src/afl_snapshot.ko
+	./test/test3
+
+
+
 all:
 	cd src && $(MAKE)
 	cd lib && $(MAKE)
+	cd test && $(MAKE)
 
 clean:
 	cd src && $(MAKE) clean
 	cd lib && $(MAKE) clean
+	cd test && $(MAKE)
 
 code-format:
 	./.custom-format.py -i src/*.c

README.md

@@ -29,6 +29,9 @@ however adding this snapshot module will still be a small improvement.
 |tiff|thumbnail|5058|3114|x1.6|
 |libxml|xmllint|7835|3450|x2.3|
 |afl++|test_persistent_new|106k|89k|x1.2|
+|afl++|emmu_fuzz|10k-20k|40|x250-x500|
+
+**TODO:** Rerun the others with improved version?
 
 ## Usage
 

load.sh

@@ -1,5 +1,7 @@
 #!/bin/sh
 
+set -e -o pipefail
+
 if [ '!' "$EUID" = 0 ] && [ '!' `id -u` = 0 ] ; then
 	echo "Warning: you need to be root to run this!"
 	# we do not exit as other mechanisms exist that allows to do this than
@@ -8,6 +10,6 @@ fi
 
 cd src/
 
-rmmod afl_snapshot
+rmmod afl_snapshot || echo "Not loaded anyways..."
 make
 insmod afl_snapshot.ko && echo Successfully loaded the snapshot module

src/Makefile

@@ -46,10 +46,9 @@ endif
 LINUX_DIR ?= /lib/modules/$(shell uname -r)/build
 
 .PHONY: all
+# 	env ARCH='$(ARCH)' LINUX_SYSTEM_MAP='$(LINUX_SYSTEM_MAP)' python3 lookup_symbols.py
 
 all:
-	env ARCH='$(ARCH)' LINUX_SYSTEM_MAP='$(LINUX_SYSTEM_MAP)' python3 lookup_symbols.py
-
 	$(MAKE) -C '$(LINUX_DIR)' M='$(M)' modules
 
 clean:

src/debug.h

@@ -54,5 +54,7 @@
 
 #endif
 
+#define PREEMPT_DEBUG(tag) SAYF("[%s():%s:%d] " tag " preempt_count() == %d\n", __FUNCTION__, __FILE__, __LINE__, preempt_count())
+
 #endif
 

src/files.c

@@ -57,7 +57,8 @@ void recover_files_snapshot(struct task_data *data) {
           DBG_PRINT("find new fds %d file* 0x%08lx\n", i, (unsigned long)file);
           // fdt->fd[i] = NULL;
           // filp_close(file, files);
-          __close_fd(files, i);
+          WARNF("closing doesn't work :(\n");
+          // __close_fd(files, i);
 
         }
 

src/ftrace_helper.h

@@ -0,0 +1,219 @@
+/*
+ * Helper library for ftrace hooking kernel functions
+ * Author: Harvey Phillips ([email protected])
+ * License: GPL
+ * */
+
+#include <linux/ftrace.h>
+#include <linux/linkage.h>
+#include <linux/slab.h>
+#include <linux/uaccess.h>
+#include <linux/version.h>
+#include "debug.h"
+#include "ftrace_util.h"
+
+#if defined(CONFIG_X86_64) && (LINUX_VERSION_CODE >= KERNEL_VERSION(4,17,0))
+#define PTREGS_SYSCALL_STUBS 1
+#endif
+
+/*
+ * On Linux kernels 5.7+, kallsyms_lookup_name() is no longer exported, 
+ * so we have to use kprobes to get the address.
+ * Full credit to @f0lg0 for the idea.
+ */
+#if LINUX_VERSION_CODE >= KERNEL_VERSION(5,7,0)
+#define KPROBE_LOOKUP 1
+#include <linux/kprobes.h>
+static struct kprobe kp = {
+    .symbol_name = "kallsyms_lookup_name"
+};
+typedef unsigned long (*kallsyms_lookup_name_t)(const char *name);
+kallsyms_lookup_name_t kallsyms_lookup_name_var;
+#define kallsyms_lookup_name kallsyms_lookup_name_var
+#endif
+
+/* x64 has to be special and require a different naming convention */
+#ifdef PTREGS_SYSCALL_STUBS
+#define SYSCALL_NAME(name) ("__x64_" name)
+#else
+#define SYSCALL_NAME(name) (name)
+#endif
+
+#define HOOK(_name, _hook, _orig)   \
+{                   \
+    .name = (_name),        \
+    .function = (_hook),        \
+    .original = (_orig),        \
+}
+
+#define SYSCALL_HOOK(_name, _hook, _orig)   \
+{                   \
+    .name = SYSCALL_NAME(_name),        \
+    .function = (_hook),        \
+    .original = (_orig),        \
+}
+
+
+/* We need to prevent recursive loops when hooking, otherwise the kernel will
+ * panic and hang. The options are to either detect recursion by looking at
+ * the function return address, or by jumping over the ftrace call. We use the 
+ * first option, by setting USE_FENTRY_OFFSET = 0, but could use the other by
+ * setting it to 1. (Oridinarily ftrace provides it's own protections against
+ * recursion, but it relies on saving return registers in $rip. We will likely
+ * need the use of the $rip register in our hook, so we have to disable this
+ * protection and implement our own).
+ * */
+#define USE_FENTRY_OFFSET 0
+#if !USE_FENTRY_OFFSET
+#pragma GCC optimize("-fno-optimize-sibling-calls")
+#endif
+
+/* We pack all the information we need (name, hooking function, original function)
+ * into this struct. This makes is easier for setting up the hook and just passing
+ * the entire struct off to fh_install_hook() later on.
+ * */
+struct ftrace_hook {
+    const char *name;
+    void *function;
+    void *original;
+
+    unsigned long address;
+    struct ftrace_ops ops;
+};
+
+/* Ftrace needs to know the address of the original function that we
+ * are going to hook. As before, we just use kallsyms_lookup_name() 
+ * to find the address in kernel memory.
+ * */
+static int fh_resolve_hook_address(struct ftrace_hook *hook)
+{
+#ifdef KPROBE_LOOKUP
+    register_kprobe(&kp);
+    kallsyms_lookup_name = (kallsyms_lookup_name_t) kp.addr;
+    unregister_kprobe(&kp);
+#endif
+    hook->address = kallsyms_lookup_name(hook->name);
+
+    if (!hook->address)
+    {
+        printk(KERN_DEBUG "rootkit: unresolved symbol: %s\n", hook->name);
+        return -ENOENT;
+    }
+
+#if USE_FENTRY_OFFSET
+    *((unsigned long*) hook->original) = hook->address + MCOUNT_INSN_SIZE;
+#else
+    *((unsigned long*) hook->original) = hook->address;
+#endif
+
+    return 0;
+}
+
+/* See comment below within fh_install_hook() */
+static void notrace fh_ftrace_thunk(unsigned long ip, unsigned long parent_ip, struct ftrace_ops *ops, ftrace_regs_ptr regs)
+{
+    struct ftrace_hook *hook = container_of(ops, struct ftrace_hook, ops);
+    struct pt_regs* pregs = ftrace_get_regs(regs);
+
+#if USE_FENTRY_OFFSET
+    pregs->ip = (unsigned long) hook->function;
+#else
+    if(!within_module(parent_ip, THIS_MODULE))
+        pregs->ip = (unsigned long) hook->function;
+#endif
+}
+
+/* Assuming we've already set hook->name, hook->function and hook->original, we 
+ * can go ahead and install the hook with ftrace. This is done by setting the 
+ * ops field of hook (see the comment below for more details), and then using
+ * the built-in ftrace_set_filter_ip() and register_ftrace_function() functions
+ * provided by ftrace.h
+ * */
+int fh_install_hook(struct ftrace_hook *hook)
+{
+    int err;
+    err = fh_resolve_hook_address(hook);
+    if(err)
+        return err;
+    SAYF("Successfully resolved address 0x%lx for function %s\n", hook->address, hook->name);
+
+    /* For many of function hooks (especially non-trivial ones), the $rip
+     * register gets modified, so we have to alert ftrace to this fact. This
+     * is the reason for the SAVE_REGS and IP_MODIFY flags. However, we also
+     * need to OR the RECURSION_SAFE flag (effectively turning if OFF) because
+     * the built-in anti-recursion guard provided by ftrace is useless if
+     * we're modifying $rip. This is why we have to implement our own checks
+     * (see USE_FENTRY_OFFSET). */
+    hook->ops.func = fh_ftrace_thunk;
+    hook->ops.flags = FTRACE_OPS_FL_SAVE_REGS
+            | FTRACE_OPS_FL_RECURSION_SAFE
+            | FTRACE_OPS_FL_IPMODIFY;
+
+    err = ftrace_set_filter_ip(&hook->ops, hook->address, 0, 0);
+    if(err)
+    {
+        printk(KERN_DEBUG "rootkit: ftrace_set_filter_ip() failed: %d\n", err);
+        return err;
+    }
+
+    err = register_ftrace_function(&hook->ops);
+    if(err)
+    {
+        printk(KERN_DEBUG "rootkit: register_ftrace_function() failed: %d\n", err);
+        return err;
+    }
+
+    return 0;
+}
+
+/* Disabling our function hook is just a simple matter of calling the built-in
+ * unregister_ftrace_function() and ftrace_set_filter_ip() functions (note the
+ * opposite order to that in fh_install_hook()).
+ * */
+void fh_remove_hook(struct ftrace_hook *hook)
+{
+    int err;
+    err = unregister_ftrace_function(&hook->ops);
+    if(err)
+    {
+        printk(KERN_DEBUG "rootkit: unregister_ftrace_function() failed: %d\n", err);
+    }
+
+    err = ftrace_set_filter_ip(&hook->ops, hook->address, 1, 0);
+    if(err)
+    {
+        printk(KERN_DEBUG "rootkit: ftrace_set_filter_ip() failed: %d\n", err);
+    }
+}
+
+/* To make it easier to hook multiple functions in one module, this provides
+ * a simple loop over an array of ftrace_hook struct
+ * */
+int fh_install_hooks(struct ftrace_hook *hooks, size_t count)
+{
+    int err;
+    size_t i;
+
+    for (i = 0 ; i < count ; i++)
+    {
+        err = fh_install_hook(&hooks[i]);
+        if(err)
+            goto error;
+    }
+    return 0;
+
+error:
+    while (i != 0)
+    {
+        fh_remove_hook(&hooks[--i]);
+    }
+    return err;
+}
+
+void fh_remove_hooks(struct ftrace_hook *hooks, size_t count)
+{
+    size_t i;
+
+    for (i = 0 ; i < count ; i++)
+        fh_remove_hook(&hooks[i]);
+}

src/ftrace_util.h

@@ -0,0 +1,18 @@
+#ifndef __FTRACE_UTIL_H
+#define __FTRACE_UTIL_H
+
+#include <linux/ftrace.h>
+#include <linux/version.h>
+
+// In 5.11+, ftrace hooks take ftrace_regs as argument.
+// Hacky way to fix this for older kernels.
+#if LINUX_VERSION_CODE < KERNEL_VERSION(5,11,0)
+typedef struct pt_regs* ftrace_regs_ptr;
+#define ftrace_get_regs(reg_ptr) reg_ptr;
+#define FTRACE_OPS_FL_RECURSION 0
+#else
+typedef struct ftrace_regs* ftrace_regs_ptr;
+#define FTRACE_OPS_FL_RECURSION_SAFE 0
+#endif
+
+#endif /* __FTRACE_UTIL_H */

src/hook.c

@@ -1,27 +1,35 @@
 #include <linux/kernel.h>
 #include <linux/kprobes.h>
+#include <linux/ftrace.h>
 #include <linux/module.h>
 #include <linux/slab.h>
-
+#include <linux/version.h>
+#include "debug.h"
+#include "ftrace_util.h"
 // TODO(andrea) switch from Kprobes to Ftrace
 
 struct hook {
 
   struct kprobe    kp;
+  struct ftrace_ops fops;
   struct list_head l;
 
 };
 
 LIST_HEAD(hooks);
 
 int try_hook(const char *func_name, void *handler) {
-
+  SAYF("Hooking function %s\n", func_name);
   struct hook *hook = kmalloc(sizeof(struct hook), GFP_KERNEL | __GFP_ZERO);
   INIT_LIST_HEAD(&hook->l);
   hook->kp.symbol_name = func_name;
   hook->kp.pre_handler = handler;
-
-  int ret = register_kprobe(&hook->kp);
+  hook->fops.flags = FTRACE_OPS_FL_SAVE_REGS | FTRACE_OPS_FL_IPMODIFY | FTRACE_OPS_FL_RECURSION;
+  hook->fops.func = handler;
+  ftrace_set_filter(&hook->fops, func_name, strlen(func_name), 0);
+  int ret = register_ftrace_function(&hook->fops);
+  SAYF("Hooked function: %d\n", ret);
+  // int ret = register_kprobe(&hook->kp);
   if (!ret) { list_add(&hook->l, &hooks); }
 
   return true;
@@ -35,7 +43,8 @@ void unhook(const char *func_name) {
 
     if (!strcmp(hook->kp.symbol_name, func_name)) {
 
-      unregister_kprobe(&hook->kp);
+      // unregister_kprobe(&hook->kp);
+      unregister_ftrace_function(&hook->fops);
 
     }
 
@@ -48,7 +57,8 @@ void unhook_all(void) {
   struct hook *hook = NULL;
   list_for_each_entry(hook, &hooks, l) {
 
-    unregister_kprobe(&hook->kp);
+    // unregister_kprobe(&hook->kp);
+    unregister_ftrace_function(&hook->fops);
 
   }
 

src/lookup_symbols.py

@@ -13,6 +13,8 @@
 
 system_map = map(lambda x: x.split(), fd.read().split('\n'))
 
+# print("system_map:", list(system_map))
+
 register_chrdev_region = None
 sys_call_table = None
 sys_read = None