Add TB pointers to post gen block hooks

.github/workflows/build_and_test.yaml

@@ -19,5 +19,7 @@ jobs:
     container: registry.gitlab.com/qemu-project/qemu/qemu/ubuntu2204:latest
     steps:
       - uses: actions/checkout@v4
+      - name: Check LibAFL format
+        run: apt update && apt install -y clang-format && ./scripts/libafl-format.sh check
       - name: Build QEMU
         run: mkdir -p build && cd build && ../configure --enable-werror --disable-docs --enable-fdt=system && make -j $(expr $(nproc) + 1)

accel/tcg/cpu-exec.c

@@ -42,6 +42,13 @@
 #include "internal-common.h"
 #include "internal-target.h"
 
+//// --- Begin LibAFL code ---
+#include "libafl/exit.h"
+#include "libafl/tcg.h"
+
+#include "libafl/hooks/tcg/block.h"
+//// --- End LibAFL code ---
+
 /* -icount align implementation. */
 
 typedef struct SyncClocks {
@@ -286,6 +293,14 @@ static inline TranslationBlock *tb_lookup(CPUState *cpu, vaddr pc,
     return tb;
 }
 
+//// --- Begin LibAFL code ---
+TranslationBlock *libafl_tb_lookup(CPUState *cpu, vaddr pc,
+                                          uint64_t cs_base, uint32_t flags,
+                                          uint32_t cflags) {
+    return tb_lookup(cpu, pc, cs_base, flags, cflags);
+}
+//// --- End LibAFL code ---
+
 static void log_cpu_exec(vaddr pc, CPUState *cpu,
                          const TranslationBlock *tb)
 {
@@ -674,6 +689,13 @@ static inline void tb_add_jump(TranslationBlock *tb, int n,
     return;
 }
 
+//// --- Begin LibAFL code ---
+void libafl_tb_add_jump(TranslationBlock *tb, int n,
+                               TranslationBlock *tb_next) {
+    tb_add_jump(tb, n, tb_next);
+}
+//// --- End LibAFL code ---
+
 static inline bool cpu_handle_halt(CPUState *cpu)
 {
 #ifndef CONFIG_USER_ONLY
@@ -708,12 +730,6 @@ static inline void cpu_handle_debug_exception(CPUState *cpu)
     }
 }
 
-//// --- Begin LibAFL code ---
-
-#include "libafl/exit.h"
-
-//// --- End LibAFL code ---
-
 static inline bool cpu_handle_exception(CPUState *cpu, int *ret)
 {
     //// --- Begin LibAFL code ---
@@ -1017,6 +1033,23 @@ cpu_exec_loop(CPUState *cpu, SyncClocks *sc)
 
                 mmap_lock();
                 tb = tb_gen_code(cpu, pc, cs_base, flags, cflags);
+
+                //// --- Begin LibAFL code ---
+#ifndef CONFIG_USER_ONLY
+                /*
+                 * We don't take care of direct jumps when address mapping
+                 * changes in system emulation.  So it's not safe to make a
+                 * direct jump to a TB spanning two pages because the mapping
+                 * for the second page can change.
+                 */
+                if (tb_page_addr1(tb) != -1) {
+                    last_tb = NULL;
+                }
+#endif
+                libafl_qemu_hook_block_post_run(tb, last_tb, pc, tb_exit);
+
+                //// --- End LibAFL code ---
+
                 mmap_unlock();
 
                 /*
@@ -1028,18 +1061,23 @@ cpu_exec_loop(CPUState *cpu, SyncClocks *sc)
                 jc->array[h].pc = pc;
                 qatomic_set(&jc->array[h].tb, tb);
             }
-
+//// --- Begin LibAFL code ---
+            else {
+//// --- End LibAFL code ---
 #ifndef CONFIG_USER_ONLY
-            /*
-             * We don't take care of direct jumps when address mapping
-             * changes in system emulation.  So it's not safe to make a
-             * direct jump to a TB spanning two pages because the mapping
-             * for the second page can change.
-             */
-            if (tb_page_addr1(tb) != -1) {
-                last_tb = NULL;
-            }
+                /*
+                 * We don't take care of direct jumps when address mapping
+                 * changes in system emulation.  So it's not safe to make a
+                 * direct jump to a TB spanning two pages because the mapping
+                 * for the second page can change.
+                 */
+                if (tb_page_addr1(tb) != -1) {
+                    last_tb = NULL;
+                }
 #endif
+//// --- Begin LibAFL code ---
+            }
+//// --- End LibAFL code ---
 
             //// --- Begin LibAFL code ---
 

accel/tcg/translate-all.c

@@ -715,12 +715,6 @@ TranslationBlock *tb_gen_code(CPUState *cpu,
     }
     tb->tc.size = gen_code_size;
 
-//// --- Begin LibAFL code ---
-
-    libafl_qemu_hook_block_post_run(tb, pc);
-
-//// --- End LibAFL code ---
-
     /*
      * For CF_PCREL, attribute all executions of the generated code
      * to its first mapping.

include/libafl/hooks/tcg/block.h

@@ -12,7 +12,10 @@
 
 typedef uint64_t (*libafl_block_pre_gen_cb)(uint64_t data, target_ulong pc);
 typedef void (*libafl_block_post_gen_cb)(uint64_t data, target_ulong pc,
-                                         target_ulong block_length);
+                                         target_ulong block_length,
+                                         TranslationBlock* tb,
+                                         TranslationBlock* last_tb,
+                                         int tb_exit);
 
 typedef void (*libafl_block_exec_cb)(uint64_t data, uint64_t id);
 
@@ -47,4 +50,6 @@ bool libafl_qemu_block_hook_set_jit(
 int libafl_qemu_remove_block_hook(size_t num, int invalidate);
 
 void libafl_qemu_hook_block_pre_run(target_ulong pc);
-void libafl_qemu_hook_block_post_run(TranslationBlock* tb, vaddr pc);
+void libafl_qemu_hook_block_post_run(TranslationBlock* tb,
+                                     TranslationBlock* last_tb, vaddr pc,
+                                     int tb_exit);

include/libafl/tcg.h

@@ -8,3 +8,11 @@
 
 void tcg_gen_callN(void* func, TCGHelperInfo* info, TCGTemp* ret,
                    TCGTemp** args);
+
+TranslationBlock* libafl_tb_lookup(CPUState* cpu, vaddr pc, uint64_t cs_base,
+                                   uint32_t flags, uint32_t cflags);
+
+TranslationBlock* libafl_tb_gen_code(CPUState* cpu, vaddr pc, uint64_t cs_base,
+                                     uint32_t flags, int cflags);
+
+void libafl_tb_add_jump(TranslationBlock* tb, int n, TranslationBlock* tb_next);

include/libafl/user.h

@@ -6,6 +6,10 @@
 
 #include "exec/cpu-defs.h"
 
+typedef void (*libafl_qemu_on_signal_hdlr)(int target_sig);
+
+extern libafl_qemu_on_signal_hdlr libafl_signal_hdlr;
+
 struct libafl_mapinfo {
     target_ulong start;
     target_ulong end;
@@ -39,6 +43,7 @@ IntervalTreeNode* libafl_maps_first(IntervalTreeRoot* map_info);
 IntervalTreeNode* libafl_maps_next(IntervalTreeNode* pageflags_maps_node,
                                    IntervalTreeRoot* proc_maps_node,
                                    struct libafl_mapinfo* ret);
+bool libafl_is_valid_addr(target_ulong addr);
 
 uint64_t libafl_load_addr(void);
 struct image_info* libafl_get_image_info(void);
@@ -52,6 +57,8 @@ int _libafl_qemu_user_init(int argc, char** argv, char** envp);
 bool libafl_get_return_on_crash(void);
 void libafl_set_return_on_crash(bool return_on_crash);
 
+void libafl_set_on_signal_handler(libafl_qemu_on_signal_hdlr hdlr);
+
 #ifdef AS_LIB
 void libafl_qemu_init(int argc, char** argv);
 #endif

libafl/hooks/tcg/block.c

@@ -52,12 +52,14 @@ bool libafl_qemu_block_hook_set_jit(size_t num, libafl_block_jit_cb jit_cb)
     return false;
 }
 
-void libafl_qemu_hook_block_post_run(TranslationBlock* tb, vaddr pc)
+void libafl_qemu_hook_block_post_run(TranslationBlock* tb,
+                                     TranslationBlock* last_tb, vaddr pc,
+                                     int tb_exit)
 {
     struct libafl_block_hook* hook = libafl_block_hooks;
     while (hook) {
         if (hook->post_gen_cb)
-            hook->post_gen_cb(hook->data, pc, tb->size);
+            hook->post_gen_cb(hook->data, pc, tb->size, tb, last_tb, tb_exit);
         hook = hook->next;
     }
 }

libafl/meson.build

@@ -5,6 +5,7 @@ specific_ss.add(files(
                     'jit.c',
                     'utils.c',
                     'gdb.c',
+                    'tcg.c',
 
                     # TCG-related hooks
                     'hooks/tcg/backdoor.c',

libafl/tcg.c

@@ -0,0 +1,14 @@
+#include "libafl/tcg.h"
+
+#include "exec/exec-all.h"
+#include "accel/tcg/internal-common.h"
+
+TranslationBlock* libafl_tb_gen_code(CPUState* cpu, vaddr pc, uint64_t cs_base,
+                                     uint32_t flags, int cflags)
+{
+    mmap_lock();
+    TranslationBlock* tb = tb_gen_code(cpu, pc, cs_base, flags, cflags);
+    mmap_unlock();
+
+    return tb;
+}

libafl/user.c

@@ -1,6 +1,7 @@
 #include "qemu/osdep.h"
 #include "qemu.h"
 #include "loader.h"
+#include "exec/exec-all.h"
 
 #include "libafl/user.h"
 
@@ -14,6 +15,8 @@ static struct libafl_qemu_sig_ctx libafl_qemu_sig_ctx = {0};
 // if false, target crahes will raise the appropriate signal.
 static bool libafl_return_on_crash = false;
 
+libafl_qemu_on_signal_hdlr libafl_signal_hdlr = NULL;
+
 void host_signal_handler(int host_sig, siginfo_t* info, void* puc);
 
 void libafl_qemu_native_signal_handler(int host_sig, siginfo_t* info, void* puc)
@@ -65,6 +68,11 @@ void libafl_set_return_on_crash(bool return_on_crash)
 
 bool libafl_get_return_on_crash(void) { return libafl_return_on_crash; }
 
+void libafl_set_on_signal_handler(libafl_qemu_on_signal_hdlr hdlr)
+{
+    libafl_signal_hdlr = hdlr;
+}
+
 #ifdef AS_LIB
 void libafl_qemu_init(int argc, char** argv)
 {

linux-user/signal.c

@@ -788,35 +788,40 @@ void die_with_signal(int host_sig)
     // Instead, LibAFL is gonna catch the signal if it has put a handler for it
     // and decide what to do
 
-    // struct sigaction act = {
-    //     .sa_handler = SIG_DFL,
-    // };
+    struct sigaction act = {
+        .sa_handler = SIG_DFL,
+    };
 
-    /*
-     * The proper exit code for dying from an uncaught signal is -<signal>.
-     * The kernel doesn't allow exit() or _exit() to pass a negative value.
-     * To get the proper exit code we need to actually die from an uncaught
-     * signal.  Here the default signal handler is installed, we send
-     * the signal and we wait for it to arrive.
-     */
-    // sigfillset(&act.sa_mask);
-    // sigaction(host_sig, &act, NULL);
+    if (libafl_signal_hdlr) {
+        /*
+         * The proper exit code for dying from an uncaught signal is -<signal>.
+         * The kernel doesn't allow exit() or _exit() to pass a negative value.
+         * To get the proper exit code we need to actually die from an uncaught
+         * signal.  Here the default signal handler is installed, we send
+         * the signal and we wait for it to arrive.
+         */
 
-    // make sure signal is not blocked
-    sigset_t host_sig_set;
-    sigemptyset(&host_sig_set);
-    sigaddset(&host_sig_set, host_sig);
+        sigfillset(&act.sa_mask);
+        sigaction(host_sig, &act, NULL);
+    } else {
+        // make sure signal is not blocked
+        sigset_t host_sig_set;
+        sigemptyset(&host_sig_set);
+        sigaddset(&host_sig_set, host_sig);
 
-    sigprocmask(SIG_UNBLOCK, &host_sig_set, NULL);
+        sigprocmask(SIG_UNBLOCK, &host_sig_set, NULL);
+    }
 //// --- End LibAFL code ---
 
     kill(getpid(), host_sig);
 
-    /* Make sure the signal isn't masked (reusing the mask inside of act). */
 //// --- Start LibAFL code ---
-    // Unused as of now
-    // sigdelset(&act.sa_mask, host_sig);
-    // sigsuspend(&act.sa_mask);
+    /* Make sure the signal isn't masked (reusing the mask inside of act). */
+
+    if (libafl_signal_hdlr) {
+        sigdelset(&act.sa_mask, host_sig);
+        sigsuspend(&act.sa_mask);
+    }
 //// --- End LibAFL code ---
 
     /* unreachable */
@@ -860,6 +865,12 @@ void dump_core_and_abort(CPUArchState *env, int target_sig)
 
     preexit_cleanup(env, 128 + target_sig);
 
+//// --- Begin LibAFL code ---
+    if (libafl_signal_hdlr) {
+        libafl_signal_hdlr(target_sig);
+    }
+//// --- End LibAFL code ---
+
     die_with_signal(host_sig);
 }
 
