add block hook with instruction count (#2378)

Co-authored-by: Mark Leon Giraud <mark.giraud@iosb.fraunhofer.de>

Author: PhilippTakacs

include/uc_priv.h

@@ -219,6 +219,7 @@ typedef enum uc_hook_idx {
     UC_HOOK_EDGE_GENERATED_IDX,
     UC_HOOK_TCG_OPCODE_IDX,
     UC_HOOK_TLB_FILL_IDX,
+    UC_HOOK_BLOCK_ICOUNT_IDX,
 
     UC_HOOK_MAX,
 } uc_hook_idx;
@@ -533,6 +534,8 @@ static inline void hooked_regions_check(uc_engine *uc, uint64_t start,
     hooked_regions_check_single(uc->hook[UC_HOOK_CODE_IDX].head, start, length);
     hooked_regions_check_single(uc->hook[UC_HOOK_BLOCK_IDX].head, start,
                                 length);
+    hooked_regions_check_single(uc->hook[UC_HOOK_BLOCK_ICOUNT_IDX].head, start,
+                                length);
 }
 
 /*

include/unicorn/unicorn.h

@@ -408,6 +408,8 @@ typedef enum uc_hook_type {
     // Register tlb fill request hook on the virtuall addresses.
     // The callback will be triggert if the tlb cache don't contain an address.
     UC_HOOK_TLB_FILL = 1 << 17,
+    // Hook basic blocks but get the number of instructions as size
+    UC_HOOK_BLOCK_ICOUNT = 1 << 18,
 } uc_hook_type;
 
 // Hook type for all events of unmapped memory access

qemu/accel/tcg/translator.c

@@ -80,6 +80,13 @@ void translator_loop(const TranslatorOps *ops, DisasContextBase *db,
         gen_uc_tracecode(tcg_ctx, 0xf8f8f8f8, UC_HOOK_BLOCK_IDX, uc, db->pc_first);
     }
 
+    if (HOOK_EXISTS_BOUNDED(uc, UC_HOOK_BLOCK_ICOUNT, tb->pc)) {
+        ops->pc_sync(db, cpu);
+        prev_op = tcg_last_op(tcg_ctx);
+        block_hook = true;
+        gen_uc_tracecode(tcg_ctx, 0xf8f8f8f9, UC_HOOK_BLOCK_ICOUNT_IDX, uc, db->pc_first);
+    }
+
     // tcg_dump_ops(tcg_ctx, false, "translator loop");
 
     /* Start translating.  */
@@ -166,6 +173,10 @@ void translator_loop(const TranslatorOps *ops, DisasContextBase *db,
             tcg_op = QTAILQ_FIRST(&tcg_ctx->ops);
         }
 
-        tcg_op->args[1] = db->tb->size;
+        if (tcg_op->args[1] == 0xf8f8f8f8) {
+            tcg_op->args[1] = db->tb->size;
+        } else {
+            tcg_op->args[1] = db->tb->icount;
+        }
     }
 }

qemu/tcg/tcg.c

