From 4b50ca5cec7ae0a379a59b65058ae264b6cfdf64 Mon Sep 17 00:00:00 2001 From: Ryan Hileman Date: Sat, 13 May 2017 10:12:57 -0700 Subject: [PATCH] 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 | 66 ++++++++++++++++++++++----------- bindings/go/unicorn/unicorn.go | 2 +- bindings/go/unicorn/x86_test.go | 25 +++++++++++++ 3 files changed, 71 insertions(+), 22 deletions(-) diff --git a/bindings/go/unicorn/hook.go b/bindings/go/unicorn/hook.go index a51f83b8..0d7fd681 100644 --- a/bindings/go/unicorn/hook.go +++ b/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))) } diff --git a/bindings/go/unicorn/unicorn.go b/bindings/go/unicorn/unicorn.go index 9b3d6272..28bd467f 100644 --- a/bindings/go/unicorn/unicorn.go +++ b/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)) diff --git a/bindings/go/unicorn/x86_test.go b/bindings/go/unicorn/x86_test.go index c74a0611..803c6463 100644 --- a/bindings/go/unicorn/x86_test.go +++ b/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) + } +}