Avoid retaining objects that are being deallocated.

We do not want to retain objects that are being deallocated because this will cause a crash later when the retaining object releases them. An example of this is when a delegate (that is mocked) is called in a dealloc for an object. The mock retains the deallocating object as part of an invocation and then things crash later when the invocation is released as part of the mock deallocating itself.
erikdoe · Jul 14, 2020 · 711969f · 711969f
1 parent b8746e5
commit 711969f
Show file tree

Hide file tree

Showing 5 changed files with 53 additions and 8 deletions.
diff --git a/Source/OCMock/NSInvocation+OCMAdditions.m b/Source/OCMock/NSInvocation+OCMAdditions.m
@@ -66,7 +66,7 @@ - (void)retainObjectArgumentsExcludingObject:(id)objectToExclude
     NSMutableArray *retainedArguments = [[NSMutableArray alloc] init];
 
     id target = [self target];
-    if((target != nil) && (target != objectToExclude) && !object_isClass(target))
+    if((target != nil) && (target != objectToExclude) && !object_isClass(target) && !OCMIsDeallocating(target))
     {
         // Bad things will happen if the target is a block since it's not being
         // copied. There isn't a very good way to tell if an invocation's target
@@ -84,7 +84,7 @@ - (void)retainObjectArgumentsExcludingObject:(id)objectToExclude
         {
             id argument;
             [self getArgument:&argument atIndex:index];
-            if((argument != nil) && (argument != objectToExclude))
+            if((argument != nil) && (argument != objectToExclude) && !OCMIsDeallocating(argument))
             {
                 if(OCMIsBlockType(argumentType))
                 {
@@ -120,7 +120,7 @@ - (void)retainObjectArgumentsExcludingObject:(id)objectToExclude
     {
         id returnValue;
         [self getReturnValue:&returnValue];
-        if((returnValue != nil) && (returnValue != objectToExclude))
+        if((returnValue != nil) && (returnValue != objectToExclude) && !OCMIsDeallocating(returnValue))
         {
             if(OCMIsBlockType(returnType))
             {

diff --git a/Source/OCMock/OCMFunctions.m b/Source/OCMock/OCMFunctions.m
@@ -32,6 +32,14 @@ - (void)recordFailureWithDescription:(NSString *)description inFile:(NSString *)
 - (void)failWithException:(NSException *)exception;
 @end
 
+@interface NSObject(OCMKnownNSObjectMethods)
+- (BOOL)_isDeallocating;
+@end
+
+// From objc/runtime/objc-internal.h
+// Only available on macOS 10.11/iOS 9.
+extern id objc_initWeakOrNil(id *location, id newObj) __attribute__((weak_import));
+extern void objc_destroyWeak(id _Nullable * _Nonnull location) __attribute__((weak_import));
 
 #pragma mark  Functions related to ObjC type system
 
@@ -450,3 +458,9 @@ void OCMReportFailure(OCMLocation *loc, NSString *description)
     }
 
 }
+
+BOOL OCMIsDeallocating(id anObject)
+{
+    return [anObject _isDeallocating];
+}
+
diff --git a/Source/OCMock/OCMFunctionsPrivate.h b/Source/OCMock/OCMFunctionsPrivate.h
@@ -46,6 +46,8 @@ OCPartialMockObject *OCMGetAssociatedMockForObject(id anObject);
 
 void OCMReportFailure(OCMLocation *loc, NSString *description);
 
+BOOL OCMIsDeallocating(id anObject);
+
 BOOL OCMIsNonEscapingBlock(id block);
 
 

diff --git a/Source/OCMock/OCMStubRecorder.m b/Source/OCMock/OCMStubRecorder.m
@@ -24,6 +24,7 @@
 #import "OCMBlockCaller.h"
 #import "OCMRealObjectForwarder.h"
 #import "OCMInvocationStub.h"
+#import "OCMFunctionsPrivate.h"
 
 
 @implementation OCMStubRecorder
@@ -51,7 +52,7 @@ - (OCMInvocationStub *)stub
 - (id)andReturn:(id)anObject
 {
     id action;
-    if(anObject == mockObject)
+    if(anObject == mockObject || OCMIsDeallocating(anObject))
     {
         action = [[[OCMNonRetainingObjectReturnValueProvider alloc] initWithValue:anObject] autorelease];
     }

diff --git a/Source/OCMockTests/OCMockObjectRuntimeTests.m b/Source/OCMockTests/OCMockObjectRuntimeTests.m
@@ -54,14 +54,15 @@ - (NSString *)stringForTypedef:(TypedefString *)string
 
 @interface TestDelegate : NSObject
 
-- (void)go;
+- (id)go:(id)sender;
 
 @end
 
 @implementation TestDelegate
 
-- (void)go
+- (id)go:(id)sender
 {
+    return sender;
 }
 
 @end
@@ -77,7 +78,24 @@ @implementation TestClassWithDelegate
 - (void)run
 {
     TestDelegate *delegate = self.delegate;
-    [delegate go];
+    [delegate go:nil];
+}
+
+@end
+
+
+@interface TestClassThatCallsDelegateOnDealloc : NSObject
+
+@property (nonatomic, weak) TestDelegate *delegate;
+
+@end
+
+@implementation TestClassThatCallsDelegateOnDealloc
+
+- (void)dealloc
+{
+    TestDelegate *delegate = self.delegate;
+    [delegate go:self];
 }
 
 @end
@@ -313,7 +331,7 @@ - (void)testWeakReferencesShouldStayAround
 
     [object run];
 
-    OCMVerify([mockDelegate go]);
+    OCMVerify([mockDelegate go:nil]);
     XCTAssertNotNil(object.delegate, @"Should still have delegate");
 }
 
@@ -329,6 +347,16 @@ - (void)testDynamicSubclassesShouldBeDisposed
     XCTAssertEqual(numClassesBefore, numClassesAfter, @"Should have disposed dynamically generated classes.");
 }
 
+- (void)testHandlesCallingMockWithSelfAsArgumentInDealloc
+{
+    // Note that this test will crash on failure.
+    id mock = [OCMockObject mockForClass:[TestDelegate class]];
+    [[mock expect] go:OCMOCK_ANY];
+    TestClassThatCallsDelegateOnDealloc *foo = [[TestClassThatCallsDelegateOnDealloc alloc] init];
+    foo.delegate = mock;
+    foo = nil;
+    [mock verify];
+}
 
 - (void)testClassesWithResolveMethodsCanBeMocked
 {