class_replaceMethod return originalImplementation = nil bug

steipete · Feb 24, 2017 · ecdd4e1 · ecdd4e1
1 parent e8a3e19
commit ecdd4e1
Showing 1 changed file with 10 additions and 3 deletions.
diff --git a/Aspects.m b/Aspects.m
@@ -391,10 +391,17 @@ static Class aspect_hookClass(NSObject *self, NSError **error) {
 static void aspect_swizzleForwardInvocation(Class klass) {
     NSCParameterAssert(klass);
     // If there is no method, replace will act like class_addMethod.
-    IMP originalImplementation = class_replaceMethod(klass, @selector(forwardInvocation:), (IMP)__ASPECTS_ARE_BEING_CALLED__, "v@:@");
+    IMP originalImplementation = class_getMethodImplementation(klass, @selector(forwardInvocation:));
+    Method method = class_getInstanceMethod(klass, @selector(forwardInvocation:));
+    method_setImplementation(method, (IMP)__ASPECTS_ARE_BEING_CALLED__);
     if (originalImplementation) {
-        class_addMethod(klass, NSSelectorFromString(AspectsForwardInvocationSelectorName), originalImplementation, "v@:@");
-    }
+          BOOL add= class_addMethod(klass, NSSelectorFromString(AspectsForwardInvocationSelectorName), originalImplementation, "v@:@");
+          if (!add) {
+              Method ORIGMethod = class_getInstanceMethod(klass, NSSelectorFromString(AspectsForwardInvocationSelectorName));
+              method_setImplementation(ORIGMethod, originalImplementation);
+          }
+  }
+
     AspectLog(@"Aspects: %@ is now aspect aware.", NSStringFromClass(klass));
 }