feat: add method unswizzling

CHANGELOG.md

@@ -14,6 +14,7 @@
 
 - `SentrySdkInfo.packages` should be an array (#4626)
 - Use the same SdkInfo for envelope header and event (#4629)
+- Add method unswizzling (#4647)
 
 ### Internal
 

Sources/Sentry/SentrySwizzle.m

@@ -40,8 +40,46 @@
 
 @implementation SentrySwizzle
 
+static NSMutableDictionary<NSValue *, NSValue *> *
+refsToOriginalImplementationsDictionary(void)
+{
+    static NSMutableDictionary *refsToOriginalImplementations;
+    static dispatch_once_t onceToken;
+    dispatch_once(&onceToken, ^{ refsToOriginalImplementations = [NSMutableDictionary new]; });
+    return refsToOriginalImplementations;
+}
+
+static void
+storeRefToOriginalImplementation(const void *key, IMP implementation)
+{
+    NSMutableDictionary<NSValue *, NSValue *> *refsToOriginalImplementations
+        = refsToOriginalImplementationsDictionary();
+    NSValue *keyValue = [NSValue valueWithPointer:key];
+    refsToOriginalImplementations[keyValue] = [NSValue valueWithPointer:implementation];
+}
+
+static void
+removeRefToOriginalImplementation(const void *key)
+{
+    NSMutableDictionary<NSValue *, NSValue *> *refsToOriginalImplementations
+        = refsToOriginalImplementationsDictionary();
+    NSValue *keyValue = [NSValue valueWithPointer:key];
+    [refsToOriginalImplementations removeObjectForKey:keyValue];
+}
+
+static IMP
+getRefToOriginalImplementation(const void *key)
+{
+    NSMutableDictionary<NSValue *, NSValue *> *refsToOriginalImplementations
+        = refsToOriginalImplementationsDictionary();
+    NSValue *keyValue = [NSValue valueWithPointer:key];
+    NSValue *originalImplementationValue = [refsToOriginalImplementations objectForKey:keyValue];
+    return (IMP)[originalImplementationValue pointerValue];
+}
+
 static void
-swizzle(Class classToSwizzle, SEL selector, SentrySwizzleImpFactoryBlock factoryBlock)
+swizzle(
+    Class classToSwizzle, SEL selector, SentrySwizzleImpFactoryBlock factoryBlock, const void *key)
 {
     Method method = class_getInstanceMethod(classToSwizzle, selector);
 
@@ -106,6 +144,33 @@
     pthread_mutex_lock(&gLock);
 
     originalIMP = class_replaceMethod(classToSwizzle, selector, newIMP, methodType);
+    if (originalIMP) {
+        storeRefToOriginalImplementation(key, originalIMP);
+    }
+
+    pthread_mutex_unlock(&gLock);
+}
+
+static void
+unswizzle(Class classToUnswizzle, SEL selector, const void *key)
+{
+    Method method = class_getInstanceMethod(classToUnswizzle, selector);
+
+    NSCAssert(NULL != method, @"Selector %@ not found in %@ methods of class %@.",
+        NSStringFromSelector(selector),
+        class_isMetaClass(classToUnswizzle) ? @"class" : @"instance", classToUnswizzle);
+
+    static pthread_mutex_t gLock = PTHREAD_MUTEX_INITIALIZER;
+
+    pthread_mutex_lock(&gLock);
+
+    IMP originalIMP = getRefToOriginalImplementation(key);
+    if (originalIMP) {
+        const char *methodType = method_getTypeEncoding(method);
+        class_replaceMethod(classToUnswizzle, selector, originalIMP, methodType);
+
+        removeRefToOriginalImplementation(key);
+    }
 
     pthread_mutex_unlock(&gLock);
 }
@@ -164,7 +229,7 @@
             }
         }
 
-        swizzle(classToSwizzle, selector, factoryBlock);
+        swizzle(classToSwizzle, selector, factoryBlock, key);
 
         if (key) {
             [swizzledClassesForKey(key) addObject:classToSwizzle];
@@ -174,6 +239,29 @@
     return YES;
 }
 
++ (BOOL)unswizzleInstanceMethod:(SEL)selector inClass:(Class)classToUnswizzle key:(const void *)key
+{
+    NSAssert(key != NULL, @"Key may not be NULL.");
+
+    if (key == NULL) {
+        NSLog(@"Key may not be NULL.");
+        return NO;
+    }
+
+    @synchronized(swizzledClassesDictionary()) {
+        NSSet<Class> *swizzledClasses = swizzledClassesForKey(key);
+        if (![swizzledClasses containsObject:classToUnswizzle]) {
+            return NO;
+        }
+
+        unswizzle(classToUnswizzle, selector, key);
+
+        [swizzledClassesForKey(key) removeObject:classToUnswizzle];
+    }
+
+    return YES;
+}
+
 + (void)swizzleClassMethod:(SEL)selector
                    inClass:(Class)classToSwizzle
              newImpFactory:(SentrySwizzleImpFactoryBlock)factoryBlock

Sources/Sentry/include/HybridPublic/SentrySwizzle.h

@@ -83,6 +83,18 @@
         _SentrySWWrapArg(SentrySWArguments), _SentrySWWrapArg(SentrySWReplacement),                \
         SentrySwizzleMode, key)
 
