-
Notifications
You must be signed in to change notification settings - Fork 1
Expand file tree
/
Copy pathmockery.go
More file actions
174 lines (143 loc) · 4.25 KB
/
Copy pathmockery.go
File metadata and controls
174 lines (143 loc) · 4.25 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
package mockery
import (
"errors"
"fmt"
"reflect"
"sync"
"unsafe"
)
type Patch struct {
targetBytes []byte
target *reflect.Value
replacement *reflect.Value
}
var (
lock = sync.Mutex{}
patches = make(map[uintptr]*Patch)
)
type value struct {
_ uintptr
ptr unsafe.Pointer
}
func getPtr(v *reflect.Value) unsafe.Pointer {
return (*value)(unsafe.Pointer(v)).ptr
}
// 把target方法替换为replacement方法
func PatchMethod(target, replacement interface{}) (*Patch, error) {
return patchValue(getValueFrom(target), getValueFrom(replacement))
}
func PatchInstanceMethodByName(target reflect.Type, methodName string, replacement interface{}) (*Patch, error) {
return PatchInstanceMethod(target, methodName, replacement)
}
// 把target结构体的methodName方法替换为replacement方法
func PatchInstanceMethod(target reflect.Type, methodName string, replacement interface{}) (*Patch, error) {
m, ok := target.MethodByName(methodName)
if !ok && target.Kind() == reflect.Struct {
target = reflect.PtrTo(target)
m, ok = target.MethodByName(methodName)
}
if !ok {
return nil, fmt.Errorf("method '%s' not found", methodName)
}
return PatchMethodByReflect(m, replacement)
}
// 把target反射方法替换为replacement方法
func PatchMethodByReflect(target reflect.Method, replacement interface{}) (*Patch, error) {
return PatchMethodByReflectValue(target.Func, replacement)
}
// 把target反射方法使用自定义函数来替换
func PatchMethodWithMakeFunc(target reflect.Method, fn func(args []reflect.Value) (results []reflect.Value)) (*Patch, error) {
return PatchMethodByReflect(target, reflect.MakeFunc(target.Type, fn))
}
// 把target反射值使用自定义函数来替换
func PatchMethodWithMakeFuncValue(target reflect.Value, fn func(args []reflect.Value) (results []reflect.Value)) (*Patch, error) {
return PatchMethodByReflectValue(target, reflect.MakeFunc(target.Type(), fn))
}
// 把target反射值替换为replacement方法
func PatchMethodByReflectValue(target reflect.Value, replacement interface{}) (*Patch, error) {
r := getValueFrom(replacement)
return patchValue(&target, r)
}
func (p *Patch) Patch() error {
if p == nil {
return errors.New("patch is nil")
}
if err := isPatchable(p.target, p.replacement); err != nil {
return err
}
if err := applyPatch(p); err != nil {
return err
}
return nil
}
func (p *Patch) Unpatch() error {
if p == nil {
return errors.New("patch is nil")
}
return unpatchValue(*p.target)
}
// interface{} to *reflect.Value
func getValueFrom(data interface{}) *reflect.Value {
if v, ok := data.(reflect.Value); ok {
return &v
} else {
v = reflect.ValueOf(data)
return &v
}
}
func isPatchable(target, replacement *reflect.Value) error {
lock.Lock()
defer lock.Unlock()
if target.Kind() != reflect.Func {
return errors.New("the target is not a Func")
}
if replacement.Kind() != reflect.Func {
return errors.New("the replacement is not a Func")
}
if target.Type() != replacement.Type() {
return fmt.Errorf("the target and redirection doesn't have the same type: %s != %s", target.Type(), replacement.Type())
}
if _, ok := patches[target.Pointer()]; ok {
return errors.New("the target is already patched")
}
return nil
}
func applyPatch(patch *Patch) (err error) {
lock.Lock()
defer lock.Unlock()
patch.targetBytes, err = replaceFunction(patch.target.Pointer(), (uintptr)(getPtr(patch.replacement)))
if err != nil {
return
}
patches[patch.target.Pointer()] = patch
return nil
}
func patchValue(target, replacement *reflect.Value) (*Patch, error) {
if err := isPatchable(target, replacement); err != nil {
return nil, err
}
patch := &Patch{target: target, replacement: replacement}
if err := applyPatch(patch); err != nil {
return nil, err
}
return patch, nil
}
func unpatchValue(target reflect.Value) error {
lock.Lock()
defer lock.Unlock()
patch, ok := patches[target.Pointer()]
if !ok {
return errors.New("the target is not patched")
}
if patch.targetBytes == nil || len(patch.targetBytes) == 0 {
return errors.New("the target is not patched")
}
if err := unpatch(target.Pointer(), patch); err != nil {
return err
}
delete(patches, target.Pointer())
return nil
}
func unpatch(target uintptr, p *Patch) error {
return copyToLocation(target, p.targetBytes)
}