feat: add method unswizzling

CHANGELOG.md

@@ -9,6 +9,7 @@
 ### Internal
 
 - Update to Xcode 16.2 in workflows (#4673)
+- Add method unswizzling (#4647)
 
 ## 8.43.0
 

Sources/Sentry/SentrySwizzle.m

@@ -1,4 +1,5 @@
 #import "SentrySwizzle.h"
+#import "SentryLog.h"
 
 #import <objc/runtime.h>
 #include <pthread.h>
@@ -20,7 +21,7 @@ - (SentrySwizzleOriginalIMP)getOriginalImplementation
 {
     NSAssert(_impProviderBlock, @"_impProviderBlock can't be missing");
     if (!_impProviderBlock) {
-        NSLog(@"_impProviderBlock can't be missing");
+        SENTRY_LOG_ERROR(@"_impProviderBlock can't be missing");
         return NULL;
     }
 
@@ -40,17 +41,104 @@ - (SentrySwizzleOriginalIMP)getOriginalImplementation
 
 @implementation SentrySwizzle
 
+// This lock is shared by all swizzling and unswizzling calls to ensure that
+// only one thread is modifying the class at a time.
+static pthread_mutex_t gLock = PTHREAD_MUTEX_INITIALIZER;
+
+#if TEST || TESTCI
+/**
+ * - Returns: a dictionary that maps keys to the references to the original implementations.
+ */
+static NSMutableDictionary<NSValue *, NSValue *> *
+refsToOriginalImplementationsDictionary(void)
+{
+    static NSMutableDictionary *refsToOriginalImplementations;
+    static dispatch_once_t onceToken;
+    dispatch_once(&onceToken, ^{ refsToOriginalImplementations = [NSMutableDictionary new]; });
+    return refsToOriginalImplementations;
+}
+
+/**
+ * Adds a reference to the original implementation to the dictionary.
+ *
+ * If the key is NULL, it will log an error and NOT store the reference.
+ *
+ * - Parameter key: The key for which to store the reference to the original implementation.
+ * - Parameter implementation: Reference to the original implementation to store.
+ */
+static void
+storeRefToOriginalImplementation(const void *key, IMP implementation)
+{
+    NSCAssert(key != NULL, @"Key may not be NULL.");
+    if (key == NULL) {
+        SENTRY_LOG_ERROR(@"Key may not be NULL.");
+        return;
+    }
+    NSMutableDictionary<NSValue *, NSValue *> *refsToOriginalImplementations
+        = refsToOriginalImplementationsDictionary();
+    NSValue *keyValue = [NSValue valueWithPointer:key];
+    refsToOriginalImplementations[keyValue] = [NSValue valueWithPointer:implementation];
+}
+
+/**
+ * Removes a reference to the original implementation from the dictionary.
+ *
+ * If the key is NULL, it will log an error and do nothing.
+ *
+ * - Parameter key: The key for which to remove the reference to the original implementation.
+ */
+static void
+removeRefToOriginalImplementation(const void *key)
+{
+    NSCAssert(key != NULL, @"Key may not be NULL.");
+    if (key == NULL) {
+        SENTRY_LOG_ERROR(@"Key may not be NULL.");
+        return;
+    }
+    NSMutableDictionary<NSValue *, NSValue *> *refsToOriginalImplementations
+        = refsToOriginalImplementationsDictionary();
+    NSValue *keyValue = [NSValue valueWithPointer:key];
+    [refsToOriginalImplementations removeObjectForKey:keyValue];
+}
+
+/**
+ * Returns the original implementation for the given key.
+ *
+ * If the key is NULL, it will log an error and return NULL.
+ * If no original implementation is found, it will return NULL.
+ *
+ * - Parameter key: The key for which to get the original implementation.
+ * - Returns: The original implementation for the given key.
+ */
+static IMP
+getRefToOriginalImplementation(const void *key)
+{
+    NSCAssert(key != NULL, @"Key may not be NULL.");
+    if (key == NULL) {
+        SENTRY_LOG_ERROR(@"Key may not be NULL.");
+        return NULL;
+    }
+    NSMutableDictionary<NSValue *, NSValue *> *refsToOriginalImplementations
+        = refsToOriginalImplementationsDictionary();
+    NSValue *keyValue = [NSValue valueWithPointer:key];
+    NSValue *originalImplementationValue = [refsToOriginalImplementations objectForKey:keyValue];
+    if (originalImplementationValue == nil) {
+        return NULL;
+    }
+    return (IMP)[originalImplementationValue pointerValue];
+}
+#endif // TEST || TESTCI
+
 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);
 
     NSCAssert(NULL != method, @"Selector %@ not found in %@ methods of class %@.",
         NSStringFromSelector(selector), class_isMetaClass(classToSwizzle) ? @"class" : @"instance",
         classToSwizzle);
 
-    static pthread_mutex_t gLock = PTHREAD_MUTEX_INITIALIZER;
-
     // To keep things thread-safe, we fill in the originalIMP later,
     // with the result of the class_replaceMethod call below.
     __block IMP originalIMP = NULL;
@@ -106,10 +194,50 @@ @implementation SentrySwizzle
     pthread_mutex_lock(&gLock);
 
     originalIMP = class_replaceMethod(classToSwizzle, selector, newIMP, methodType);
+#if TEST || TESTCI
+    if (originalIMP) {
+        if (key != NULL) {
+            storeRefToOriginalImplementation(key, originalIMP);
+        } else {
+            SENTRY_LOG_WARN(
+                @"Swizzling without a key is not recommended, because they can not be unswizzled.");
+        }
+    }
+#endif // TEST || TESTCI
 
     pthread_mutex_unlock(&gLock);
 }
 
+#if TEST || TESTCI
+static void
+unswizzle(Class classToUnswizzle, SEL selector, const void *key)
+{
+    NSCAssert(key != NULL, @"Key may not be NULL.");
+    if (key == NULL) {
+        SENTRY_LOG_WARN(@"Key may not be NULL.");
+        return;
+    }
+
+    Method method = class_getInstanceMethod(classToUnswizzle, selector);
+
+    NSCAssert(NULL != method, @"Selector %@ not found in %@ methods of class %@.",
+        NSStringFromSelector(selector),
+        class_isMetaClass(classToUnswizzle) ? @"class" : @"instance", classToUnswizzle);
+
+    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);
+}
+#endif // TEST || TESTCI
+
 static NSMutableDictionary<NSValue *, NSMutableSet<Class> *> *
 swizzledClassesDictionary(void)
 {
@@ -143,7 +271,7 @@ + (BOOL)swizzleInstanceMethod:(SEL)selector
         @"Key may not be NULL if mode is not SentrySwizzleModeAlways.");
 
     if (key == NULL && mode != SentrySwizzleModeAlways) {
-        NSLog(@"Key may not be NULL if mode is not SentrySwizzleModeAlways.");
+        SENTRY_LOG_WARN(@"Key may not be NULL if mode is not SentrySwizzleModeAlways.");
         return NO;
     }
 
