From 2b283a9e3cb120b168fbd1b046d923f0899daafc Mon Sep 17 00:00:00 2001
From: Administrator <root@example.com>
Date: Mon, 16 Aug 2021 21:02:01 +0300
Subject: [PATCH 1/9] v1.1.0: 	Add ftrace support 	Add reflective symbols
 extractor (work on 5.10+) 	Fix horrible bug which fault on
 do_exit_group() because of invalid return size (long/int) -- try make it
 universal (reflective) 	Minimal security fixes like do NOT trying to
 insert LKM after building...

---
 README.md      |  22 ++++-
 load.sh        |   3 +-
 src/Makefile   |  32 ++----
 src/debug.h    |   2 +-
 src/hook.c     |   8 +-
 src/memory.c   | 112 +++++++++++++++------
 src/module.c   | 262 ++++++++++++++++++++++++++++---------------------
 src/snapshot.c |  13 +--
 src/snapshot.h |  17 ++--
 test/test1.c   |  37 ++++---
 10 files changed, 309 insertions(+), 199 deletions(-)

diff --git a/README.md b/README.md
index 1fe921e..7cd5ee9 100644
--- a/README.md
+++ b/README.md
@@ -32,9 +32,14 @@ however adding this snapshot module will still be a small improvement.
 
 ## Usage
 
-Load it using `./load.sh`, unload it using `./unload.sh`.
+!WARNING! This LKM is in alpha testing state. DO NOT LOAD IT ON YOU'RE REAL MACHINE WITHOUT TESTING!!!
+!DANGER!! It can crash the kernel and you will lose all you're unsaved data (open tabs, notes, etc.)
+
+At the moment it builds and run ok on at least:
+Debian buster Linux stand 4.19.160 #2 SMP Mon Dec 28 11:58:39 EET 2020 x86_64 GNU/Linux
+Debian bullseye Linux l0c4lh05t 5.10.24 #8 SMP Sun Jun 13 01:31:09 EEST 2021 x86_64 GNU/Linux
+Both on real hardware and under qemu vm.
 