@@ -1336,13 +1347,13 @@ static void handle_pending_signal(CPUArchState *cpu_env, int sig,
     if (unlikely(qemu_loglevel_mask(LOG_STRACE))) {
         print_taken_signal(sig, &unswapped);
     }
-    
+
     //// --- Start LibAFL code ---
-    
+
     if (libafl_force_dfl && (sig == SIGABRT || sig == SIGSEGV || sig == SIGILL || sig == SIGBUS)) {
         handler = TARGET_SIG_DFL;
     }
-    
+
     //// --- End LibAFL code ---
 
     if (handler == TARGET_SIG_DFL) {

linux-user/syscall.c

@@ -13925,6 +13925,14 @@ IntervalTreeNode * libafl_maps_next(IntervalTreeNode *pageflags_maps_node, Inter
     }
 }
 
+//// --- Begin LibAFL code ---
+bool libafl_is_valid_addr(target_ulong addr) {
+    int flags = page_get_flags(addr);
+
+    return (flags & PAGE_VALID) && (flags & PAGE_READ);
+}
+//// --- End LibAFL code ---
+
 //// --- End LibAFL code ---
 
 abi_long do_syscall(CPUArchState *cpu_env, int num, abi_long arg1,

scripts/libafl-format.sh

@@ -11,5 +11,5 @@ elif [ "$1" != "" ]; then
 fi
 
 cd "$SCRIPT_DIR" || exit 1
-find "$ROOT_DIR/libafl" -name "*.c" -exec clang-format $CHECK -style="file:$ROOT_DIR/libafl/.clang-format" -i {} \;
-find "$ROOT_DIR/include/libafl" -name "*.h" -exec clang-format $CHECK -style=file:"$ROOT_DIR/libafl/.clang-format" -i {} \;
+find "$ROOT_DIR/libafl" -name "*.c" -print0 | xargs -0 clang-format $CHECK -style="file:$ROOT_DIR/libafl/.clang-format" -i || exit 1
+find "$ROOT_DIR/include/libafl" -name "*.h" -print0 | xargs -0 clang-format $CHECK -style=file:"$ROOT_DIR/libafl/.clang-format" -i || exit 1