@@ -164,7 +292,7 @@ + (BOOL)swizzleInstanceMethod:(SEL)selector
             }
         }
 
-        swizzle(classToSwizzle, selector, factoryBlock);
+        swizzle(classToSwizzle, selector, factoryBlock, key);
 
         if (key) {
             [swizzledClassesForKey(key) addObject:classToSwizzle];
@@ -174,6 +302,30 @@ + (BOOL)swizzleInstanceMethod:(SEL)selector
     return YES;
 }
 
+#if TEST || TESTCI
++ (BOOL)unswizzleInstanceMethod:(SEL)selector inClass:(Class)classToUnswizzle key:(const void *)key
+{
+    NSAssert(key != NULL, @"Key may not be NULL.");
+    if (key == NULL) {
+        SENTRY_LOG_WARN(@"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;
+}
+#endif // TEST || TESTCI
+
 + (void)swizzleClassMethod:(SEL)selector
                    inClass:(Class)classToSwizzle
              newImpFactory:(SentrySwizzleImpFactoryBlock)factoryBlock

Sources/Sentry/include/HybridPublic/SentrySwizzle.h

@@ -83,6 +83,23 @@
         _SentrySWWrapArg(SentrySWArguments), _SentrySWWrapArg(SentrySWReplacement),                \
         SentrySwizzleMode, key)
 
+#if TEST || TESTCI
+/**
+ * Unswizzles the instance method of the class.
+ *
+ * @warning To reduce the risk of breaking functionality with unswizzling, this method is not
+ * considered safe-to-use in production and only available in test targets.
+ *
+ * @param classToUnswizzle The class with the method that should be unswizzled.
+ * @param selector Selector of the method that should be unswizzled.
+ * @param key The key to unswizzle the method with.
+ *
+ * @return @c YES if successfully unswizzled and @c NO if the method was not swizzled.
+ */
+#    define SentryUnswizzleInstanceMethod(classToUnswizzle, selector, key)                         \
+        _SentryUnswizzleInstanceMethod(classToUnswizzle, selector, key)
+#endif // TEST || TESTCI
+
 #pragma mark └ Swizzle Class Method
 
 /**
@@ -302,6 +319,23 @@ typedef NS_ENUM(NSUInteger, SentrySwizzleMode) {
                          mode:(SentrySwizzleMode)mode
                           key:(const void *)key;
 
+#if TEST || TESTCI
+/**
+ * Unswizzles the instance method of the class.
+ *
+ * @warning To reduce the risk of breaking functionality with unswizzling, this method is not
+ * considered safe-to-use in production and only available in test targets.
+ *
+ * @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;
+#endif // TEST || TESTCI
+
 #pragma mark └ Swizzle Class method
 
 /**
@@ -396,6 +430,22 @@ typedef NS_ENUM(NSUInteger, SentrySwizzleMode) {
                          mode:SentrySwizzleMode                                                    \
                           key:KEY];
 
+#if TEST || TESTCI
+/**
+ * Macro to unswizzle an instance method.
+ *
+ * @warning To reduce the risk of breaking functionality with unswizzling, this macro is not
+ * considered safe-to-use in production and only available in test targets.
+ *
+ * @param classToUnswizzle The class to unswizzle the method from.
+ * @param selector The selector of the method to unswizzle.
+ * @param KEY The key to unswizzle the method with.
+ * @return @c YES if the method was successfully unswizzled, @c NO otherwise.
+ */
+#    define _SentryUnswizzleInstanceMethod(classToUnswizzle, selector, KEY)                        \
+        [SentrySwizzle unswizzleInstanceMethod:selector inClass:[classToUnswizzle class] key:KEY]
+#endif // TEST || TESTCI
+
 #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,84 @@ - (void)testSwizzleDontCallOriginalImplementation
     XCTAssertThrows([a methodForSwizzlingWithoutCallOriginal]);
 }
 