@@ -686,6 +686,7 @@ void uc_add_inline_hook(uc_engine *uc, struct hook *hk, void** args, int args_le
     // Only UC_HOOK_BLOCK and UC_HOOK_CODE is generated into tcg code and can be inlined.
     switch (hk->type) {
     case UC_HOOK_BLOCK:
+    case UC_HOOK_BLOCK_ICOUNT:
     case UC_HOOK_CODE:
         // (*uc_cb_hookcode_t)(uc_engine *uc, uint64_t address, uint32_t size, void *user_data);
         sizemask = dh_sizemask(void, 0) | dh_sizemask(ptr, 1) | dh_sizemask(i64, 2) | dh_sizemask(i32, 3) | dh_sizemask(ptr, 4);

tests/unit/test_ctl.c

@@ -429,6 +429,48 @@ static void test_add_block_hook(void)
     OK(uc_close(uc));
 }
 
+static void test_add_block_icount_hook_block_cb(uc_engine *uc, uint64_t address, uint32_t size, void *user_data)
+{
+    uint64_t *icount = user_data;
+    *icount += size;
+}
+
+static void test_add_block_icount_hook(void)
+{
+    uc_engine *uc;
+    uint64_t icount = 0;
+    uc_hook syscall_hook;
+    uc_hook block_hook;
+    /* nop
+     * syscall
+     */
+    char code[] = "\x90\x0F\x05";
+
+    uc_common_setup(&uc, UC_ARCH_X86, UC_MODE_64, code, sizeof(code) - 1);
+    OK(uc_hook_add(uc, &syscall_hook, UC_HOOK_INSN, &test_add_block_hook_syscall_cb, NULL, 1, 0, UC_X86_INS_SYSCALL));
+    OK(uc_emu_start(uc, code_start, 0, 0, 0));
+    OK(uc_hook_add(uc, &block_hook, UC_HOOK_BLOCK_ICOUNT, &test_add_block_icount_hook_block_cb, &icount, code_start, code_start+0x1000));
+    OK(uc_emu_start(uc, code_start, 0, 0, 0));
+    TEST_CHECK(icount == 2);
+    TEST_CHECK(uc_hook_add(uc, &block_hook, UC_HOOK_BLOCK, &test_add_block_icount_hook_block_cb, &icount, 0, code_start+0x1000) == UC_ERR_HOOK_EXIST);
+    TEST_CHECK(uc_hook_add(uc, &block_hook, UC_HOOK_BLOCK,
+                           &test_add_block_icount_hook_block_cb, &icount,
+                           code_start + 0x1000,
+                           code_start + 0x2000) == UC_ERR_HOOK_EXIST);
+    TEST_CHECK(uc_hook_add(uc, &block_hook,
+                           UC_HOOK_BLOCK | UC_HOOK_BLOCK_ICOUNT,
+                           &test_add_block_icount_hook_block_cb, &icount,
+                           code_start, code_start) == UC_ERR_HOOK);
+    OK(uc_hook_del(uc, block_hook));
+    OK(uc_hook_add(uc, &block_hook, UC_HOOK_BLOCK_ICOUNT,
+                   &test_add_block_icount_hook_block_cb, &icount, code_start,
+                   code_start));
+    TEST_CHECK(uc_hook_add(uc, &block_hook, UC_HOOK_BLOCK,
+                           &test_add_block_icount_hook_block_cb, &icount,
+                           code_start, code_start) == UC_ERR_HOOK_EXIST);
+    OK(uc_close(uc));
+}
+
 TEST_LIST = {
     {"test_uc_ctl_mode", test_uc_ctl_mode},
     {"test_uc_ctl_page_size", test_uc_ctl_page_size},
@@ -448,4 +490,5 @@ TEST_LIST = {
     {"test_tlb_clear", test_tlb_clear},
     {"test_noexec", test_noexec},
     {"test_add_block_hook", test_add_block_hook},
+    {"test_add_block_icount_hook", test_add_block_icount_hook},
     {NULL, NULL}};

uc.c

@@ -1927,6 +1927,30 @@ uc_err uc_hook_add(uc_engine *uc, uc_hook *hh, int type, void *callback,
 
     UC_INIT(uc);
 
+    int block_type = type & (UC_HOOK_BLOCK | UC_HOOK_BLOCK_ICOUNT);
+    if (block_type == (UC_HOOK_BLOCK | UC_HOOK_BLOCK_ICOUNT)) {
+        restore_jit_state(uc);
+        return UC_ERR_HOOK;
+    }
+    if (block_type != 0) {
+        int opposite_idx = block_type == UC_HOOK_BLOCK
+                               ? UC_HOOK_BLOCK_ICOUNT_IDX
+                               : UC_HOOK_BLOCK_IDX;
+        for (struct list_item *c = uc->hook[opposite_idx].head; c;
+             c = c->next) {
+            struct hook *other = c->data;
+            if (other->to_delete) {
+                continue;
+            }
+            bool overlaps = begin > end || other->begin > other->end ||
+                            (begin <= other->end && other->begin <= end);
+            if (overlaps) {
+                restore_jit_state(uc);
+                return UC_ERR_HOOK_EXIST;
+            }
+        }
+    }
+
     struct hook *hook = calloc(1, sizeof(struct hook));
     if (hook == NULL) {
         restore_jit_state(uc);
@@ -2013,7 +2037,8 @@ uc_err uc_hook_add(uc_engine *uc, uc_hook *hh, int type, void *callback,
         return UC_ERR_OK;
     }
 
-    if (type & UC_HOOK_CODE || type & UC_HOOK_BLOCK) {
+    if (type & UC_HOOK_CODE || type & UC_HOOK_BLOCK ||
+        type & UC_HOOK_BLOCK_ICOUNT) {
         if (end <= begin) {
             uc->tb_flush(uc);
         } else {
@@ -2073,7 +2098,8 @@ uc_err uc_hook_del(uc_engine *uc, uc_hook hh)
     // and store the type mask in the hook pointer.
     for (i = 0; i < UC_HOOK_MAX; i++) {
         if (list_exists(&uc->hook[i], (void *)hook)) {
-            if (hook->type & UC_HOOK_CODE || hook->type & UC_HOOK_BLOCK) {
+            if (hook->type & UC_HOOK_CODE || hook->type & UC_HOOK_BLOCK ||
+                hook->type & UC_HOOK_BLOCK_ICOUNT) {
                 g_hash_table_foreach(hook->hooked_regions,
                                      hook_invalidate_region, uc);
             }
@@ -2092,7 +2118,8 @@ UNICORN_EXPORT
 uc_err uc_hook_set_user_data(uc_engine *uc, uc_hook hh, void *user_data)
 {
     struct hook *hook = (struct hook *)hh;
-    if (hook->type == UC_HOOK_BLOCK || hook->type == UC_HOOK_CODE) {
+    if (hook->type == UC_HOOK_BLOCK || hook->type == UC_HOOK_CODE ||
+        hook->type == UC_HOOK_BLOCK_ICOUNT) {
         if (uc->nested_level) {
             return UC_ERR_ARG;
         }