Update the GULObjectSwizzler to handle NSProxy objects (#2053)

tejasd · web-flow · commit 829b9f557cb5 · 2018-11-13T09:35:35.000-08:00
diff --git a/GoogleUtilities/CHANGELOG.md b/GoogleUtilities/CHANGELOG.md
@@ -2,6 +2,9 @@
 - Fixed an issue where GoogleUtilities would leak non-background URL sessions.
   (#2061)
 
+- Fixed a crash caused due to `NSURLConnection` delegates being wrapped in an
+  `NSProxy`. (#1936)
+
 # 5.3.4
 - Fixed a crash caused by unprotected access to sessions in
   `GULNetworkURLSession` (#1964).
diff --git a/GoogleUtilities/Example/Tests/Swizzler/GULObjectSwizzlerTest.m b/GoogleUtilities/Example/Tests/Swizzler/GULObjectSwizzlerTest.m
@@ -13,8 +13,10 @@
 // limitations under the License.
 
 #import <XCTest/XCTest.h>
+#import <objc/runtime.h>
 
 #import <GoogleUtilities/GULObjectSwizzler.h>
+#import <GoogleUtilities/GULProxy.h>
 #import <GoogleUtilities/GULSwizzledObject.h>
 
 @interface GULObjectSwizzlerTest : XCTestCase
@@ -23,6 +25,16 @@ @interface GULObjectSwizzlerTest : XCTestCase
 
 @implementation GULObjectSwizzlerTest
 
+/** Used as a donor method to add a method that doesn't exist on the superclass. */
+- (NSString *)donorDescription {
+  return @"SwizzledDonorDescription";
+}
+
+/** Used as a donor method to add a method that exists on the superclass. */
+- (NSString *)description {
+  return @"SwizzledDescription";
+}
+
 /** Exists just as a donor method. */
 - (void)donorMethod {
 }
@@ -226,6 +238,7 @@ - (void)testSetGetAssociatedObjectRetainNonatomic {
   XCTAssertEqualObjects(returnedObject, associatedObject);
 }
 
+/** Tests getting and setting an associated object with an invalid association type. */
 - (void)testSetGetAssociatedObjectWithoutProperAssociation {
   NSObject *object = [[NSObject alloc] init];
   NSDictionary *associatedObject = [[NSDictionary alloc] init];
@@ -235,4 +248,105 @@ - (void)testSetGetAssociatedObjectWithoutProperAssociation {
   XCTAssertEqualObjects(returnedObject, associatedObject);
 }
 
+/** Tests using the GULObjectSwizzler to swizzle an object wrapped in an NSProxy. */
+- (void)testSwizzleProxiedObject {
+  NSObject *object = [[NSObject alloc] init];
+  GULProxy *proxyObject = [GULProxy proxyWithDelegate:object];
+  GULObjectSwizzler *swizzler = [[GULObjectSwizzler alloc] initWithObject:proxyObject];
+
+  XCTAssertNoThrow([swizzler swizzle]);
+
+  XCTAssertNotEqual(object_getClass(proxyObject), [GULProxy class]);
+  XCTAssertTrue([object_getClass(proxyObject) isSubclassOfClass:[GULProxy class]]);
+
+  XCTAssertTrue([proxyObject respondsToSelector:@selector(gul_objectSwizzler)]);
+  XCTAssertNoThrow([proxyObject performSelector:@selector(gul_objectSwizzler)]);
+
+  XCTAssertTrue([proxyObject respondsToSelector:@selector(gul_class)]);
+  XCTAssertNoThrow([proxyObject performSelector:@selector(gul_class)]);
+}
+
+/** Tests overriding a method that already exists on a proxied object works as expected. */
+- (void)testSwizzleProxiedObjectInvokesInjectedMethodWhenOverridingMethod {
+  NSObject *object = [[NSObject alloc] init];
+  GULProxy *proxyObject = [GULProxy proxyWithDelegate:object];
+
+  GULObjectSwizzler *swizzler = [[GULObjectSwizzler alloc] initWithObject:proxyObject];
+  [swizzler copySelector:@selector(description)
+               fromClass:[GULObjectSwizzlerTest class]
+         isClassSelector:NO];
+  [swizzler swizzle];
+
+  XCTAssertEqual([proxyObject performSelector:@selector(description)], @"SwizzledDescription");
+}
+
+/** Tests adding a method that doesn't exist on a proxied object works as expected. */
+- (void)testSwizzleProxiedObjectInvokesInjectedMethodWhenAddingMethod {
+  NSObject *object = [[NSObject alloc] init];
+  GULProxy *proxyObject = [GULProxy proxyWithDelegate:object];
+
+  GULObjectSwizzler *swizzler = [[GULObjectSwizzler alloc] initWithObject:proxyObject];
+  [swizzler copySelector:@selector(donorDescription)
+               fromClass:[GULObjectSwizzlerTest class]
+         isClassSelector:NO];
+  [swizzler swizzle];
+
+  XCTAssertEqual([proxyObject performSelector:@selector(donorDescription)],
+                 @"SwizzledDonorDescription");
+}
+
+/** Tests KVOing a proxy object that we've ISA Swizzled works as expected. */
+- (void)testRespondsToSelectorWorksEvenIfSwizzledProxyIsKVOd {
+  NSObject *object = [[NSObject alloc] init];
+  GULProxy *proxyObject = [GULProxy proxyWithDelegate:object];
+
+  GULObjectSwizzler *swizzler = [[GULObjectSwizzler alloc] initWithObject:proxyObject];
+  [swizzler copySelector:@selector(donorDescription)
+               fromClass:[GULObjectSwizzlerTest class]
+         isClassSelector:NO];
+  [swizzler swizzle];
+
+  [(NSObject *)proxyObject addObserver:self
+                            forKeyPath:NSStringFromSelector(@selector(description))
+                               options:0
+                               context:NULL];
+
+  XCTAssertTrue([proxyObject respondsToSelector:@selector(donorDescription)]);
+  XCTAssertEqual([proxyObject performSelector:@selector(donorDescription)],
+                 @"SwizzledDonorDescription");
+
+  [(NSObject *)proxyObject removeObserver:self
+                               forKeyPath:NSStringFromSelector(@selector(description))];
+}
+
+/** Tests that -[NSObjectProtocol resopondsToSelector:] works as expected after someone else ISA
+ *  swizzles a proxy object that we've also ISA Swizzled.
+ */
+- (void)testRespondsToSelectorWorksEvenIfSwizzledProxyISASwizzledBySomeoneElse {
+  NSObject *object = [[NSObject alloc] init];
+  GULProxy *proxyObject = [GULProxy proxyWithDelegate:object];
+
+  GULObjectSwizzler *swizzler = [[GULObjectSwizzler alloc] initWithObject:proxyObject];
+  [swizzler copySelector:@selector(donorDescription)
+               fromClass:[GULObjectSwizzlerTest class]
+         isClassSelector:NO];
+  [swizzler swizzle];
+
+  // Someone else ISA Swizzles the same object after GULObjectSwizzler.
+  Class originalClass = object_getClass(proxyObject);
+  NSString *newClassName =
+      [NSString stringWithFormat:@"gul_test_%p_%@", proxyObject, NSStringFromClass(originalClass)];
+  Class generatedClass = objc_allocateClassPair(originalClass, newClassName.UTF8String, 0);
+  objc_registerClassPair(generatedClass);
+  object_setClass(proxyObject, generatedClass);
+
+  XCTAssertTrue([proxyObject respondsToSelector:@selector(donorDescription)]);
+  XCTAssertEqual([proxyObject performSelector:@selector(donorDescription)],
+                 @"SwizzledDonorDescription");
+
+  // Clean up.
+  object_setClass(proxyObject, originalClass);
+  objc_disposeClassPair(generatedClass);
+}
+
 @end
diff --git a/GoogleUtilities/ISASwizzler/GULObjectSwizzler.m b/GoogleUtilities/ISASwizzler/GULObjectSwizzler.m
@@ -80,7 +80,7 @@ - (instancetype)initWithObject:(id)object {
     __strong id swizzledObject = object;
     if (swizzledObject) {
       _swizzledObject = swizzledObject;
-      _originalClass = [swizzledObject class];
+      _originalClass = object_getClass(object);
       NSString *newClassName = [NSString
           stringWithFormat:@"fir_%p_%@", swizzledObject, NSStringFromClass(_originalClass)];
       _generatedClass = objc_allocateClassPair(_originalClass, newClassName.UTF8String, 0);
@@ -134,7 +134,7 @@ - (void)swizzle {
 
     [GULSwizzledObject copyDonorSelectorsUsingObjectSwizzler:self];
 
-    NSAssert(_originalClass == [_swizzledObject class],
+    NSAssert(_originalClass == object_getClass(swizzledObject),
              @"The original class is not the reported class now.");
     NSAssert(class_getInstanceSize(_originalClass) == class_getInstanceSize(_generatedClass),
              @"The instance size of the generated class must be equal to the original class.");
@@ -154,4 +154,8 @@ - (void)dealloc {
   objc_disposeClassPair(_generatedClass);
 }
 
+- (BOOL)isSwizzlingProxyObject {
+  return [_swizzledObject isProxy];
+}
+
 @end
diff --git a/GoogleUtilities/ISASwizzler/GULSwizzledObject.m b/GoogleUtilities/ISASwizzler/GULSwizzledObject.m
@@ -12,8 +12,10 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#import "Private/GULSwizzledObject.h"
+#import <objc/runtime.h>
+
 #import "Private/GULObjectSwizzler.h"
+#import "Private/GULSwizzledObject.h"
 
 NSString *kSwizzlerAssociatedObjectKey = @"gul_objectSwizzler";
 
@@ -26,6 +28,16 @@ @implementation GULSwizzledObject
 + (void)copyDonorSelectorsUsingObjectSwizzler:(GULObjectSwizzler *)objectSwizzler {
   [objectSwizzler copySelector:@selector(gul_objectSwizzler) fromClass:self isClassSelector:NO];
   [objectSwizzler copySelector:@selector(gul_class) fromClass:self isClassSelector:NO];
+
+  // This is needed because NSProxy objects usually override -[NSObjectProtocol respondsToSelector:]
+  // and ask this question to the underlying object. Since we don't swizzle the underlying object
+  // but swizzle the proxy, when someone calls -[NSObjectProtocol respondsToSelector:] on the proxy,
+  // the answer ends up being NO even if we added new methods to the subclass through ISA Swizzling.
+  // To solve that, we override -[NSObjectProtocol respondsToSelector:] in such a way that takes
+  // into account the fact that we've added new methods.
+  if ([objectSwizzler isSwizzlingProxyObject]) {
+    [objectSwizzler copySelector:@selector(respondsToSelector:) fromClass:self isClassSelector:NO];
+  }
 }
 
 - (instancetype)init {
@@ -43,4 +55,10 @@ - (Class)gul_class {
   return [[self gul_objectSwizzler] generatedClass];
 }
 
+// Only added to a class when we detect it is a proxy.
+- (BOOL)respondsToSelector:(SEL)aSelector {
+  Class gulClass = [[self gul_objectSwizzler] generatedClass];
+  return [gulClass instancesRespondToSelector:aSelector] || [super respondsToSelector:aSelector];
+}
+
 @end
diff --git a/GoogleUtilities/ISASwizzler/Private/GULObjectSwizzler.h b/GoogleUtilities/ISASwizzler/Private/GULObjectSwizzler.h
@@ -115,6 +115,9 @@ typedef OBJC_ENUM(uintptr_t, GUL_ASSOCIATION){
  *  the class pair. */
 - (void)swizzle;
 
+/** @return The value of -[objectBeingSwizzled isProxy] */
+- (BOOL)isSwizzlingProxyObject;
+
 @end
 
 NS_ASSUME_NONNULL_END
diff --git a/GoogleUtilities/SwizzlerTestHelpers/GULProxy.h b/GoogleUtilities/SwizzlerTestHelpers/GULProxy.h
@@ -0,0 +1,24 @@
+/*
+ * Copyright 2018 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#import <Foundation/Foundation.h>
+
+/** An example NSProxy that could be used to wrap an object that we have to ISA Swizzle. */
+@interface GULProxy : NSProxy
+
++ (instancetype)proxyWithDelegate:(id)delegate;
+
+@end
diff --git a/GoogleUtilities/SwizzlerTestHelpers/GULProxy.m b/GoogleUtilities/SwizzlerTestHelpers/GULProxy.m
@@ -0,0 +1,95 @@
+/*
+ * Copyright 2018 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#import "GULProxy.h"
+
+@interface GULProxy ()
+
+@property(nonatomic, strong) id delegateObject;
+
+@end
+
+@implementation GULProxy
+
+- (instancetype)initWithDelegate:(id)delegate {
+  _delegateObject = delegate;
+  return self;
+}
+
++ (instancetype)proxyWithDelegate:(id)delegate {
+  return [[GULProxy alloc] initWithDelegate:delegate];
+}
+
+- (id)forwardingTargetForSelector:(SEL)selector {
+  return _delegateObject;
+}
+
+- (void)forwardInvocation:(NSInvocation *)invocation {
+  if (_delegateObject != nil) {
+    [invocation setTarget:_delegateObject];
+    [invocation invoke];
+  }
+}
+
+- (NSMethodSignature *)methodSignatureForSelector:(SEL)selector {
+  return [_delegateObject instanceMethodSignatureForSelector:selector];
+}
+
+- (BOOL)respondsToSelector:(SEL)aSelector {
+  return [_delegateObject respondsToSelector:aSelector];
+}
+
+- (BOOL)isEqual:(id)object {
+  return [_delegateObject isEqual:object];
+}
+
+- (NSUInteger)hash {
+  return [_delegateObject hash];
+}
+
+- (Class)superclass {
+  return [_delegateObject superclass];
+}
+
+- (Class) class {
+  return [_delegateObject class];
+}
+
+    - (BOOL)isKindOfClass : (Class)aClass {
+  return [_delegateObject isKindOfClass:aClass];
+}
+
+- (BOOL)isMemberOfClass:(Class)aClass {
+  return [_delegateObject isMemberOfClass:aClass];
+}
+
+- (BOOL)conformsToProtocol:(Protocol *)aProtocol {
+  return [_delegateObject conformsToProtocol:aProtocol];
+}
+
+- (BOOL)isProxy {
+  return YES;
+}
+
+- (NSString *)description {
+  return [_delegateObject description];
+}
+
+- (NSString *)debugDescription {
+  return [_delegateObject debugDescription];
+}
+
+@end