+/**
+ * Unswizzles the instance method of the class.
+ *
+ * @param classToUnswizzle The class with the method that should be unswizzled.
+ *
+ * @param selector Selector of the method that should be unswizzled.
+ *
+ * @return @c YES if successfully unswizzled and @c NO if the method was not swizzled.
+ */
+#define SentryUnswizzleInstanceMethod(classToUnswizzle, selector, key)                             \
+    _SentryUnswizzleInstanceMethod(classToUnswizzle, selector, key)
+
 #pragma mark └ Swizzle Class Method
 
 /**
@@ -302,6 +314,20 @@ typedef NS_ENUM(NSUInteger, SentrySwizzleMode) {
                          mode:(SentrySwizzleMode)mode
                           key:(const void *)key;
 
+/**
+ * Unswizzles the instance method of the class.
+ *
+ * @param selector Selector of the method that should be unswizzled.
+ *
+ * @param classToUnswizzle The class with the method that should be unswizzled.
+ *
+ * @param key The key is used in combination with the mode to indicate whether the
+ * swizzling should be done for the given class.
+ *
+ * @return @c YES if successfully unswizzled and @c NO if the method was not swizzled.
+ */
++ (BOOL)unswizzleInstanceMethod:(SEL)selector inClass:(Class)classToUnswizzle key:(const void *)key;
+
 #pragma mark └ Swizzle Class method
 
 /**
@@ -396,6 +422,9 @@ typedef NS_ENUM(NSUInteger, SentrySwizzleMode) {
                          mode:SentrySwizzleMode                                                    \
                           key:KEY];
 
+#define _SentryUnswizzleInstanceMethod(classToUnswizzle, selector, KEY)                            \
+    [SentrySwizzle unswizzleInstanceMethod:selector inClass:[classToUnswizzle class] key:KEY]
+
 #define _SentrySwizzleClassMethod(                                                                 \
     classToSwizzle, selector, SentrySWReturnType, SentrySWArguments, SentrySWReplacement)          \
     [SentrySwizzle                                                                                 \

Tests/SentryTests/Helper/SentrySwizzleTests.m

@@ -83,6 +83,10 @@ - (void)methodForSwizzlingWithoutCallOriginal
 {
 };
 
+- (void)methodForUnswizzling
+{
+};
+
 - (NSString *)string
 {
     return @"ABC";
@@ -353,4 +357,23 @@ - (void)testSwizzleDontCallOriginalImplementation
     XCTAssertThrows([a methodForSwizzlingWithoutCallOriginal]);
 }
 
+- (void)testUnswizzleInstanceMethod
+{
+    SEL methodForUnswizzling = NSSelectorFromString(@"methodForUnswizzling");
+
+    SentrySwizzleTestClass_A *object = [SentrySwizzleTestClass_B new];
+    swizzleVoidMethod(
+        [SentrySwizzleTestClass_A class], methodForUnswizzling, ^{ SentryTestsLog(@"A"); },
+        SentrySwizzleModeAlways, (void *)methodForUnswizzling);
+    [object methodForUnswizzling];
+    ASSERT_LOG_IS(@"A");
+
+    [SentryTestsLog clear];
+
+    [SentrySwizzle unswizzleInstanceMethod:methodForUnswizzling
+                                   inClass:[SentrySwizzleTestClass_A class]
+                                       key:(void *)methodForUnswizzling];
+    [object methodForUnswizzling];
+    ASSERT_LOG_IS(@"");
+}
 @end

scripts/.clang-format-version

@@ -1 +1 @@
-19.1.5
+19.1.6