+- (void)testUnswizzleInstanceMethod
+{
+    // -- Arrange --
+    SEL methodForUnswizzling = NSSelectorFromString(@"methodForUnswizzling");
+    SentrySwizzleTestClass_A *object = [SentrySwizzleTestClass_B new];
+
+    // Swizzle the method once
+    swizzleVoidMethod(
+        [SentrySwizzleTestClass_A class], methodForUnswizzling, ^{ SentryTestsLog(@"A"); },
+        SentrySwizzleModeAlways, (void *)methodForUnswizzling);
+
+    // Smoke test the swizzling
+    [object methodForUnswizzling];
+    ASSERT_LOG_IS(@"A");
+    [SentryTestsLog clear];
+
+    // -- Act --
+    [SentrySwizzle unswizzleInstanceMethod:methodForUnswizzling
+                                   inClass:[SentrySwizzleTestClass_A class]
+                                       key:(void *)methodForUnswizzling];
+    [object methodForUnswizzling];
+
+    // -- Assert --
+    ASSERT_LOG_IS(@"");
+}
+
+- (void)testUnswizzleInstanceMethod_methodNotSwizzled_shouldWork
+{
+    // -- Arrange --
+    SEL methodForUnswizzling = NSSelectorFromString(@"methodForUnswizzling");
+    SentrySwizzleTestClass_A *object = [SentrySwizzleTestClass_A new];
+
+    // Smoke-test the swizzling
+    [object methodForUnswizzling];
+    ASSERT_LOG_IS(@"");
+    [SentryTestsLog clear];
+
+    // -- Act --
+    [SentrySwizzle unswizzleInstanceMethod:methodForUnswizzling
+                                   inClass:[SentrySwizzleTestClass_A class]
+                                       key:(void *)methodForUnswizzling];
+    [object methodForUnswizzling];
+
+    // -- Assert --
+    ASSERT_LOG_IS(@"");
+}
+
+- (void)testUnswizzleInstanceMethod_unswizzlingMethodMultipleTimes_shouldWork
+{
+    // -- Arrange --
+    SEL methodForUnswizzling = NSSelectorFromString(@"methodForUnswizzling");
+    SentrySwizzleTestClass_A *object = [SentrySwizzleTestClass_A new];
+
+    swizzleVoidMethod(
+        [SentrySwizzleTestClass_A class], methodForUnswizzling, ^{ SentryTestsLog(@"A"); },
+        SentrySwizzleModeAlways, (void *)methodForUnswizzling);
+
+    // Smoke test the swizzling
+    [object methodForUnswizzling];
+    ASSERT_LOG_IS(@"A");
+    [SentryTestsLog clear];
+
+    [SentrySwizzle unswizzleInstanceMethod:methodForUnswizzling
+                                   inClass:[SentrySwizzleTestClass_A class]
+                                       key:(void *)methodForUnswizzling];
+    [object methodForUnswizzling];
+    ASSERT_LOG_IS(@"");
+    [SentryTestsLog clear];
+
+    // -- Act --
+    // Unswizzle again should not cause issues
+    [SentrySwizzle unswizzleInstanceMethod:methodForUnswizzling
+                                   inClass:[SentrySwizzleTestClass_A class]
+                                       key:(void *)methodForUnswizzling];
+    [object methodForUnswizzling];
+
+    // -- Assert -
+    ASSERT_LOG_IS(@"");
+}
+
 @end