Go: improve hook callback speed by 30% and add a HOOK_CODE benchmark (#835)

* add x86 hook benchmark

* Go: improve hook callback speed by 30%

bindings/go/unicorn/hook.go

@@ -19,54 +19,82 @@ type HookData struct {
 
 type Hook uint64
 
-var hookDataLock sync.RWMutex
-var hookDataMap = make(map[uintptr]*HookData)
-
-func getHookData(user unsafe.Pointer) *HookData {
-	hookDataLock.RLock()
-	defer hookDataLock.RUnlock()
-	return hookDataMap[uintptr(user)]
+type fastHookMap struct {
+	vals []*HookData
+	sync.RWMutex
 }
 
+func (m *fastHookMap) insert(h *HookData) uintptr {
+	// don't change this to defer
+	m.Lock()
+	for i, v := range m.vals {
+		if v == nil {
+			m.vals[i] = h
+			m.Unlock()
+			return uintptr(i)
+		}
+	}
+	i := len(m.vals)
+	m.vals = append(m.vals, h)
+	m.Unlock()
+	return uintptr(i)
+}
+
+func (m *fastHookMap) get(i unsafe.Pointer) *HookData {
+	m.RLock()
+	// TODO: nil check?
+	v := m.vals[uintptr(i)]
+	m.RUnlock()
+	return v
+}
+
+func (m *fastHookMap) remove(i uintptr) {
+	m.Lock()
+	m.vals[i] = nil
+	m.Unlock()
+}
+
+var hookMap fastHookMap
+
 //export hookCode
 func hookCode(handle unsafe.Pointer, addr uint64, size uint32, user unsafe.Pointer) {
-	hook := getHookData(user)
+	hook := hookMap.get(user)
 	hook.Callback.(func(Unicorn, uint64, uint32))(hook.Uc, uint64(addr), uint32(size))
 }
 
 //export hookMemInvalid
 func hookMemInvalid(handle unsafe.Pointer, typ C.uc_mem_type, addr uint64, size int, value int64, user unsafe.Pointer) bool {
-	hook := getHookData(user)
+	hook := hookMap.get(user)
 	return hook.Callback.(func(Unicorn, int, uint64, int, int64) bool)(hook.Uc, int(typ), addr, size, value)
 }
 
 //export hookMemAccess
 func hookMemAccess(handle unsafe.Pointer, typ C.uc_mem_type, addr uint64, size int, value int64, user unsafe.Pointer) {
-	hook := getHookData(user)
+	hook := hookMap.get(user)
 	hook.Callback.(func(Unicorn, int, uint64, int, int64))(hook.Uc, int(typ), addr, size, value)
 }
 
 //export hookInterrupt
 func hookInterrupt(handle unsafe.Pointer, intno uint32, user unsafe.Pointer) {
-	hook := getHookData(user)
+	hook := hookMap.get(user)
 	hook.Callback.(func(Unicorn, uint32))(hook.Uc, intno)
 }
 
 //export hookX86In
 func hookX86In(handle unsafe.Pointer, port, size uint32, user unsafe.Pointer) uint32 {
-	hook := getHookData(user)
+	hook := hookMap.get(user)
 	return hook.Callback.(func(Unicorn, uint32, uint32) uint32)(hook.Uc, port, size)
 }
 
 //export hookX86Out
 func hookX86Out(handle unsafe.Pointer, port, size, value uint32, user unsafe.Pointer) {
-	hook := getHookData(user)
+	hook := hookMap.get(user)
 	hook.Callback.(func(Unicorn, uint32, uint32, uint32))(hook.Uc, port, size, value)
 }
 
 //export hookX86Syscall
 func hookX86Syscall(handle unsafe.Pointer, user unsafe.Pointer) {
-	hook := getHookData(user)
+	hook := hookMap.get(user)
 	hook.Callback.(func(Unicorn))(hook.Uc)
 }
 
@@ -105,15 +133,13 @@ func (u *uc) HookAdd(htype int, cb interface{}, begin, end uint64, extra ...int)
 	}
 	var h2 C.uc_hook
 	data := &HookData{u, cb}
-	uptr := uintptr(unsafe.Pointer(data))
+	uptr := hookMap.insert(data)
 	if insnMode {
 		C.uc_hook_add_insn(u.handle, &h2, C.uc_hook_type(htype), callback, C.uintptr_t(uptr), C.uint64_t(begin), C.uint64_t(end), insn)
 	} else {
 		C.uc_hook_add_wrap(u.handle, &h2, C.uc_hook_type(htype), callback, C.uintptr_t(uptr), C.uint64_t(begin), C.uint64_t(end))
 	}
-	hookDataLock.Lock()
-	hookDataMap[uptr] = data
-	hookDataLock.Unlock()
+	// TODO: could move Hook and uptr onto HookData and just return it
 	u.hooks[Hook(h2)] = uptr
 	return Hook(h2), nil
 }
@@ -121,9 +147,7 @@ func (u *uc) HookAdd(htype int, cb interface{}, begin, end uint64, extra ...int)
 func (u *uc) HookDel(hook Hook) error {
 	if uptr, ok := u.hooks[hook]; ok {
 		delete(u.hooks, hook)
-		hookDataLock.Lock()
-		delete(hookDataMap, uptr)
-		hookDataLock.Unlock()
+		hookMap.remove(uptr)
 	}
 	return errReturn(C.uc_hook_del(u.handle, C.uc_hook(hook)))
 }

bindings/go/unicorn/unicorn.go

@@ -96,7 +96,7 @@ func (u *uc) Close() (err error) {
 	u.final.Do(func() {
 		if u.handle != nil {
 			for _, uptr := range u.hooks {
-				delete(hookDataMap, uptr)
+				hookMap.remove(uptr)
 			}
 			u.hooks = nil
 			err = errReturn(C.uc_close(u.handle))

bindings/go/unicorn/x86_test.go

@@ -158,3 +158,28 @@ func TestX86Mmr(t *testing.T) {
 		t.Fatalf("mmr read failed: %#v", mmr)
 	}
 }
+
+func BenchmarkX86Hook(b *testing.B) {
+	// loop rax times
+	code := "\x48\xff\xc8\x48\x83\xf8\x00\x0f\x8f\xf3\xff\xff\xff"
+	mu, err := MakeUc(MODE_64, code)
+	if err != nil {
+		b.Fatal(err)
+	}
+	count := 0
+	mu.HookAdd(HOOK_CODE, func(_ Unicorn, addr uint64, size uint32) {
+		count++
+	}, 1, 0)
+	mu.RegWrite(X86_REG_RAX, uint64(b.N))
+	b.ResetTimer()
+	if err := mu.Start(ADDRESS, ADDRESS+uint64(len(code))); err != nil {
+		b.Fatal(err)
+	}
+	rax, _ := mu.RegRead(X86_REG_RAX)
+	if rax != 0 {
+		b.Errorf("benchmark fell short: rax (%d) != 0", rax)
+	}
+	if count != b.N*3 {
+		b.Fatalf("benchmark fell short: %d < %d", count, b.N)
+	}
+}