Merge pull request #17 from Mutated1994/master

fix: Fix panic when calling PatchGuard.Restore
cch123 · Aug 1, 2021 · 16a9bd9 · 16a9bd9
2 parents d792ef7 + d203a03
commit 16a9bd9
Show file tree

Hide file tree

Showing 7 changed files with 59 additions and 11 deletions.
diff --git a/examples/patch_func.go b/examples/patch_func.go
@@ -20,6 +20,14 @@ func patchFunc() {
 	fmt.Println("unpatch, then output:")
 	heyHey()
 
+	patchGuard.Restore()
+	fmt.Println("restore, then output:")
+	heyHey()
+
+	patchGuard.Unpatch()
+	fmt.Println("unpatch, then output:")
+	heyHey()
+
 	fmt.Println()
 }
 

diff --git a/examples/patch_func_symbol.go b/examples/patch_func_symbol.go
@@ -20,6 +20,14 @@ func patchFuncSymbol() {
 	fmt.Println("unpatch, then output:")
 	heyHeyHey()
 
+	patchGuard.Restore()
+	fmt.Println("restore, then output:")
+	heyHeyHey()
+
+	patchGuard.Unpatch()
+	fmt.Println("unpatch, then output:")
+	heyHeyHey()
+
 	fmt.Println()
 }
 

diff --git a/examples/patch_instance_func.go b/examples/patch_instance_func.go
@@ -28,5 +28,13 @@ func patchInstanceFunc() {
 	fmt.Println("unpatch, then output:")
 	p.speak()
 
+	patchGuard.Restore()
+	fmt.Println("restore, then output:")
+	p.speak()
+
+	patchGuard.Unpatch()
+	fmt.Println("unpatch, then output:")
+	p.speak()
+
 	fmt.Println()
 }
diff --git a/examples/patch_instance_func_symbol.go b/examples/patch_instance_func_symbol.go
@@ -23,5 +23,13 @@ func patchInstanceFuncSymbol() {
 	fmt.Println("unpatch, then output:")
 	p.speak()
 
+	patchGuard.Restore()
+	fmt.Println("restore, then output:")
+	p.speak()
+
+	patchGuard.Unpatch()
+	fmt.Println("unpatch, then output:")
+	p.speak()
+
 	fmt.Println()
 }
diff --git a/examples/patch_struct_method.go b/examples/patch_struct_method.go
@@ -34,5 +34,13 @@ func patchStructMethod() {
 	fmt.Println("unpatch, then output:")
 	fmt.Println(f.MyFunc(nil))
 
+	patchGuard.Restore()
+	fmt.Println("restore, then output:")
+	fmt.Println(f.MyFunc(nil))
+
+	patchGuard.Unpatch()
+	fmt.Println("unpatch, then output:")
+	fmt.Println(f.MyFunc(nil))
+
 	fmt.Println()
-}
+}
diff --git a/examples/patch_struct_method_symbol.go b/examples/patch_struct_method_symbol.go
@@ -24,6 +24,13 @@ func patchStructMethodSymbol() {
 	fmt.Println("unpatch, then output:")
 	fmt.Println(f.MyFunc(nil))
 
+	patchGuard.Restore()
+	fmt.Println("restore, then output:")
+	fmt.Println(f.MyFunc(nil))
+
+	patchGuard.Unpatch()
+	fmt.Println("unpatch, then output:")
+	fmt.Println(f.MyFunc(nil))
+
 	fmt.Println()
 }
-
diff --git a/internal/bouk/monkey.go b/internal/bouk/monkey.go
@@ -32,31 +32,32 @@ func getPtr(v reflect.Value) unsafe.Pointer {
 type PatchGuard struct {
 	target      reflect.Value
 	replacement reflect.Value
+	isSymbol    bool
 }
 
 func (g *PatchGuard) Unpatch() {
 	unpatchValue(g.target)
 }
 
 func (g *PatchGuard) Restore() {
-	patchValue(g.target, g.replacement)
+	patchValue(g.target, g.replacement, g.isSymbol)
 }
 
 // Patch replaces a function with another
 func Patch(target, replacement interface{}) *PatchGuard {
 	t := reflect.ValueOf(target)
 	r := reflect.ValueOf(replacement)
-	patchValue(t, r)
+	patchValue(t, r, false)
 
-	return &PatchGuard{t, r}
+	return &PatchGuard{t, r, false}
 }
 
 func PatchSymbol(target, replacement interface{}) *PatchGuard {
 	t := reflect.ValueOf(target)
 	r := reflect.ValueOf(replacement)
 	patchSymbolValue(t, r)
 
-	return &PatchGuard{t, r}
+	return &PatchGuard{t, r, true}
 }
 
 // PatchInstanceMethod replaces an instance method methodName for the type target with replacement
@@ -67,24 +68,24 @@ func PatchInstanceMethod(target reflect.Type, methodName string, replacement int
 		panic(fmt.Sprintf("unknown method %s", methodName))
 	}
 	r := reflect.ValueOf(replacement)
-	patchValue(m.Func, r)
+	patchValue(m.Func, r, false)
 
-	return &PatchGuard{m.Func, r}
+	return &PatchGuard{m.Func, r, false}
 }
 
-func patchValue(target, replacement reflect.Value) {
+func patchValue(target, replacement reflect.Value, isSymbol bool) {
 	lock.Lock()
 	defer lock.Unlock()
 
-	if target.Kind() != reflect.Func {
+	if !isSymbol && target.Kind() != reflect.Func {
 		panic("target has to be a Func")
 	}
 
 	if replacement.Kind() != reflect.Func {
 		panic("replacement has to be a Func")
 	}
 
-	if target.Type() != replacement.Type() {
+	if !isSymbol && target.Type() != replacement.Type() {
 		panic(fmt.Sprintf("target and replacement have to have the same type %s != %s", target.Type(), replacement.Type()))
 	}