-`./load.sh` will compile the module for you, you need also python3.
 
 While the module is loaded, [AFL++](https://github.com/AFLplusplus/AFLplusplus)
 will detect it and automatically switch from fork() to snapshot mode.
@@ -93,5 +98,14 @@ Remove the snapshot, you can not call `afl_snapshot_take` in another program poi
 
  + support for multithreaded applications
  + file descriptors state restore (lseek)
- + switch from kprobe to ftrace for hooking (faster)
- 
+ + switch from ftrace to jmp for hooking (faster)
+ + add support of task control from other task (can be achived via find_vpid(pid))
+
+### Chandgelog
+
+v1.1.0:
+	Add ftrace support
+	Add reflective symbols extractor (work on 5.10+)
+	Fix horrible bug which fault on do_exit_group() because of invalid return size (long/int) -- try make it universal (reflective)
+	Minimal security fixes like do NOT trying to insert LKM after building...
+v1.0.0 -- Initial release
diff --git a/load.sh b/load.sh
index 10787ba..27d6ee5 100755
--- a/load.sh
+++ b/load.sh
@@ -10,4 +10,5 @@ cd src/
 
 rmmod afl_snapshot
 make
-insmod afl_snapshot.ko && echo Successfully loaded the snapshot module
+echo "DO NOT INSERT THIS LKM IN YOU'RE REAL MACHINE WITHOUT TESTING! YOU CAN LOSE YOU'RE RUNTIME DATA!!!"
+#insmod afl_snapshot.ko && echo Successfully loaded the snapshot module
diff --git a/src/Makefile b/src/Makefile
index 739e698..47c4a08 100644
--- a/src/Makefile
+++ b/src/Makefile
@@ -9,46 +9,34 @@ else
   ARCH := ia32
 endif
 
-MAP = /boot/System.map-$(shell uname -r)
-SYMS = /proc/kallsyms
-
-ifndef LINUX_SYSTEM_MAP
-  ifdef KALLSYMS
-    override LINUX_SYSTEM_MAP = /proc/kallsyms
-  else
-    ifeq "$(shell test -r $(MAP) && echo 1)" "1"
-      override LINUX_SYSTEM_MAP = $(MAP)
-    else
-      ifeq "$(shell test -r $(SYMS) && echo 1)" "1"
-        override LINUX_SYSTEM_MAP = $(SYMS)
-      else
-        $(error no readable $(MAP) and no readable $(SYMS) found)
-      endif
-    endif
-  endif
-endif
+#MAP = /opt/DEBIAN/fs/boot/System.map-4.19.160
+#/boot/System.map-$(shell uname -r)
+#SYMS = /opt/DEBIAN/fs/proc/kallsyms
 
 obj-m += afl_snapshot.o
 afl_snapshot-objs := memory.o files.o threads.o task_data.o snapshot.o hook.o module.o
 
 ccflags-y := \
-  -ggdb3 \
+  -g3 -ggdb3 \
   -std=gnu99 \
-  -Wframe-larger-than=1000000000 \
+  -Wframe-larger-than=100000000 \
   -I$(M)/../include \
   -Wno-declaration-after-statement \
+  -fdata-sections \
   $(CCFLAGS)
 
 ifdef DEBUG
   ccflags-y += -DDEBUG
 endif
 
-LINUX_DIR ?= /lib/modules/$(shell uname -r)/build
+#LINUX_DIR ?= /lib/modules/$(shell uname -r)/build
+LINUX_DIR ?= /opt/DEBIAN/stuff/linux-4.19.160
+#/lib/modules/$(shell uname -r)/build
 
 .PHONY: all
 
 all:
-	env ARCH='$(ARCH)' LINUX_SYSTEM_MAP='$(LINUX_SYSTEM_MAP)' python3 lookup_symbols.py
+	#env ARCH='$(ARCH)' LINUX_SYSTEM_MAP='$(LINUX_SYSTEM_MAP)' python3 lookup_symbols.py
 
 	$(MAKE) -C '$(LINUX_DIR)' M='$(M)' modules
 
diff --git a/src/debug.h b/src/debug.h
index 2e3de44..5ff6036 100644
--- a/src/debug.h
+++ b/src/debug.h
@@ -5,7 +5,7 @@
 #include <linux/kern_levels.h>
 
 /* Output macros */
-
+#define DEBUG DEBUG
 #define HEXDUMP(type, prefix, ptr, size)                 \
   do {                                                   \
                                                          \
diff --git a/src/hook.c b/src/hook.c
index 9b39f79..7eb0eaf 100644
--- a/src/hook.c
+++ b/src/hook.c
@@ -1,10 +1,9 @@
+#ifdef USE_KPROBES
 #include <linux/kernel.h>
 #include <linux/kprobes.h>
 #include <linux/module.h>
 #include <linux/slab.h>
 
-// TODO(andrea) switch from Kprobes to Ftrace
-
 struct hook {
 
   struct kprobe    kp;
@@ -15,8 +14,8 @@ struct hook {
 LIST_HEAD(hooks);
 
 int try_hook(const char *func_name, void *handler) {
-
-  struct hook *hook = kmalloc(sizeof(struct hook), GFP_KERNEL | __GFP_ZERO);
+  /* Sometimes non-atomic allocations can fall, when calling from userspace (init) context */
+  struct hook *hook = kmalloc(sizeof(struct hook), GFP_ATOMIC | __GFP_ZERO);
   INIT_LIST_HEAD(&hook->l);
   hook->kp.symbol_name = func_name;
   hook->kp.pre_handler = handler;
@@ -53,3 +52,4 @@ void unhook_all(void) {
   }
 
 }
+#endif
diff --git a/src/memory.c b/src/memory.c
index 24b381d..babc7fc 100644
--- a/src/memory.c
+++ b/src/memory.c
@@ -5,7 +5,8 @@
 
 static DEFINE_PER_CPU(struct task_struct *, last_task) = NULL;
 static DEFINE_PER_CPU(struct task_data *, last_data) = NULL;
-
+extern long (*do_wp_page_orig)(long);
+extern long (*page_add_new_anon_rmap_orig)(long,long,long,long);
 pmd_t *get_page_pmd(unsigned long addr) {
 
   pgd_t *pgd;
@@ -66,7 +67,7 @@ pte_t *walk_page_table(unsigned long addr) {
   pgd = pgd_offset(mm, addr);
   if (pgd_none(*pgd) || pgd_bad(*pgd)) {
 
-    // DBG_PRINT("Invalid pgd.");
+    //DBG_PRINT("Invalid pgd.");
     goto out;
 
   }
@@ -74,7 +75,7 @@ pte_t *walk_page_table(unsigned long addr) {
   p4d = p4d_offset(pgd, addr);
   if (p4d_none(*p4d) || p4d_bad(*p4d)) {
 
-    // DBG_PRINT("Invalid p4d.");
+    //DBG_PRINT("Invalid p4d.");
     goto out;
 
   }
@@ -82,7 +83,7 @@ pte_t *walk_page_table(unsigned long addr) {
   pud = pud_offset(p4d, addr);
   if (pud_none(*pud) || pud_bad(*pud)) {
 
-    // DBG_PRINT("Invalid pud.");
+    //DBG_PRINT("Invalid pud.");
     goto out;
 
   }
@@ -90,7 +91,7 @@ pte_t *walk_page_table(unsigned long addr) {
   pmd = pmd_offset(pud, addr);
   if (pmd_none(*pmd) || pmd_bad(*pmd)) {
 
-    // DBG_PRINT("Invalid pmd.");
+    //DBG_PRINT("Invalid pmd.");
     goto out;
 
   }
@@ -98,7 +99,7 @@ pte_t *walk_page_table(unsigned long addr) {
   ptep = pte_offset_map(pmd, addr);
   if (!ptep) {
 
-    // DBG_PRINT("[NEW] Invalid pte.");
+    //DBG_PRINT("[NEW] Invalid pte.");
     goto out;
 
   }
@@ -107,6 +108,37 @@ pte_t *walk_page_table(unsigned long addr) {
   return ptep;
 
 }
+/* TODO: Allow control restoring from other task
+ * For example from AFLplusplus, or from forkserver
+***/
+pte_t *memwalk(uint64_t addr, struct mm_struct *mm){
+        pgd_t *pgd;// = 0x01;
+        p4d_t *p4d;// = 0x02;
+        pud_t *pud;// = 0x03;
+        pmd_t *pmd;// = 0x04;
+        pte_t *pte;// = 0x05;
+	//printk("memwalking(%#lx, %#lx\n", addr, mm);
+        pgd = pgd_offset((mm) ?: (current->mm), addr);
+        if(!pgd_none(*pgd) && !pgd_bad(*pgd)){
+                p4d = p4d_offset(pgd, addr);
+                if(!p4d_none(*p4d) && !p4d_bad(*p4d)){
+                        pud = pud_offset(p4d, addr);
+                        if(!pud_none(*pud) && !pud_bad(*pud)){
+                                pmd = pmd_offset(pud, addr);
+                                if(!pmd_none(*pmd) && !pmd_bad(*pmd)){
+                                        pte = pte_offset_map(pmd, addr);
+                                        if(!pte_none(*pte)){
+                                                //printk("pgd = %#llx, p4d = %#llx, pud = %#llx, pmd = %#llx, pte = %#llx\n",
+                                                //        *pgd, *p4d, *pud, *pmd, *pte);
+						return pte;
+                                        } printk("Failed to extract pte of addr+mm\n");
+                                        pte_unmap(pte);
+					return 0x00;
+                                }
+                        }
+                }
+        }
+}
 
 // TODO lock thee lists
 
@@ -242,7 +274,8 @@ void make_snapshot_page(struct task_data *data, struct mm_struct *mm,
   struct snapshot_page *sp;
   struct page *         page;
 
-  pte = walk_page_table(addr);
+  pte = memwalk(addr, ( data && data->tsk && data->tsk->mm ) ? data->tsk->mm : NULL );
+
   if (!pte) goto out;
 
   page = pte_page(*pte);
@@ -477,7 +510,14 @@ void do_recover_none_pte(struct snapshot_page *sp) {
 }
 
 void recover_memory_snapshot(struct task_data *data) {
-
+  if( !data ){
+      printk("WARNING: recover_memory_snapshot() called with NULL arg!\n");
+      return;
+  }
+  if( !data->tsk ){
+      printk("WARNING in recover_memory_snapshot() arg `data` don't have tsk field!\n");
+      return;
+  }
   struct snapshot_page *sp, *prev_sp = NULL;
   struct mm_struct *    mm = data->tsk->mm;
   pte_t *               pte, entry;
@@ -497,13 +537,14 @@ void recover_memory_snapshot(struct task_data *data) {
       if (pte) {
 
         /* Private rw page */
-        DBG_PRINT("private writable addr: 0x%08lx\n", sp->page_base);
+        DBG_PRINT("private writable addr: %#lx\n", sp->page_base);
         ptep_set_wrprotect(mm, sp->page_base, pte);
         set_snapshot_page_private(sp);
 
         /* flush tlb to make the pte change effective */
         k_flush_tlb_mm_range(mm, sp->page_base, sp->page_base + PAGE_SIZE,
                              PAGE_SHIFT, false);
+
         DBG_PRINT("writable now: %d\n", pte_write(*pte));
 
         pte_unmap(pte);
@@ -573,14 +614,24 @@ void clean_memory_snapshot(struct task_data *data) {
 
 }
 
+/*
 static long return_0_stub_func(void) {
 
   return 0;
 
-}
-
-int wp_page_hook(struct kprobe *p, struct pt_regs *regs) {
-
+}*/
+
+/* do_wp_page got seven args on <= 4.7 kernels...
+ * two args on == 4.8* kernels
+ * and one arg on >= 4.10 kernels...
+ * we can use #define to solve this
+ * or we can make universal, but expensive variant
+ *	and always push seven arguments to original function
+ * or we can just say that for afl-lkm u need 4.10+ kernel...
+***/
+asmlinkage int wp_page_hook(long rdi) {
+//( (( (void(*)() )(0xffffffff810a6140) )))();
+//printk("do_wp_page_orig == %#llx\n", (long*)do_wp_page_orig);
   struct vm_fault *     vmf;
   struct mm_struct *    mm;
   struct task_data *    data;
@@ -589,7 +640,7 @@ int wp_page_hook(struct kprobe *p, struct pt_regs *regs) {
   pte_t                 entry;
   char *                vfrom;
 
-  vmf = (struct vm_fault *)regs->di;
+  vmf = (struct vm_fault *)rdi;
   mm = vmf->vma->vm_mm;
   ss_page = NULL;
 
@@ -616,16 +667,16 @@ int wp_page_hook(struct kprobe *p, struct pt_regs *regs) {
 
   } else
 
-    return 0;  // continue
+    return do_wp_page_orig(rdi);//0;  // continue
 
   if (!ss_page) {
 
     // not a snapshot'ed page
-    return 0;  // continue
+    return do_wp_page_orig(rdi);//0;  // continue
 
   }
 
-  if (ss_page->dirty) return 0;
+  if (ss_page->dirty) return do_wp_page_orig(rdi);//0;
 
   ss_page->dirty = true;
 
@@ -674,28 +725,31 @@ int wp_page_hook(struct kprobe *p, struct pt_regs *regs) {
 
     pte_unmap_unlock(vmf->pte, vmf->ptl);
 
-    // skip original function
-    regs->ip = (long unsigned int)&return_0_stub_func;
-    return 1;
+    /* // skip original function
+     * regs->ip = (long unsigned int)&return_0_stub_func;
+     * return 1; */
 
-  }
+    /* since we using ftrace -- we don't need to touch $rip, yes? */
+    return 0x00;
 
-  return 0;  // continue
+  }
+  return do_wp_page_orig(rdi);  // continue
 
 }
 
 // actually hooking page_add_new_anon_rmap, but we really only care about calls
 // from do_anonymous_page
-int do_anonymous_hook(struct kprobe *p, struct pt_regs *regs) {
+asmlinkage int do_anonymous_hook(long rdi, long rsi, long rdx, long rcx) {
 
   struct vm_area_struct *vma;
   struct mm_struct *     mm;
   struct task_data *     data;
   struct snapshot_page * ss_page;
   unsigned long          address;
-
-  vma = (struct vm_area_struct *)regs->si;
-  address = regs->dx;
+//((((void(*)())(0xffffffff810a6140))))();
+//printk("page_add_new_anon_rmap_orig == %#llx\n", (long*)page_add_new_anon_rmap_orig);
+  vma = (struct vm_area_struct *)rsi;
+  address = (unsigned)rdx;
   mm = vma->vm_mm;
   ss_page = NULL;
 
@@ -722,14 +776,14 @@ int do_anonymous_hook(struct kprobe *p, struct pt_regs *regs) {
 
   } else {
 
-    return 0;
+    return page_add_new_anon_rmap_orig(rdi, rsi, rdx, rcx);//0;
 
   }
 
   if (!ss_page) {
 
     /* not a snapshot'ed page */
-    return 0;
+    return page_add_new_anon_rmap_orig(rdi, rsi, rdx, rcx);//0;
 
   }
 
@@ -738,7 +792,7 @@ int do_anonymous_hook(struct kprobe *p, struct pt_regs *regs) {
   // HAVE PTE NOW
   ss_page->has_had_pte = true;
 
-  return 0;
+  return page_add_new_anon_rmap_orig(rdi, rsi, rdx, rcx);//0;
 
 }
 
diff --git a/src/module.c b/src/module.c
index e3303b6..a5285ff 100644
--- a/src/module.c
+++ b/src/module.c
@@ -12,12 +12,14 @@
 #include <linux/uaccess.h>
 #include <linux/kallsyms.h>
 #include <linux/version.h>
+#include <linux/stop_machine.h>
 
 #include "task_data.h"  // mm associated data
-#include "hook.h"       // function hooking
+#ifdef USE_KPROBES
+# include "hook.h"       // function hooking
+#endif
 #include "snapshot.h"   // main implementation
 #include "debug.h"
-#include "symbols.h"
 
 #include "afl_snapshot.h"
 
@@ -25,9 +27,32 @@
 #define CLASS_NAME "afl_snapshot"
 
 MODULE_LICENSE("GPL");
-MODULE_AUTHOR("kallsyms & andreafioraldi");
+MODULE_AUTHOR("andreafioraldi && Th3C4t");
 MODULE_DESCRIPTION("Fast process snapshots for fuzzing");
-MODULE_VERSION("1.0.0");
+MODULE_VERSION("1.1.0");
+#ifndef __used
+# define __used		__attribute__((used))
+#endif
+#ifndef __unused
+# define __unused	__attribute__((unused))
+#endif
+#ifndef __weak
+# define __weak		__attribute__((weak))
+#endif
+
+#ifdef WITHOUT_KERNEL_HEADERS
+ extern  __weak uint64_t kallsyms_lookup_name;
+#else
+# if LINUX_VERSION_CODE > KERNEL_VERSION(4,19,0)
+#  warning "Reflective mod enable"
+# else
+   extern  typeof(kallsyms_lookup_name) kallsyms_lookup_name;
+# endif
+#endif
+
+typedef typeof(void *(*)(void *rdi, ...)) func_t;
+func_t	_stop_machine = ((typeof(_stop_machine))&stop_machine);
+func_t	_kallsyms_lookup_name = NULL;
 
 void (*k_flush_tlb_mm_range)(struct mm_struct *mm, unsigned long start,
                              unsigned long end, unsigned int stride_shift,
@@ -129,58 +154,32 @@ static struct file_operations dev_fops = {
 
 };
 
-#ifdef ARCH_HAS_SYSCALL_WRAPPER
-typedef int (*syscall_handler_t)(struct pt_regs *);
-
-// The original syscall handler that we removed to override exit_group()
+typedef long (*syscall_handler_t)(long rdi);
 syscall_handler_t orig_sct_exit_group = NULL;
 
-// TODO: non-x86 architectures syscall_table entries don't take pt_regs,
-// they take normal args
-// https://grok.osiris.cyber.nyu.edu/xref/linux/include/linux/syscalls.h?r=83fa805b#235
-// but x86 is (of course) different, taking a pt_regs, then passing extracted
-// values to the actual __do_sys*
-// https://grok.osiris.cyber.nyu.edu/xref/linux/arch/x86/include/asm/syscall_wrapper.h?r=6e484764#161
-
-asmlinkage int sys_exit_group(struct pt_regs *regs) {
-
-  if (exit_snapshot()) return orig_sct_exit_group(regs);
-
-  return 0;
-
+asmlinkage long sys_exit_group(long rdi) {
+  return exit_snapshot() ? orig_sct_exit_group(rdi) : 0x00;
 }
-#else
-typedef long (*syscall_handler_t)(int error_code);
-
-// The original syscall handler that we removed to override exit_group()
-syscall_handler_t orig_sct_exit_group = NULL;
-
-asmlinkage long sys_exit_group(int error_code) {
 
-  if (exit_snapshot()) return orig_sct_exit_group(error_code);
-
-  return 0;
-
-}
-#endif
 #if LINUX_VERSION_CODE >= KERNEL_VERSION(5,8,0) /* rename since Linux 5.8 */
 #define probe_kernel_read copy_from_kernel_nofault
 #endif
 
+static uint64_t *_sys_call_table = NULL;
+
+/* :XXX: if this really needed? */
 static void **get_syscall_table(void) {
 
   void **syscall_table = NULL;
 
-  syscall_table = (void**)SYMADDR_sys_call_table;
-
-  if (syscall_table) { return syscall_table; }
-
   int                i;
-  unsigned long long s0 = SYMADDR_sys_read;
-  unsigned long long s1 = SYMADDR_sys_read;
+  unsigned long long s0 = (typeof(s0))_kallsyms_lookup_name("sys_read") ?:
+	 (typeof(s0))_kallsyms_lookup_name("__x64_sys_read");
+  unsigned long long s1 = s0;
+  if( !s1 || !s0 ) return NULL;
 
-  unsigned long long *data =
-      (unsigned long long *)(SYMADDR__etext & ~0x7);
+  unsigned long long *data = (uint64_t*)(((uint64_t)_kallsyms_lookup_name("_etext"))&~0x7);
+  //      (unsigned long long *)(SYMADDR__etext & ~0x7);
   for (i = 0; (unsigned long long)(&data[i]) < ULLONG_MAX; i++) {
 
     unsigned long long d;
@@ -200,59 +199,114 @@ static void **get_syscall_table(void) {
 
 }
 
-static void _write_cr0(unsigned long val) {
-
-  asm volatile("mov %0,%%cr0" : "+r"(val));
-
-}
-
-static void enable_write_protection(void) {
-
-  _write_cr0(read_cr0() | (1 << 16));
-
+static void __used disable_write_protection(volatile void *addr){
+   unsigned int level;
+   uint64_t *pte = (uint64_t*)lookup_address((unsigned long)addr, &level);
+   if (pte[0x00] & ~((uint64_t)1<<1)) pte[0x00] |= ((unsigned long long int)1<<1);
 }
 
-static void disable_write_protection(void) {
-
-  _write_cr0(read_cr0() & (~(1 << 16)));
-
+static void __used enable_write_protection(volatile void *addr){
+   unsigned int level;
+   uint64_t *pte = (uint64_t*)lookup_address((unsigned long)addr, &level);
+   pte[0x00] = pte[0x00] & ~((uint64_t)1<<1);
 }
 
-static void **syscall_table_ptr;
-
 static void patch_syscall_table(void) {
-
-  disable_write_protection();
-  orig_sct_exit_group = syscall_table_ptr[__NR_exit_group];
-  syscall_table_ptr[__NR_exit_group] = &sys_exit_group;
-  enable_write_protection();
+  __asm__("cli;   \n\t");
+  disable_write_protection(_sys_call_table);
+  orig_sct_exit_group = (void*)_sys_call_table[__NR_exit_group];
+  _sys_call_table[__NR_exit_group] = (uint64_t)&sys_exit_group;
+  enable_write_protection(_sys_call_table);
+  __asm__("sti;   \n\t");
 
 }
 
 static void unpatch_syscall_table(void) {
+  __asm__("cli;   \n\t");
+  disable_write_protection(_sys_call_table);
+  _sys_call_table[__NR_exit_group] = (uint64_t)orig_sct_exit_group;
+  enable_write_protection(_sys_call_table);
+  __asm__("sti;   \n\t");
+}
+#ifndef USE_KPROBES
+# include "ftrace_helper.c"
+#endif
 
-  disable_write_protection();
-  syscall_table_ptr[__NR_exit_group] = orig_sct_exit_group;
-  enable_write_protection();
+extern __weak __unused int do_wp_page(struct vm_fault *vmf); /* it doesn't have headers? */
+asmlinkage typeof(&page_add_new_anon_rmap) page_add_new_anon_rmap_orig = NULL;
+asmlinkage typeof(&do_wp_page) do_wp_page_orig = NULL;
+asmlinkage typeof(&do_exit) do_exit_orig = NULL;
+static struct ftrace_hook hooks[] = {
+        HOOK("page_add_new_anon_rmap", do_anonymous_hook, &page_add_new_anon_rmap_orig),
+        HOOK("do_wp_page", wp_page_hook, &do_wp_page_orig),
+	HOOK("do_exit", exit_hook, &do_exit_orig),
+};
 
-}
 
 int snapshot_initialize_k_funcs() {
+  func_t chek_kaddr = (void*)_kallsyms_lookup_name("func_ptr_is_kernel_text") ?:
+	(void*)_kallsyms_lookup_name("kernel_text_address");
 
-  k_flush_tlb_mm_range = (void *)SYMADDR_flush_tlb_mm_range;
-  k_zap_page_range = (void *)SYMADDR_zap_page_range;
+  k_flush_tlb_mm_range = (void *)_kallsyms_lookup_name("flush_tlb_mm_range");
+  k_zap_page_range = (void *)_kallsyms_lookup_name("zap_page_range");
 
-  if (!k_flush_tlb_mm_range || !k_zap_page_range) { return -ENOENT; }
+  if (!chek_kaddr(k_flush_tlb_mm_range) ||
+	!chek_kaddr(k_zap_page_range)) { return -ENOENT; }
 
-  SAYF("All loaded");
+  SAYF("All loaded :3");
 
   return 0;
 
 }
 
-static int __init mod_init(void) {
+#define NEEDLE "kallsyms_lookup_name+0x0"
+/* pre_init -- func, which ask Mr. Torvald to give us needed adresses
+ * $rax == ptr to kallsyms_lookup_name
+***/
+static void *pre_init(void){
+	u_int64_t before = 0x00, after = 0x00;
+	char kln2buf[0x200] = { 0x00 };
+
+	before = (u_int64_t)sprint_symbol - 0x1000;
+	after  = (u_int64_t)sprint_symbol + 0x1000;
+
+	for ( u_int64_t now = before; now <= after; now++ ){
+	        sprint_symbol(kln2buf, now);
+	        if ( strncmp(kln2buf, NEEDLE, sizeof(NEEDLE)-1 ) == 0x00 ){
+			SAYF("Thank you, Mr. Torvald :)");
+                        /* is KLN() instrumented via 5b NOP for ftrace() in prologue? */
+	                return ( *(uint32_t*)now == (uint32_t)0x00441f0f ) ?
+	                        (((void*)now) + 0x05) : (void*)now;
+	        } kln2buf[0x00] = 0x00;
+	}
+	return 0x00;
+}
+
+volatile static unsigned long long int __used slim_white_line(volatile void *args){
+	/* doesn't work under the void state */
+        //int ok = !!!fh_install_hooks(hooks, ARRAY_SIZE(hooks));
+        //if( !ok ) return !ok;
+        patch_syscall_table();
+        return 0x00;//ok;
+}
+static int __attribute__((section(".text"))) mod_init(void) {
+
+  int ok = 0x00;
+
+  //printk("%#llx, %#llx, %#llx\n", &page_add_new_anon_rmap_orig, &do_wp_page_orig, &do_exit_orig);
 
   SAYF("Loading AFL++ snapshot LKM");
+#if LINUX_VERSION_CODE > KERNEL_VERSION(4,19,0)
+  _kallsyms_lookup_name = pre_init();
+#else
+  _kallsyms_lookup_name = kallsyms_lookup_name;
+#endif
+  if( !_kallsyms_lookup_name ){ /* should not happend */
+    printk("Very very bad, I don't have kallsyms_lookup_name()\n"
+		"Don't look at me, go and do something!\n");
+	 return -ENOENT;
+  }
+  _stop_machine = _kallsyms_lookup_name("stop_machine");
 
   mod_kobj = kobject_create_and_add("afl_snapshot", kernel_kobj);
   if (!mod_kobj) return -ENOMEM;
@@ -291,65 +345,47 @@ static int __init mod_init(void) {
 
   SAYF("The major device number is %d", mod_major_num);
 
-  // syscall_table overwrites
-  syscall_table_ptr = get_syscall_table();
-  if (!syscall_table_ptr) {
+  // syscall_table prepare
+  _sys_call_table = _kallsyms_lookup_name("sys_call_table") ?: get_syscall_table();
+  if (!_sys_call_table) {
 
     FATAL("Unable to locate syscall_table");
     return -ENOENT;
 
   }
 
-  patch_syscall_table();
-
-  // func hooks
-  if (!try_hook("do_wp_page", &wp_page_hook)) {
-
-    FATAL("Unable to hook do_wp_page");
-    unpatch_syscall_table();
-
-    return -ENOENT;
-
-  }
-
-  if (!try_hook("page_add_new_anon_rmap", &do_anonymous_hook)) {
-
-    FATAL("Unable to hook page_add_new_anon_rmap");
-
-    unhook_all();
-    unpatch_syscall_table();
-    return -ENOENT;
-
-  }
-
-  if (!try_hook("do_exit", &exit_hook)) {
-
-    FATAL("Unable to hook do_exit");
-
-    unhook_all();
-    unpatch_syscall_table();
-    return -ENOENT;
+  /* breakpoint for GDB */
+  //( ( void(*)(void) ) (long)(_kallsyms_lookup_name("sys_ni_syscall")) ) ();
 
-  }
+  (ok = !!!_stop_machine(slim_white_line, NULL, NULL)) ?
+	printk("Well done!\n") : printk("Failed :(\n");
+  if( !ok ) return -ok;
 
-  // initialize snapshot non-exported funcs
+  fh_install_hooks(hooks, ARRAY_SIZE(hooks));
   return snapshot_initialize_k_funcs();
-
+  // initialize snapshot non-exported funcs
 }
-
-static void __exit mod_exit(void) {
+volatile static unsigned long long int __used fat_black_line(volatile void *args){
+	//fh_remove_hooks(hooks, ARRAY_SIZE(hooks));
+	unpatch_syscall_table();
+	return 0x00;
+}
+static void __attribute__((section(".text"))) mod_exit(void) {
 
   SAYF("Unloading AFL++ snapshot LKM");
+  fh_remove_hooks(hooks, ARRAY_SIZE(hooks));
+  _stop_machine(fat_black_line, NULL, NULL);
 
   kobject_put(mod_kobj);
-
+  printk("Objects dropped\n");
   device_destroy(mod_class, MKDEV(mod_major_num, 0));
+  printk("Devices dropped\n");
   class_unregister(mod_class);
   class_destroy(mod_class);
+  printk("Classes dropped\n");
   unregister_chrdev(mod_major_num, DEVICE_NAME);
-
-  unhook_all();
-  unpatch_syscall_table();
+  printk("Chardev dropped\n");
+  SAYF("[+] We're done here. Have a nice day!");
 
 }
 
diff --git a/src/snapshot.c b/src/snapshot.c
index faa32f3..5ebc6b1 100644
--- a/src/snapshot.c
+++ b/src/snapshot.c
@@ -3,12 +3,12 @@
 #include "task_data.h"
 #include "snapshot.h"
 
-int exit_hook(struct kprobe *p, struct pt_regs *regs) {
-
+/*int exit_hook(struct kprobe *p, struct pt_regs *regs) {*/
+asmlinkage long exit_hook(long rdi) {
   clean_snapshot();
-
-  return 0;
-
+//  printk("do_exit_orig == %#llx\n", (long*)do_exit_orig);
+  do_exit_orig(rdi); /* no return */
+  return ((long(*)())(0xdeadbeef))(); /* to keep gcc happy */
 }
 
 void initialize_snapshot(struct task_data *data, int config) {
@@ -86,7 +86,7 @@ void recover_snapshot(void) {
 
 }
 
-int exit_snapshot(void) {
+long exit_snapshot(void) {
 
   struct task_data *data = get_task_data(current);
   if (data && (data->config & AFL_SNAPSHOT_EXIT) && have_snapshot(data)) {
@@ -102,6 +102,7 @@ int exit_snapshot(void) {
 
 }
 
+/* :TODO: this should be static inline at least? */
 void clean_snapshot(void) {
 
   struct task_data *data = get_task_data(current);
diff --git a/src/snapshot.h b/src/snapshot.h
index afb03f2..52ef4bd 100644
--- a/src/snapshot.h
+++ b/src/snapshot.h
@@ -1,6 +1,6 @@
 #ifndef __AFL_SNAPSHOT_H__
 #define __AFL_SNAPSHOT_H__
-
+#include <linux/kernel.h>
 #include <asm/mmu.h>
 #include <asm/page.h>
 #include <linux/auxvec.h>
@@ -89,8 +89,8 @@
 #include <asm/pgalloc.h>
 #include <asm/pgtable.h>
 #include <asm/pgtable_types.h>
-#include <asm/tlb.h>
-#include <asm/tlbflush.h>
+//#include <asm/tlb.h>
+//#include <asm/tlbflush.h>
 #include <asm/uaccess.h>
 
 #include "afl_snapshot.h"
@@ -205,14 +205,15 @@ void recover_threads_snapshot(struct task_data *data);
 
 int snapshot_initialize_k_funcs(void);
 
-int wp_page_hook(struct kprobe *p, struct pt_regs *regs);
-int do_anonymous_hook(struct kprobe *p, struct pt_regs *regs);
-int exit_hook(struct kprobe *p, struct pt_regs *regs);
+int wp_page_hook(long rdi);//(struct kprobe *p, struct pt_regs *regs);
+int do_anonymous_hook(long rdi, long rsi, long rdx, long rcx);//(struct kprobe *p, struct pt_regs *regs);
+long exit_hook(long rdi);//(struct kprobe *p, struct pt_regs *regs);
+extern typeof(&do_exit) do_exit_orig;
 
 int  take_snapshot(int config);
 void recover_snapshot(void);
-void clean_snapshot(void);
-int  exit_snapshot(void);
+inline void clean_snapshot(void);
+long  exit_snapshot(void);
 
 void exclude_vmrange(unsigned long start, unsigned long end);
 void include_vmrange(unsigned long start, unsigned long end);
diff --git a/test/test1.c b/test/test1.c
index 6681d16..b6e6a79 100644
--- a/test/test1.c
+++ b/test/test1.c
@@ -1,4 +1,5 @@
 #define _GNU_SOURCE
+#include <stdint.h>
 #include <fcntl.h>
 #include <malloc.h>
 #include <setjmp.h>
@@ -17,22 +18,26 @@
 
 int* shm_addr;
 int* none_addr;
-
+void *chunk;
 int pippo = 1;
 
 void test2() {
-
+  *shm_addr = 0xffffffff;
   if (afl_snapshot_take(AFL_SNAPSHOT_NOSTACK) == 1)
     fprintf(stderr, "first time!\n");
 
 loop:
-
+  memset(chunk, 0x43, 0x1001);
+  fprintf(stderr, "*chunk = %#llx\n", *(uint64_t*)chunk);
   *none_addr += 1;
-  *shm_addr += 1;
-  fprintf(stderr, ">> %d     %p = %d    %p = %d\n", pippo, none_addr, *none_addr, shm_addr, *shm_addr);
+  fprintf(stderr, ">> %d     %p = %#llx    %p = %#llx\n", pippo, none_addr, *none_addr, shm_addr, *shm_addr);
+  *shm_addr = 0xdeadbeef;
   ++pippo;
-
+  memset(chunk, 0x44, 0x1001);
+  fprintf(stderr, "*chunk = %#llx\n", *(uint64_t*)chunk);
   afl_snapshot_restore();
+  fprintf(stderr, "*chunk = %#llx\n", *(uint64_t*)chunk);
+  if( *none_addr > 0x100 ) exit(0x00);
   goto loop;
 
 }
@@ -40,18 +45,28 @@ void test2() {
 int main() {
 
   afl_snapshot_init();
-  
+  puts("1");
   shm_addr = mmap(0, 0x10000, PROT_READ | PROT_WRITE | PROT_EXEC,
              MAP_SHARED | MAP_ANONYMOUS, 0, 0);
-
+  memset(shm_addr, 0x41, 0x10000);
+  puts("2");
   none_addr = mmap((void *)0, 0x1000, PROT_READ | PROT_WRITE,
               MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
-
+  memset(none_addr, 0xff, 0x1000);
+  puts("3");
   afl_snapshot_exclude_vmrange((unsigned long)none_addr, (unsigned long)(none_addr + (0x1000/4)));
+  puts("4");
   afl_snapshot_include_vmrange((unsigned long)shm_addr, (unsigned long)(shm_addr + (0x10000/4)));
-
+  puts("5");
   *shm_addr = 0;
-
+  puts("6");
+  chunk = malloc(0x1001);
+  memset(chunk, 0x41, 0x1001);
+  fprintf(stderr, "*chunk = %#llx\n", *(uint64_t*)chunk);
+  afl_snapshot_include_vmrange(chunk, chunk+0x1001);
+  memset(chunk, 0x42, 0x1001);
+  fprintf(stderr, "*chunk = %#llx\n", *(uint64_t*)chunk);
+  puts("7");
   test2();
 
   return 0;

From 32986b3f2261b379dc658b6fcdb07499b18b53d1 Mon Sep 17 00:00:00 2001
From: Administrator <root@example.com>
Date: Mon, 16 Aug 2021 21:05:18 +0300
Subject: [PATCH 2/9] v1.1.0: 	Add ftrace support 	Add reflective symbols
 extractor (work on 5.10+) 	Fix horrible bug which fault on
 do_exit_group() because of invalid return size (long/int) -- try make it
 universal (reflective) 	Minimal security fixes like do NOT trying to
 insert LKM after building...

---
 src/lookup_symbols.py | 64 -------------------------------------------
 1 file changed, 64 deletions(-)
 delete mode 100644 src/lookup_symbols.py

diff --git a/src/lookup_symbols.py b/src/lookup_symbols.py
deleted file mode 100644
index f7e8da9..0000000
--- a/src/lookup_symbols.py
+++ /dev/null
@@ -1,64 +0,0 @@
-#!/usr/bin/env python3
-
-import os
-
-arch = os.getenv("ARCH")
-system_map_fname = os.getenv("LINUX_SYSTEM_MAP")
-assert system_map_fname and "Specify the LINUX_SYSTEM_MAP env var with the path to your current kernel system map"
-
-try:
-    fd = open(system_map_fname)
-except:
-    raise RuntimeError(system_map_fname + ' not found, please specify another system map file using the LINUX_SYSTEM_MAP env var')
-
-system_map = map(lambda x: x.split(), fd.read().split('\n'))
-
-register_chrdev_region = None
-sys_call_table = None
-sys_read = None
-sys_write = None
-flush_tlb_mm_range = None
-zap_page_range = None
-_etext = None
-
-for e in system_map:
-    if len(e) < 3: continue
-    if e[2] == 'register_chrdev_region':
-        register_chrdev_region = int(e[0], 16)
-    elif e[2] == 'sys_call_table':
-        sys_call_table = int(e[0], 16)
-    elif e[2] == '__' + arch + '_sys_read' or e[2] == 'sys_read':
-        sys_read = int(e[0], 16)
-    elif e[2] == '__' + arch + '_sys_write' or e[2] == 'sys_write':
-        sys_write = int(e[0], 16)
-    elif e[2] == 'flush_tlb_mm_range':
-        flush_tlb_mm_range = int(e[0], 16)
-    elif e[2] == 'zap_page_range':
-        zap_page_range = int(e[0], 16)
-    elif e[2] == '_etext':
-        _etext = int(e[0], 16)
-
-assert register_chrdev_region != None
-assert sys_call_table != None
-assert sys_read != None
-assert sys_write != None
-assert flush_tlb_mm_range != None
-assert zap_page_range != None
-assert _etext != None
-
-with open('symbols.h', 'w') as f:
-    f.write('''#ifndef __AFL_SNAPSHOT_SYMBOLS_H__
-#define __AFL_SNAPSHOT_SYMBOLS_H__
-
-#define SYMADDR(offset) ((unsigned long)register_chrdev_region - \\
-                         0x%x + (offset))
-
-#define SYMADDR_sys_call_table SYMADDR(0x%x)
-#define SYMADDR_sys_read SYMADDR(0x%x)
-#define SYMADDR_sys_write SYMADDR(0x%x)
-#define SYMADDR_flush_tlb_mm_range SYMADDR(0x%x)
-#define SYMADDR_zap_page_range SYMADDR(0x%x)
-#define SYMADDR__etext SYMADDR(0x%x)
-
-#endif
-''' % (register_chrdev_region, sys_call_table, sys_read, sys_write, flush_tlb_mm_range, zap_page_range, _etext))

From bdfad147b29fb34ed875233d33a9a9de85396833 Mon Sep 17 00:00:00 2001
From: Administrator <root@example.com>
Date: Mon, 16 Aug 2021 21:07:49 +0300
Subject: [PATCH 3/9] v1.1.0: 	Add ftrace support 	Add reflective symbols
 extractor (work on 5.10+) 	Fix horrible bug which fault on
 do_exit_group() because of invalid return size (long/int) -- try make it
 universal (reflective) 	Minimal security fixes like do NOT trying to
 insert LKM after building...

---
 src/ftrace_helper.c | 175 ++++++++++++++++++++++++++++++++++++++++++++
 test/make.sh        |   1 +
 2 files changed, 176 insertions(+)
 create mode 100644 src/ftrace_helper.c
 create mode 100755 test/make.sh

diff --git a/src/ftrace_helper.c b/src/ftrace_helper.c
new file mode 100644
index 0000000..e034954
--- /dev/null
+++ b/src/ftrace_helper.c
@@ -0,0 +1,175 @@
+
+// (Th3C4t) Switch from ftrace to jmp via prologue patch
+// e9 3b c1 ff 7d           jmp    0x7dffc140
+
+
+
+/*
+ * Helper library for ftrace hooking kernel functions
+ * Author: Harvey Phillips (xcellerator@gmx.com)
+ * License: GPL
+ * */
+
+#include <linux/ftrace.h>
+#include <linux/linkage.h>
+#include <linux/slab.h>
+#include <linux/uaccess.h>
+
+#define HOOK(_name, _hook, _orig)	\
+{					\
+	.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)
+{
+	hook->address = (long unsigned)_kallsyms_lookup_name((void*)hook->name);
+
+	if (!hook->address)
+	{
+		printk(KERN_DEBUG "AFL-kit: 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;*/ ( *(uint32_t*)hook->address == (uint32_t)0x00441f0f ) ?
+                                (((uint64_t)hook->address) + 0x05) : (uint64_t)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, struct pt_regs *regs)
+{
+	struct ftrace_hook *hook = container_of(ops, struct ftrace_hook, ops);
+
+#if USE_FENTRY_OFFSET
+	regs->ip = (unsigned long) hook->function;
+#else
+	if(!within_module(parent_ip, THIS_MODULE))
+		regs->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;
+
+	/* 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 "AFL-kit: ftrace_set_filter_ip() failed: %d\n", err);
+		return err;
+	}
+
+	err = register_ftrace_function(&hook->ops);
+	if(err){
+		printk(KERN_DEBUG "AFL-kit: 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;
+	*((unsigned long*) hook->original) = hook->address;
+	err = unregister_ftrace_function(&hook->ops);
+	if(err)
+		printk(KERN_DEBUG "AFL-kit: unregister_ftrace_function() failed: %d\n", err);
+
+	err = ftrace_set_filter_ip(&hook->ops, hook->address, 1, 0);
+	if(err)
+		printk(KERN_DEBUG "AFL-kit: 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]);
+}
diff --git a/test/make.sh b/test/make.sh
new file mode 100755
index 0000000..66633a4
--- /dev/null
+++ b/test/make.sh
@@ -0,0 +1 @@
+cc -I../include test1.c ../lib/libaflsnapshot.o -o test1

From b691deb4ca9c6d58937ed5646f1ee74d26447316 Mon Sep 17 00:00:00 2001
From: Administrator <root@example.com>
Date: Mon, 16 Aug 2021 21:59:51 +0300
Subject: [PATCH 4/9] v1.1.0 alpha release

---
 README.md | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index 7cd5ee9..ed72bad 100644
--- a/README.md
+++ b/README.md
@@ -104,8 +104,8 @@ Remove the snapshot, you can not call `afl_snapshot_take` in another program poi
 ### Chandgelog
 
 v1.1.0:
-	Add ftrace support
-	Add reflective symbols extractor (work on 5.10+)
-	Fix horrible bug which fault on do_exit_group() because of invalid return size (long/int) -- try make it universal (reflective)
-	Minimal security fixes like do NOT trying to insert LKM after building...
++	Add ftrace support
++	Add reflective symbols extractor (work on 5.10+)
++	Fix horrible bug which fault on do_exit_group() because of invalid return size (long/int) -- try make it universal (reflective)
++	Minimal security fixes like do NOT trying to insert LKM after building...
 v1.0.0 -- Initial release

From 8ca317fcb356d23ff4edee6b7924b1d29a3c2e06 Mon Sep 17 00:00:00 2001
From: Administrator <root@example.com>
Date: Mon, 16 Aug 2021 22:01:39 +0300
Subject: [PATCH 5/9] v1.1.0 -- alpha release

---
 README.md | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index ed72bad..6dbb204 100644
--- a/README.md
+++ b/README.md
@@ -99,7 +99,7 @@ Remove the snapshot, you can not call `afl_snapshot_take` in another program poi
  + support for multithreaded applications
  + file descriptors state restore (lseek)
  + switch from ftrace to jmp for hooking (faster)
- + add support of task control from other task (can be achived via find_vpid(pid))
+ + add support of tasks snapshot control from other process (can be achived via find_vpid(pid))
 
 ### Chandgelog
 
@@ -108,4 +108,6 @@ v1.1.0:
 +	Add reflective symbols extractor (work on 5.10+)
 +	Fix horrible bug which fault on do_exit_group() because of invalid return size (long/int) -- try make it universal (reflective)
 +	Minimal security fixes like do NOT trying to insert LKM after building...
-v1.0.0 -- Initial release
+
+v1.0.0:
++	Initial release

From 8a9fdd14e84781fefa2bfcde43ec740b6e28e079 Mon Sep 17 00:00:00 2001
From: Administrator <root@example.com>
Date: Tue, 17 Aug 2021 00:12:36 +0300
Subject: [PATCH 6/9] minimal typo fixes

---
 src/Makefile |  6 ++++--
 src/debug.h  |  2 +-
 src/memory.c |  2 +-
 test/test1.c | 21 +++++++++++----------
 4 files changed, 17 insertions(+), 14 deletions(-)

diff --git a/src/Makefile b/src/Makefile
index 47c4a08..f90a0f5 100644
--- a/src/Makefile
+++ b/src/Makefile
@@ -16,14 +16,16 @@ endif
 obj-m += afl_snapshot.o
 afl_snapshot-objs := memory.o files.o threads.o task_data.o snapshot.o hook.o module.o
 
+
 ccflags-y := \
-  -g3 -ggdb3 \
   -std=gnu99 \
   -Wframe-larger-than=100000000 \
   -I$(M)/../include \
   -Wno-declaration-after-statement \
   -fdata-sections \
-  $(CCFLAGS)
+  $(CCFLAGS) \
+  -O2
+#  -g3 -ggdb3 
 
 ifdef DEBUG
   ccflags-y += -DDEBUG
diff --git a/src/debug.h b/src/debug.h
index 5ff6036..5517b58 100644
--- a/src/debug.h
+++ b/src/debug.h
@@ -5,7 +5,7 @@
 #include <linux/kern_levels.h>
 
 /* Output macros */
-#define DEBUG DEBUG
+//#define DEBUG DEBUG
 #define HEXDUMP(type, prefix, ptr, size)                 \
   do {                                                   \
                                                          \
diff --git a/src/memory.c b/src/memory.c
index babc7fc..eb34fdf 100644
--- a/src/memory.c
+++ b/src/memory.c
@@ -137,7 +137,7 @@ pte_t *memwalk(uint64_t addr, struct mm_struct *mm){
                                 }
                         }
                 }
-        }
+        } return 0x00;
 }
 
 // TODO lock thee lists
diff --git a/test/test1.c b/test/test1.c
index b6e6a79..92b0e15 100644
--- a/test/test1.c
+++ b/test/test1.c
@@ -20,30 +20,31 @@ int* shm_addr;
 int* none_addr;
 void *chunk;
 int pippo = 1;
+FILE *stdf = NULL;
 
 void test2() {
   *shm_addr = 0xffffffff;
-  if (afl_snapshot_take(AFL_SNAPSHOT_NOSTACK) == 1)
-    fprintf(stderr, "first time!\n");
+  if (afl_snapshot_take( AFL_SNAPSHOT_REGS ) == 1)
+    fprintf(stdf, "first time!\n");
 
 loop:
   memset(chunk, 0x43, 0x1001);
-  fprintf(stderr, "*chunk = %#llx\n", *(uint64_t*)chunk);
+  fprintf(stdf, "*chunk = %#llx\n", *(uint64_t*)chunk);
   *none_addr += 1;
-  fprintf(stderr, ">> %d     %p = %#llx    %p = %#llx\n", pippo, none_addr, *none_addr, shm_addr, *shm_addr);
+  fprintf(stdf, ">> %d     %p = %#llx    %p = %#llx\n", pippo, none_addr, *none_addr, shm_addr, *shm_addr);
   *shm_addr = 0xdeadbeef;
   ++pippo;
   memset(chunk, 0x44, 0x1001);
-  fprintf(stderr, "*chunk = %#llx\n", *(uint64_t*)chunk);
+  fprintf(stdf, "*chunk = %#llx\n", *(uint64_t*)chunk);
   afl_snapshot_restore();
-  fprintf(stderr, "*chunk = %#llx\n", *(uint64_t*)chunk);
-  if( *none_addr > 0x100 ) exit(0x00);
+  fprintf(stdf, "*chunk = %#llx\n", *(uint64_t*)chunk);
+  if( *none_addr > 0x10000 ) exit(0x00);
   goto loop;
 
 }
 
 int main() {
-
+  stdf = fopen("/dev/shm/log", "a");
   afl_snapshot_init();
   puts("1");
   shm_addr = mmap(0, 0x10000, PROT_READ | PROT_WRITE | PROT_EXEC,
@@ -62,10 +63,10 @@ int main() {
   puts("6");
   chunk = malloc(0x1001);
   memset(chunk, 0x41, 0x1001);
-  fprintf(stderr, "*chunk = %#llx\n", *(uint64_t*)chunk);
+  fprintf(stdf, "*chunk = %#llx\n", *(uint64_t*)chunk);
   afl_snapshot_include_vmrange(chunk, chunk+0x1001);
   memset(chunk, 0x42, 0x1001);
-  fprintf(stderr, "*chunk = %#llx\n", *(uint64_t*)chunk);
+  fprintf(stdf, "*chunk = %#llx\n", *(uint64_t*)chunk);
   puts("7");
   test2();
 

From 8c1f9177adfecb95452dd4814bfc6fae6a6a5e0c Mon Sep 17 00:00:00 2001
From: Administrator <root@example.com>
Date: Tue, 17 Aug 2021 02:24:55 +0300
Subject: [PATCH 7/9] add module description

---
 .gitignore      |  1 +
 src/files.c     |  2 +-
 src/module.c    |  6 ++++--
 src/snapshot.c  | 28 +++-------------------------
 src/task_data.h |  4 ++--
 test/test1.c    | 22 ++++++++++++++--------
 6 files changed, 25 insertions(+), 38 deletions(-)

diff --git a/.gitignore b/.gitignore
index 52e576f..818301c 100644
--- a/.gitignore
+++ b/.gitignore
@@ -52,3 +52,4 @@ modules.order
 Module.symvers
 Mkfile.old
 dkms.conf
+.idea/
diff --git a/src/files.c b/src/files.c
index bf80cb1..6af5be0 100644
--- a/src/files.c
+++ b/src/files.c
@@ -9,7 +9,7 @@ void take_files_snapshot(struct task_data *data) {
   struct fdtable *     fdt = rcu_dereference_raw(files->fdt);
   int                  size, i;
 
-  size = (fdt->max_fds - 1) / BITS_PER_LONG + 1;
+  size = (fdt->max_fds - 1) / BITS_PER_LONG + 1; // NOLINT(cppcoreguidelines-narrowing-conversions)
 
   if (data->snapshot_open_fds == NULL)
     data->snapshot_open_fds =
diff --git a/src/module.c b/src/module.c
index a5285ff..7c67cae 100644
--- a/src/module.c
+++ b/src/module.c
@@ -27,8 +27,10 @@
 #define CLASS_NAME "afl_snapshot"
 
 MODULE_LICENSE("GPL");
-MODULE_AUTHOR("andreafioraldi && Th3C4t");
-MODULE_DESCRIPTION("Fast process snapshots for fuzzing");
+MODULE_AUTHOR("kallsyms && andreafioraldi");
+MODULE_DESCRIPTION("Fast process snapshots for fuzzing\n"
+		"thx to xcellerator for ftrace helper funcs\n"
+		"adoptation by Th3C4t\n");
 MODULE_VERSION("1.1.0");
 #ifndef __used
 # define __used		__attribute__((used))
diff --git a/src/snapshot.c b/src/snapshot.c
index 5ebc6b1..05b126f 100644
--- a/src/snapshot.c
+++ b/src/snapshot.c
@@ -6,22 +6,19 @@
 /*int exit_hook(struct kprobe *p, struct pt_regs *regs) {*/
 asmlinkage long exit_hook(long rdi) {
   clean_snapshot();
-//  printk("do_exit_orig == %#llx\n", (long*)do_exit_orig);
-  do_exit_orig(rdi); /* no return */
-  return ((long(*)())(0xdeadbeef))(); /* to keep gcc happy */
+  //  printk("do_exit_orig == %#llx\n", (long*)do_exit_orig);
+  do_exit_orig(rdi);                   /* no return */
+  return ((long (*)())(0xdeadbeef))(); /* to keep gcc happy */
 }
 
 void initialize_snapshot(struct task_data *data, int config) {
-
   struct pt_regs *regs = task_pt_regs(current);
 
   data->config = config;
 
   if (!had_snapshot(data)) {
-
     set_had_snapshot(data);
     hash_init(data->ss.ss_page);
-
   }
 
   set_snapshot(data);
@@ -33,11 +30,9 @@ void initialize_snapshot(struct task_data *data, int config) {
 
   // copy current brk
   data->ss.oldbrk = current->mm->brk;
-
 }
 
 int take_snapshot(int config) {
-
   struct task_data *data = ensure_task_data(current);
 
   if (!have_snapshot(data)) {  // first execution
@@ -47,64 +42,49 @@ int take_snapshot(int config) {
     take_files_snapshot(data);
 
     return 1;
-
   }
 
   return 0;
-
 }
 
 void recover_state(struct task_data *data) {
-
   if (data->config & AFL_SNAPSHOT_REGS) {
-
     struct pt_regs *regs = task_pt_regs(current);
 
     // restore regs context
     *regs = data->ss.regs;
-
   }
 
   // restore brk
   if (current->mm->brk > data->ss.oldbrk) current->mm->brk = data->ss.oldbrk;
-
 }
 
 void restore_snapshot(struct task_data *data) {
-
   recover_threads_snapshot(data);
   recover_memory_snapshot(data);
   recover_files_snapshot(data);
   recover_state(data);
-
 }
 
 void recover_snapshot(void) {
-
   struct task_data *data = get_task_data(current);
   restore_snapshot(data);
-
 }
 
 long exit_snapshot(void) {
-
   struct task_data *data = get_task_data(current);
   if (data && (data->config & AFL_SNAPSHOT_EXIT) && have_snapshot(data)) {
-
     restore_snapshot(data);
     return 0;
-
   }
 
   if (data && had_snapshot(data)) clean_snapshot();
 
   return 1;
-
 }
 
 /* :TODO: this should be static inline at least? */
 void clean_snapshot(void) {
-
   struct task_data *data = get_task_data(current);
   if (!data) { return; }
 
@@ -113,6 +93,4 @@ void clean_snapshot(void) {
   clear_snapshot(data);
 
   remove_task_data(data);
-
 }
-
diff --git a/src/task_data.h b/src/task_data.h
index b8efbbb..16a8e6a 100644
--- a/src/task_data.h
+++ b/src/task_data.h
@@ -54,13 +54,13 @@ static inline void set_snapshot(struct task_data *data) {
 
 static inline bool have_snapshot(struct task_data *data) {
 
-  return !!(data->ss.status & SNAPSHOT_MADE);
+  return (data->ss.status & SNAPSHOT_MADE) != 0;
 
 }
 
 static inline bool had_snapshot(struct task_data *data) {
 
-  return !!(data->ss.status & SNAPSHOT_HAD);
+  return (data->ss.status & SNAPSHOT_HAD) != 0;
 
 }
 
diff --git a/test/test1.c b/test/test1.c
index 92b0e15..e1d4dc8 100644
--- a/test/test1.c
+++ b/test/test1.c
@@ -12,7 +12,7 @@
 #include <sys/types.h>
 #include <unistd.h>
 
-#include "libaflsnapshot.h"
+#include "../libaflsnapshot.h"
 
 // gcc -I ../include -g test2.c ../lib/libaflsnapshot.o -o test2
 
@@ -29,20 +29,25 @@ void test2() {
 
 loop:
   memset(chunk, 0x43, 0x1001);
-  fprintf(stdf, "*chunk = %#llx\n", *(uint64_t*)chunk);
+  fprintf(stdf, "*chunk = %#lx\n", *(uint64_t*)chunk);
   *none_addr += 1;
-  fprintf(stdf, ">> %d     %p = %#llx    %p = %#llx\n", pippo, none_addr, *none_addr, shm_addr, *shm_addr);
+  fprintf(stdf, ">> %d     %p = %#x    %p = %#x\n", pippo, none_addr, *none_addr, shm_addr, *shm_addr);
   *shm_addr = 0xdeadbeef;
   ++pippo;
   memset(chunk, 0x44, 0x1001);
-  fprintf(stdf, "*chunk = %#llx\n", *(uint64_t*)chunk);
+  fprintf(stdf, "*chunk = %#lx\n", *(uint64_t*)chunk);
   afl_snapshot_restore();
-  fprintf(stdf, "*chunk = %#llx\n", *(uint64_t*)chunk);
+  fprintf(stdf, "*chunk = %#lx\n", *(uint64_t*)chunk);
   if( *none_addr > 0x10000 ) exit(0x00);
   goto loop;
 
 }
 
+void test3(){
+  s = socket
+
+}
+
 int main() {
   stdf = fopen("/dev/shm/log", "a");
   afl_snapshot_init();
@@ -63,12 +68,13 @@ int main() {
   puts("6");
   chunk = malloc(0x1001);
   memset(chunk, 0x41, 0x1001);
-  fprintf(stdf, "*chunk = %#llx\n", *(uint64_t*)chunk);
+  fprintf(stdf, "*chunk = %#lx\n", *(uint64_t*)chunk);
   afl_snapshot_include_vmrange(chunk, chunk+0x1001);
   memset(chunk, 0x42, 0x1001);
-  fprintf(stdf, "*chunk = %#llx\n", *(uint64_t*)chunk);
+  fprintf(stdf, "*chunk = %#lx\n", *(uint64_t*)chunk);
   puts("7");
-  test2();
+  //test2();
+  test3();
 
   return 0;
 

From c5a2318d72b801b8e35b4a3695dfeaaf49eaea13 Mon Sep 17 00:00:00 2001
From: Administrator <root@example.com>
Date: Tue, 17 Aug 2021 02:27:27 +0300
Subject: [PATCH 8/9] add test

---
 test/test1.c |   5 --
 test/test3.c | 240 +++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 240 insertions(+), 5 deletions(-)
 create mode 100644 test/test3.c

diff --git a/test/test1.c b/test/test1.c
index e1d4dc8..bfdfd98 100644
--- a/test/test1.c
+++ b/test/test1.c
@@ -43,11 +43,6 @@ void test2() {
 
 }
 
-void test3(){
-  s = socket
-
-}
-
 int main() {
   stdf = fopen("/dev/shm/log", "a");
   afl_snapshot_init();
diff --git a/test/test3.c b/test/test3.c
new file mode 100644
index 0000000..e67cd8a
--- /dev/null
+++ b/test/test3.c
@@ -0,0 +1,240 @@
+/*
+Manul - test file
+-------------------------------------
+Maksim Shudrak <mshudrak@salesforce.com> <mxmssh@gmail.com>
+
+Copyright 2019 Salesforce.com, inc. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at:
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <malloc.h>
+#include <errno.h>
+#include <unistd.h>
+#include <stdint.h>
+#include <signal.h>
+#include <pthread.h>
+#include <sys/types.h>
+#include <sys/socket.h>
+#include <arpa/inet.h>
+#include "../include/afl_snapshot.h"
+#include "../include/libaflsnapshot.h"
+
+static int dev_fd;
+void       LOG(const char *msg);
+char *     log_name = "/dev/shm/stage4.log";
+
+// static unsigned char *_buf = NULL;
+static int  branch_1(int in, char *_buf);
+static int  branch_2(char *buf);
+void *      ThreadMain(void *argv);
+static void p1w(void *addr) {
+  LOG("Yes, General!\n");
+  if (addr) ((volatile void (*)())addr)();
+}
+
+void LOG2WIN(const char *msg) {
+  FILE *f = fopen("/dev/shm/WIN", "a");
+  fprintf(f, "%s", msg);
+  fclose(f);
+}
+
+void LOG(const char *msg) {
+  FILE *f = fopen(log_name, "a");
+  fprintf(f, "%s", msg);
+  fclose(f);
+}
+
+void *open_file(char *name) {
+  char *buf = NULL;
+  int   size = 0;
+  FILE *fp = fopen(name, "rb");
+  if (!fp) {
+    printf("Couldn't open file specified %s", name);
+    return 0x00;
+  }
+  printf("Opening %s\n", name);
+  // obtain file size:
+  fseek(fp, 0, SEEK_END);
+  size = ftell(fp);
+  rewind(fp);
+
+  // allocate memory to contain the whole file:
+  buf = (char *)malloc(sizeof(char) * size);
+  if (buf == NULL) {
+    LOG("Unable to read file");
+    exit(-1);
+  }
+
+  // copy the file into the buffer:
+  fread(buf, 1, size, fp);
+  fclose(fp);
+  return buf;
+}
+
+int main() {
+  dev_fd = afl_snapshot_init();
+  printf("AFL-kit device opened %d\n", dev_fd);
+
+  int stop = 0x00;
+
+  signal(SIGPIPE, SIG_IGN);
+  LOG("Initializing...\n");
+
+  int server = socket(PF_INET, SOCK_STREAM, 0);
+  if (server == -1) {
+    printf("Could not create server socket: %d\n", errno);
+    return 2;
+  }
+  LOG("[+]Socket \n");
+  int retVal = 0;
+
+  int enable = 1;
+  if (setsockopt(server, SOL_SOCKET, SO_REUSEADDR, &enable, sizeof(enable)) !=
+      0) {
+    printf("Could not set SO_REUSEADDR: %d\n", errno);
+    retVal = 2;
+    goto end;
+  }
+  LOG("[+]Socket opts\n");
+
+  struct sockaddr_in addr = {0};
+  addr.sin_family = AF_INET;
+  addr.sin_port = htons(8080);
+  addr.sin_addr.s_addr = INADDR_ANY;
+
+  if (bind(server, (struct sockaddr *)&addr, sizeof(addr)) != 0) {
+    printf("Could not bind: %d\n", errno);
+    retVal = 3;
+    goto end;
+  }
+  LOG("[+]Socket binded\n");
+  if (listen(server, 20) != 0) {
+    printf("Could not listen: %d\n", errno);
+    retVal = 4;
+    goto end;
+  }
+  LOG("[+]Going to point of __noreturn :3\n");
+  while (!stop) {
+    LOG("[I]main: loop prologue\n");
+    struct sockaddr_storage theirAddr = {0};
+    socklen_t               len = sizeof(theirAddr);
+    int client = accept(server, (struct sockaddr *)&theirAddr, &len);
+    if (client == -1) {
+      printf("Accept returned %d\n", errno);
+      continue;
+    }
+    stop = (int)(long)ThreadMain((void*)(long)client);
+    LOG("[I]main: loop epilogue\n");
+  }
+
+end:
+  close(server);
+  return retVal;
+}
+
+#define HELLO "Hello, Fuzz!\n"
+void *ThreadMain(void *argv) {
+  char     name[0xff] = {0x00}, handshake[0xff] = {0x00}, *_buf = NULL;
+  size_t   hssize = 0x00;
+  uint32_t resp = 0x00;
+  int      fd = (int)(long)argv;
+  int      ok = 0x00;
+  if (fd < 0) goto fall;
+
+  ok = afl_snapshot_take(AFL_SNAPSHOT_REGS | AFL_SNAPSHOT_FDS |
+                         AFL_SNAPSHOT_MMAP);
+
+  printf("snapshot taken %d\n", ok);
+
+  strcat(handshake, HELLO);
+  strcat(handshake, "\xff\x44\x33\x33\x00");
+  hssize = strlen(handshake);
+
+  ok = send(fd, handshake, hssize, 0x00);
+  if (ok != hssize) goto fall;
+
+  ok = recv(fd, (void *)&resp, sizeof(resp), 0x00);
+  LOG2WIN("\n(1)|");
+  LOG2WIN((void *)&resp);
+  LOG2WIN("|\n");
+  printf("recved data of %d is %s\n", ok, (char*)&resp);
+
+  printf("Now, recovering snapshot\n");
+  close(fd);
+  afl_snapshot_restore();
+
+  if (ok != sizeof(resp)) goto fall;
+  if (resp != 0x333344ff) goto fall;
+  ok = send(fd, "\x01", sizeof(char), 0x00);
+  if (ok != sizeof(char)) goto fall;
+
+  ok = recv(fd, name, 0xff - 1, 0x00);
+  if (ok <= 0x00) goto fall;
+  LOG2WIN("\n(2)|");
+  LOG2WIN(name);
+  LOG2WIN("|\n");
+  name[ok] = 0x00;
+
+  /* _buf = open_file(name); */
+  _buf = malloc(0x1000);
+  if (!_buf) return (void *)-1;
+
+  recv(fd, _buf, 0x1000 - 1, MSG_DONTWAIT);
+  if (_buf[0] == 0x01) p1w((void *)(long)branch_1(_buf[1], &(_buf[2])));
+
+  if (malloc_usable_size(_buf)) {
+    memset(_buf, 0x41, malloc_usable_size(_buf));
+    free(_buf);
+    _buf = NULL;
+  }
+
+  send(fd, "\x02", sizeof(char), 0x00);
+  return NULL;
+
+fall:
+  printf("EXIT TRIGGERED!\n");
+  close(fd);
+  if (_buf) free(_buf);
+  _buf = NULL;
+  exit(0x00);
+  return 0x00;
+}
+
+static int branch_1(int in, char *_buf) {
+  int ret;
+  LOG2WIN("Wiiiin");
+  printf("in=%d, buf=%s\n", in, _buf);
+  if (in % 0xae != 0x00) ret = 0x00;
+
+  if (in % 17 == 0x00)
+    if (_buf) ret = branch_2(_buf);
+  return ret;
+}
+
+static int branch_2(char *buf) {
+  LOG("hitted brach_2\n");
+  if (buf[0] == 'P') {
+    if (buf[1] == 'W') {
+      if (buf[2] == 'N') {
+        if (buf[3] == 'I') {
+          LOG("Found it!\n");
+          return 0xdeadbeef;
+        }
+      }
+    }
+  }
+  return 0x00;
+}

From e35d1ac4e8b28e48072d1bb0a00060aae0af8481 Mon Sep 17 00:00:00 2001
From: Administrator <root@example.com>
Date: Wed, 18 Aug 2021 01:33:09 +0300
Subject: [PATCH 9/9] Multi-thread support. WIP started

---
 include/libaflsnapshot.h |   2 +-
 lib/libaflsnapshot.c     |   6 ++-
 src/Makefile             |   3 +-
 src/module.c             |  31 ++++++++----
 src/snapshot.c           | 107 +++++++++++++++++++++++++++++++++++----
 src/snapshot.h           |   6 +--
 src/task_data.c          |  51 +++++++++++++++----
 src/task_data.h          |  76 ++++++++++++++++++++++++---
 test/make.sh             |   2 +
 test/test1.c             |   8 ++-
 test/test3.c             |   2 +-
 11 files changed, 245 insertions(+), 49 deletions(-)

diff --git a/include/libaflsnapshot.h b/include/libaflsnapshot.h
index 7fc069d..2765860 100644
--- a/include/libaflsnapshot.h
+++ b/include/libaflsnapshot.h
@@ -7,7 +7,7 @@ int  afl_snapshot_init();
 void afl_snapshot_exclude_vmrange(void *start, void *end);
 void afl_snapshot_include_vmrange(void *start, void *end);
 int  afl_snapshot_do(void);
-int  afl_snapshot_take(int config);
+int  afl_snapshot_take(int pid, unsigned long config);
 void afl_snapshot_restore(void);
 void afl_snapshot_clean(void);
 
diff --git a/lib/libaflsnapshot.c b/lib/libaflsnapshot.c
index 84e9be9..a83b1fe 100644
--- a/lib/libaflsnapshot.c
+++ b/lib/libaflsnapshot.c
@@ -29,9 +29,11 @@ void afl_snapshot_include_vmrange(void *start, void *end) {
 
 }
 
-int afl_snapshot_take(int config) {
+int afl_snapshot_take(int pid, unsigned long config) {
 
-  return ioctl(dev_fd, AFL_SNAPSHOT_IOCTL_TAKE, config);
+  /* high half of config can't be used! */
+  if(config >> 0x20) return -1;
+  return ioctl(dev_fd, AFL_SNAPSHOT_IOCTL_TAKE, (((long)pid) << 0x20) | config);
 
 }
 
diff --git a/src/Makefile b/src/Makefile
index f90a0f5..c1dce8f 100644
--- a/src/Makefile
+++ b/src/Makefile
@@ -24,8 +24,7 @@ ccflags-y := \
   -Wno-declaration-after-statement \
   -fdata-sections \
   $(CCFLAGS) \
-  -O2
-#  -g3 -ggdb3 
+  -O0 -g3 -ggdb3
 
 ifdef DEBUG
   ccflags-y += -DDEBUG
diff --git a/src/module.c b/src/module.c
index 7c67cae..b823418 100644
--- a/src/module.c
+++ b/src/module.c
@@ -55,6 +55,8 @@ MODULE_VERSION("1.1.0");
 typedef typeof(void *(*)(void *rdi, ...)) func_t;
 func_t	_stop_machine = ((typeof(_stop_machine))&stop_machine);
 func_t	_kallsyms_lookup_name = NULL;
+func_t k_kern_addr_valid = NULL;
+void *k_tasklist_lock = NULL;
 
 void (*k_flush_tlb_mm_range)(struct mm_struct *mm, unsigned long start,
                              unsigned long end, unsigned int stride_shift,
@@ -79,6 +81,14 @@ static char *mod_devnode(struct device *dev, umode_t *mode) {
 long mod_dev_ioctl(struct file *filep, unsigned int cmd, unsigned long arg) {
 
   switch (cmd) {
+      /* should be on the first place when compare */
+    case AFL_SNAPSHOT_IOCTL_TAKE: {
+
+      DBG_PRINT("Calling afl_snapshot_take");
+      /* should be in void state, while doing snapshot ? */
+      return (long)_stop_machine(take_snapshot, arg, NULL);
+
+    }
 
     case AFL_SNAPSHOT_EXCLUDE_VMRANGE: {
 
@@ -86,8 +96,9 @@ long mod_dev_ioctl(struct file *filep, unsigned int cmd, unsigned long arg) {
 
       struct afl_snapshot_vmrange_args args;
       if (copy_from_user(&args, (void *)arg,
-                         sizeof(struct afl_snapshot_vmrange_args)))
+                         sizeof(struct afl_snapshot_vmrange_args))) {
         return -EINVAL;
+      }
 
       exclude_vmrange(args.start, args.end);
       return 0;
@@ -100,21 +111,15 @@ long mod_dev_ioctl(struct file *filep, unsigned int cmd, unsigned long arg) {
 
       struct afl_snapshot_vmrange_args args;
       if (copy_from_user(&args, (void *)arg,
-                         sizeof(struct afl_snapshot_vmrange_args)))
+                         sizeof(struct afl_snapshot_vmrange_args))) {
         return -EINVAL;
+      }
 
       include_vmrange(args.start, args.end);
       return 0;
 
     }
 
-    case AFL_SNAPSHOT_IOCTL_TAKE: {
-
-      DBG_PRINT("Calling afl_snapshot_take");
-
-      return take_snapshot(arg);
-
-    }
 
     case AFL_SNAPSHOT_IOCTL_DO: {
 
@@ -142,6 +147,7 @@ long mod_dev_ioctl(struct file *filep, unsigned int cmd, unsigned long arg) {
       return 0;
 
     }
+    default: break;
 
   }
 
@@ -251,9 +257,14 @@ int snapshot_initialize_k_funcs() {
 
   k_flush_tlb_mm_range = (void *)_kallsyms_lookup_name("flush_tlb_mm_range");
   k_zap_page_range = (void *)_kallsyms_lookup_name("zap_page_range");
+  k_tasklist_lock = (void *)_kallsyms_lookup_name("tasklist_lock");
+  k_kern_addr_valid = (void*)_kallsyms_lookup_name("kern_addr_valid");
 
   if (!chek_kaddr(k_flush_tlb_mm_range) ||
-	!chek_kaddr(k_zap_page_range)) { return -ENOENT; }
+	!chek_kaddr(k_zap_page_range)   ||
+        !chek_kaddr(k_kern_addr_valid)  ||
+        !k_tasklist_lock)
+  { return -ENOENT; }
 
   SAYF("All loaded :3");
 
diff --git a/src/snapshot.c b/src/snapshot.c
index 05b126f..ae3b8fd 100644
--- a/src/snapshot.c
+++ b/src/snapshot.c
@@ -4,15 +4,15 @@
 #include "snapshot.h"
 
 /*int exit_hook(struct kprobe *p, struct pt_regs *regs) {*/
-asmlinkage long exit_hook(long rdi) {
+asmlinkage __naked long exit_hook(long rdi) {
   clean_snapshot();
   //  printk("do_exit_orig == %#llx\n", (long*)do_exit_orig);
   do_exit_orig(rdi);                   /* no return */
   return ((long (*)())(0xdeadbeef))(); /* to keep gcc happy */
 }
 
-void initialize_snapshot(struct task_data *data, int config) {
-  struct pt_regs *regs = task_pt_regs(current);
+void initialize_snapshot(struct task_data *data, unsigned long config) {
+
 
   data->config = config;
 
@@ -25,25 +25,114 @@ void initialize_snapshot(struct task_data *data, int config) {
   // INIT_LIST_HEAD(&(data->ss.ss_mmap));
   data->ss.ss_mmap = NULL;
 
-  // copy current regs context
-  data->ss.regs = *regs;
-
   // copy current brk
   data->ss.oldbrk = current->mm->brk;
 }
 
-int take_snapshot(int config) {
-  struct task_data *data = ensure_task_data(current);
+#if 0
+inline static dump_task_stuff(){
+  initialize_snapshot(data, config);
+  take_memory_snapshot(data);
+  take_files_snapshot(data);
+}
+#endif
+
+/*
+ * thread_walk_cb -- search all threads of needle PID
+ * $rax == &task_struct of thread, of found, or NULL;
+ */
+static struct task_struct *thread_walk_cb(struct task_struct *tsk, pid_t needle){
+  pid_t pid_nr = tsk->pid;
+  pid_t parrent_nr = tsk->parent ? tsk->parent->pid : 0x00;
+  pid_t leader_nr = tsk->group_leader->pid;
+  char name[TASK_COMM_LEN];
+
+  get_task_comm(name, tsk); name[TASK_COMM_LEN] = 0x00;
+  printk("I'm kernel_thread, my pid %d, my parent's pid %d, my name %s "
+      "our group_leader pid is %d\n",
+         pid_nr, parrent_nr, name, leader_nr);
+  if( needle == parrent_nr ){
+    printk("Found needed thread, returning &task_struct\n");
+    return tsk;
+  }
+  return 0x00;
+}
+inline static int wpt_cb(struct task_struct *t, void *args){
+  printk("%s, my name %s, my pid %d\n", __func__, t->comm, t->pid);
+  return 0x00;
+}
+
+#if 0
+[  105.169403] [AFL++] SAY: Entry to snapshot: requested pid 565, my name test4 my group_leader pid is 565
+[  105.171610] Group thread founded!
+               name test4, tid 565, pid 565
+[  105.172717] I'm kernel_thread, my pid 565, my parent's pid 549, my name test4 our group_leader pid is 565
+[  105.174309] Group thread founded!
+               name test4, tid 567, pid 565
+[  105.175460] I'm kernel_thread, my pid 565, my parent's pid 549, my name test4 our group_leader pid is 565
+#endif
+
+/*
+ * take_snapshot -- handle ioctl request from userspace
+ * @config -- this is [int][int]. Lowest part as it was,
+ *    highest part == target pid_t, or 0x00 -- than current will be used
+ */
+long take_snapshot(unsigned long config) {
+  /* reflective overloaded */
+  struct task_struct *main_tsk = NULL;
+
+  pid_t main_nr = (int)(config >> 0x20);
+
+  struct pid *main_pid = (main_nr) ?
+          find_get_pid(main_nr)    :
+          get_task_pid((main_tsk = current->group_leader), PIDTYPE_PID);
+  if( !main_pid ) return -1;
+
+  struct task_data *data =
+      ensure_task_data( (main_tsk) ?: /* already ptr to group leader */
+          (main_tsk = (get_pid_task(main_pid, 0x00)->group_leader)) );
+  /* above should be safe, coz &pid always have &task, and &task always have -> group_leader */
+
+  SAYF("Entry to snapshot: requested pid %d, my name %s "
+         "my group_leader pid is %d\n",
+       main_nr, main_tsk->comm, main_pid->numbers->nr);
 
   if (!have_snapshot(data)) {  // first execution
+    struct task_struct *p, *t;
+    /* do not push read_lock, since we are in the void */
+    //read_lock(k_tasklist_lock);
+    for_each_process_thread(p, t) {
+      if (p == main_tsk) {
+        if (p == t) {
+          printk(
+              "Group leader founded!\n"
+              "name %s, tid %d, pid %d\n",
+              t->comm, t->pid, t->group_leader->pid);
+
+        } else {
+          printk(
+              "Group thread founded!\n"
+              "name %s, tid %d, pid %d\n",
+              t->comm, t->pid, t->group_leader->pid);
+          /* initialize snapshot data, etc */
+          // thread_walk_cb(p, main_nr);
+        }
+      }
+    }
+
+
+    #warning "main_pid unneeded?"
+    put_pid(main_pid);
+    return 0x01;
 
     initialize_snapshot(data, config);
     take_memory_snapshot(data);
     take_files_snapshot(data);
+    //read_unlock(k_tasklist_lock);
 
+    /* :XXX: 0x00/-E conv. ? */
     return 1;
   }
-
   return 0;
 }
 
diff --git a/src/snapshot.h b/src/snapshot.h
index 52ef4bd..ca4fac5 100644
--- a/src/snapshot.h
+++ b/src/snapshot.h
@@ -175,8 +175,6 @@ struct snapshot {
 
   struct snapshot_vma *ss_mmap;
 
-  struct pt_regs regs;
-
   DECLARE_HASHTABLE(ss_page, SNAPSHOT_HASHTABLE_SZ);
 
 };
@@ -210,13 +208,13 @@ int do_anonymous_hook(long rdi, long rsi, long rdx, long rcx);//(struct kprobe *
 long exit_hook(long rdi);//(struct kprobe *p, struct pt_regs *regs);
 extern typeof(&do_exit) do_exit_orig;
 
-int  take_snapshot(int config);
+long  take_snapshot(unsigned long config);
 void recover_snapshot(void);
 inline void clean_snapshot(void);
 long  exit_snapshot(void);
 
 void exclude_vmrange(unsigned long start, unsigned long end);
 void include_vmrange(unsigned long start, unsigned long end);
-
+extern void *k_tasklist_lock;
 #endif
 
diff --git a/src/task_data.c b/src/task_data.c
index 01c21c0..fcaa755 100644
--- a/src/task_data.c
+++ b/src/task_data.c
@@ -11,18 +11,16 @@ static void task_data_free_callback(struct rcu_head *rcu) {
 
 }
 
-struct task_data *get_task_data(const struct task_struct *tsk) {
+struct task_data *get_task_data(const struct task_struct *leader) {
 
   struct task_data *data = NULL;
 
   rcu_read_lock();
   list_for_each_entry_rcu(data, &task_datas, list) {
 
-    if (data->tsk == tsk) {
-
+    if (data->leader == leader) {
       rcu_read_unlock();
       return data;
-
     }
 
   }
@@ -33,23 +31,58 @@ struct task_data *get_task_data(const struct task_struct *tsk) {
 
 }
 
-struct task_data *ensure_task_data(const struct task_struct *tsk) {
+/*
+ * add_thread_data -- add certain thread to global process structure
+ * @data -- container of process snapshot, where to add thread
+ * @thread -- certain thread's &task_struct
+ * $rax -- ptr to certain offset in vector memory, or NULL in case of Err.
+ */
+struct thread_data *add_thread_data(struct task_data *data, struct task_struct *thread){
+  if( !ISKADDR(data) || !ISKADDR(thread) ) return NULL;
+  if( data->threads_nr >= 0x10 ) return NULL;
+
+  /* extract free piece of memory */
+  data->threads_nr++;
+  struct thread_data *thr = &data->threads[data->threads_nr];
+
+  /* filling the structure */
+  thr->tsk = thread;
+  thr->pid = get_task_pid(thread, );
 
-  struct task_data *data = get_task_data(tsk);
-  if (data) return data;
+  struct pt_regs *regs = task_pt_regs(thread);
+  memcpy(&thr->regs, regs, sizeof(thr->regs));
+
+}
+
+/*
+ * add_task_data -- allocate zeroed memory to place at least one threaded process
+ *    put initial information (ptr to group leader task_struct) and increase thread counter
+ * @leader -- ptr to group leader leader
+ * $rax -- ptr to mem
+ */
+struct task_data *add_task_data(struct task_struct *leader){
+  /* ensure */
+  if( !ISKADDR(leader) || leader->group_leader != leader ) return NULL;
 
   // XXX: this is academic code (tm) so if we run out of memory, too bad!
-  data = kmalloc(sizeof(struct task_data), GFP_KERNEL | __GFP_ZERO);
+  /* add place for 16 threads, including thread_leader. This should be enough in most cases */
+  /* :TODO: Add realloc mechanism, or, maybe, count needed place amount before? */
+  struct task_data *data = kzalloc(sizeof(*data) + sizeof(struct thread_data) * 0x10,
+                                   GFP_KERNEL);
   if (!data) return NULL;
 
-  data->tsk = tsk;
+  data->leader = leader;
+  data->threads_nr++;
 
   spin_lock(&task_datas_lock);
   list_add_rcu(&data->list, &task_datas);
   spin_unlock(&task_datas_lock);
 
   return data;
+}
 
+struct task_data *ensure_task_data(const struct task_struct *leader) {
+  return get_task_data(leader) ? : NULL;
 }
 
 void remove_task_data(struct task_data *data) {
diff --git a/src/task_data.h b/src/task_data.h
index 16a8e6a..e6768a6 100644
--- a/src/task_data.h
+++ b/src/task_data.h
@@ -4,6 +4,7 @@
 #include "snapshot.h"
 #include <linux/list.h>
 #include <linux/sched.h>
+#include "gears.h"
 
 struct vmrange_node {
 
@@ -13,20 +14,83 @@ struct vmrange_node {
   struct vmrange_node *next;
 
 };
+#if 0
+POSIX.1 also requires that threads share a range of other attributes (i.e., these attributes are process-wide rather than per-thread):
 
-struct task_data {
+-  process ID
+
+-  parent process ID
+
+-  process group ID and session ID
+
+-  controlling terminal
+
+-  user and group IDs
+
+-  open file descriptors
+
+-  record locks (see fcntl(2))
+
+-  signal dispositions
+
+-  file mode creation mask (umask(2))
+
+-  current directory (chdir(2)) and root directory (chroot(2))
+
+-  interval timers (setitimer(2)) and POSIX timers (timer_create(2))
+
+-  nice value (setpriority(2))
+
+-  resource limits (setrlimit(2))
+
+-  measurements of the consumption of CPU time (times(2)) and resources (getrusage(2))
+
+As well as the stack, POSIX.1 specifies that various other attributes are distinct for each thread, including:
+
+-  thread ID (the pthread_t data type)
+
+-  signal mask (pthread_sigmask(3))
+
+-  the errno variable
+
+-  alternate signal stack (sigaltstack(2))
+
+-  real-time scheduling policy and priority (sched(7))
+
+The following Linux-specific features are also per-thread:
+
+-  capabilities (see capabilities(7))
+
+-  CPU affinity (sched_setaffinity(2))
+
+#endif
+
+struct thread_data {
 
   // what task_struct is this for?
-  const struct task_struct *tsk;
+  const struct task_struct *tsk;          // ptr to task_struct of certain thread
+  const struct pid *pid;                  // ptr to pid struct of certain thread
+  const struct mm_struct *mm;             // ptr to mm struct of certain thread
+  const struct vm_area_struct *vma;       // ptr to vma struct of certain thread
+  void *thread_pid_field_backup;          // backup of pointer needed to restore thread visibility in the system
+  struct pt_regs regs;                    // backup of cpu-state per each thread
+
+};
+
+struct task_data {
+
+  int threads_nr;                          // count of all process threads at the moment of snapshoting
+  struct task_struct *leader;              // main task, present all process live
+  long  config;                            // configuration flags captured from user request
 
   struct snapshot ss;
   unsigned long * snapshot_open_fds;
-
   struct vmrange_node *allowlist, *blocklist;
-  int                  config;
 
-  struct list_head list;
-  struct rcu_head  rcu;
+  struct list_head list;                   // will have sense when we do snapshots of few process
+  struct rcu_head  rcu;                    // will have sense when we do snapshots of few process
+
+  struct thread_data threads[];            // vector of all threads pf process at the moment of snapshotting
 
 };
 
diff --git a/test/make.sh b/test/make.sh
index 66633a4..854ca18 100755
--- a/test/make.sh
+++ b/test/make.sh
@@ -1 +1,3 @@
 cc -I../include test1.c ../lib/libaflsnapshot.o -o test1
+cc -I../include test3.c ../lib/libaflsnapshot.o -o test3
+cc -I../include test4.c ../lib/libaflsnapshot.o -o test4 -lpthread
diff --git a/test/test1.c b/test/test1.c
index bfdfd98..29b6c70 100644
--- a/test/test1.c
+++ b/test/test1.c
@@ -12,7 +12,7 @@
 #include <sys/types.h>
 #include <unistd.h>
 
-#include "../libaflsnapshot.h"
+#include "../include/libaflsnapshot.h"
 
 // gcc -I ../include -g test2.c ../lib/libaflsnapshot.o -o test2
 
@@ -24,7 +24,7 @@ FILE *stdf = NULL;
 
 void test2() {
   *shm_addr = 0xffffffff;
-  if (afl_snapshot_take( AFL_SNAPSHOT_REGS ) == 1)
+  if (afl_snapshot_take(getpid(), AFL_SNAPSHOT_REGS ) == 1)
     fprintf(stdf, "first time!\n");
 
 loop:
@@ -68,9 +68,7 @@ int main() {
   memset(chunk, 0x42, 0x1001);
   fprintf(stdf, "*chunk = %#lx\n", *(uint64_t*)chunk);
   puts("7");
-  //test2();
-  test3();
-
+  test2();
   return 0;
 
 }
diff --git a/test/test3.c b/test/test3.c
index e67cd8a..49c2736 100644
--- a/test/test3.c
+++ b/test/test3.c
@@ -154,7 +154,7 @@ void *ThreadMain(void *argv) {
   int      ok = 0x00;
   if (fd < 0) goto fall;
 
-  ok = afl_snapshot_take(AFL_SNAPSHOT_REGS | AFL_SNAPSHOT_FDS |
+  ok = afl_snapshot_take(getpid(), AFL_SNAPSHOT_REGS | AFL_SNAPSHOT_FDS |
                          AFL_SNAPSHOT_MMAP);
 
   printf("snapshot taken %d\n", ok);