OCMockObject is not thread safe

Make OCMockObject thread safe.

Synchronize on the four mutable array instance variables during access to ensure their consistency.
Retain everything retrieved from the arrays if the retrieved objects ever get used outside of the synchronized chunks.

Retain captured invocation arguments since the arguments could otherwise be deallocated at any time. The invocation can't outright retain its arguments since that would cause a retain cycle, so use a customized method to retain arguments that filters out self.
Add all the reference counting methods (-retainCount, -retain, -release, -autorelease) to OCPartialMockObject's selector black list. They're a recipe for stack overflow since partial mocks set up a forwarder for -retain. The forwarder would capture the call to -retain as an NSInvocation, OCMockObject now retains the invocation arguments, which calls -retain on the original object again, which hits the forwarder method again, and so on and so forth.

Make OCMMacroState thread safe as well.

Having a global macro state doesn't work. With locking it doesn't crash, but if two threads are recording at the same time then one of the macro states gets stepped on. Change the global state to be a per-thread state.
Avoid a use-after-dealloc race by not trying to clear the global state in -dealloc and making the macros exception safe such that the +endXXXMacro calls always happen even if the captured invocation throws an exception.

Make OCObserverMockObject thread safe.

Synchronize on the two mutable array instance variables during access to ensure their consistency.

#171

Source/OCMock.xcodeproj/project.pbxproj

@@ -181,6 +181,8 @@
 		2FA28FEAEF9333D2C214DF53 /* NSValue+OCMAdditions.m in Sources */ = {isa = PBXBuildFile; fileRef = 2FA28896E5EEFD7C2F12C2F8 /* NSValue+OCMAdditions.m */; };
 		3C0FF06A1BAA3FD10021AD20 /* OCMFunctionsPrivate.h in Headers */ = {isa = PBXBuildFile; fileRef = 03F370CA1BAA1DE800CAD3E8 /* OCMFunctionsPrivate.h */; };
 		3C0FF06B1BAA3FD20021AD20 /* OCMFunctionsPrivate.h in Headers */ = {isa = PBXBuildFile; fileRef = 03F370CA1BAA1DE800CAD3E8 /* OCMFunctionsPrivate.h */; };
+		3C76716C1BB3EBC500FDC9F4 /* TestClassWithCustomReferenceCounting.m in Sources */ = {isa = PBXBuildFile; fileRef = 3CFBDD761BB3DB200050D9C5 /* TestClassWithCustomReferenceCounting.m */; settings = {COMPILER_FLAGS = "-fno-objc-arc"; }; };
+		3CFBDD771BB3DB200050D9C5 /* TestClassWithCustomReferenceCounting.m in Sources */ = {isa = PBXBuildFile; fileRef = 3CFBDD761BB3DB200050D9C5 /* TestClassWithCustomReferenceCounting.m */; settings = {COMPILER_FLAGS = "-fno-objc-arc"; }; };
 		817EB1171BD765130047E85A /* OCMockObject.m in Sources */ = {isa = PBXBuildFile; fileRef = 03B3159E146333BF0052CD09 /* OCMockObject.m */; };
 		817EB1181BD765130047E85A /* OCClassMockObject.m in Sources */ = {isa = PBXBuildFile; fileRef = 03B3158C146333BF0052CD09 /* OCClassMockObject.m */; };
 		817EB1191BD765130047E85A /* OCPartialMockObject.m in Sources */ = {isa = PBXBuildFile; fileRef = 03B315AA146333BF0052CD09 /* OCPartialMockObject.m */; };
@@ -444,6 +446,8 @@
 		2FA28EC49F6C59B940AE6D00 /* OCMFunctions.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = OCMFunctions.h; sourceTree = "<group>"; };
 		2FA28EDBF243639C57F88A1B /* OCMArgTests.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = OCMArgTests.m; sourceTree = "<group>"; };
 		2FA28EE3142412BD601026EF /* OCMockObjectDynamicPropertyMockingTests.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = OCMockObjectDynamicPropertyMockingTests.m; sourceTree = "<group>"; };
+		3CFBDD751BB3DB200050D9C5 /* TestClassWithCustomReferenceCounting.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = TestClassWithCustomReferenceCounting.h; sourceTree = "<group>"; };
+		3CFBDD761BB3DB200050D9C5 /* TestClassWithCustomReferenceCounting.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = TestClassWithCustomReferenceCounting.m; sourceTree = "<group>"; };
 		817EB1621BD765130047E85A /* OCMock.framework */ = {isa = PBXFileReference; explicitFileType = wrapper.framework; includeInIndex = 0; path = OCMock.framework; sourceTree = BUILT_PRODUCTS_DIR; };
 		D31108AD1828DB8700737925 /* OCMockLibTests.xctest */ = {isa = PBXFileReference; explicitFileType = wrapper.cfbundle; includeInIndex = 0; path = OCMockLibTests.xctest; sourceTree = BUILT_PRODUCTS_DIR; };
 		D31108B71828DB8700737925 /* OCMockLibTests-Info.plist */ = {isa = PBXFileReference; lastKnownFileType = text.plist.xml; path = "OCMockLibTests-Info.plist"; sourceTree = "<group>"; };
@@ -612,6 +616,8 @@
 				03B316291463350E0052CD09 /* OCObserverMockObjectTests.m */,
 				03B3161F1463350E0052CD09 /* NSInvocationOCMAdditionsTests.m */,
 				2FA28DEDB9163597B7C49F3D /* NSMethodSignatureOCMAdditionsTests.m */,
+				3CFBDD751BB3DB200050D9C5 /* TestClassWithCustomReferenceCounting.h */,
+				3CFBDD761BB3DB200050D9C5 /* TestClassWithCustomReferenceCounting.m */,
 				03565A3418F0566F003AE91E /* Supporting Files */,
 			);
 			path = OCMockTests;
@@ -1244,6 +1250,7 @@
 				03565A4418F05721003AE91E /* OCMockObjectProtocolMocksTests.m in Sources */,
 				03565A4B18F05721003AE91E /* NSInvocationOCMAdditionsTests.m in Sources */,
 				03565A4618F05721003AE91E /* OCMockObjectHamcrestTests.mm in Sources */,
+				3CFBDD771BB3DB200050D9C5 /* TestClassWithCustomReferenceCounting.m in Sources */,
 				03E98D5018F310EE00522D42 /* OCMockObjectMacroTests.m in Sources */,
 				03565A4A18F05721003AE91E /* OCObserverMockObjectTests.m in Sources */,
 				03565A4318F05721003AE91E /* OCMockObjectClassMethodMockingTests.m in Sources */,
@@ -1305,6 +1312,7 @@
 			isa = PBXSourcesBuildPhase;
 			buildActionMask = 2147483647;
 			files = (
+				3C76716C1BB3EBC500FDC9F4 /* TestClassWithCustomReferenceCounting.m in Sources */,
 				031E50591BB4A56300E257C3 /* OCMBoxedReturnValueProviderTests.m in Sources */,
 				03C9CA1E18F05A84006DF94D /* OCMArgTests.m in Sources */,
 				0322DA66191188D100CACAF1 /* OCMockObjectVerifyAfterRunTests.m in Sources */,

Source/OCMock/NSInvocation+OCMAdditions.h

@@ -20,7 +20,7 @@
 
 + (NSInvocation *)invocationForBlock:(id)block withArguments:(NSArray *)arguments;
 
-- (BOOL)hasCharPointerArgument;
+- (void)retainObjectArgumentsExcluding:(id)objectToExclude;
 
 - (id)getArgumentAtIndexAsObject:(NSInteger)argIndex;
 

Source/OCMock/NSInvocation+OCMAdditions.m

@@ -14,7 +14,9 @@
  *  under the License.
  */
 
+#import <objc/runtime.h>
 #import "NSInvocation+OCMAdditions.h"
+#import "OCMFunctions.h"
 #import "OCMFunctionsPrivate.h"
 #import "NSMethodSignature+OCMAdditions.h"
 
@@ -41,16 +43,74 @@ + (NSInvocation *)invocationForBlock:(id)block withArguments:(NSArray *)argument
 }
 
 
-- (BOOL)hasCharPointerArgument
+static NSString *const OCMRetainedObjectArgumentsKey = @"OCMRetainedObjectArgumentsKey";
+
+- (void)retainObjectArgumentsExcluding:(id)objectToExclude
 {
-    NSMethodSignature *signature = [self methodSignature];
-    for(NSUInteger i = 0; i < [signature numberOfArguments]; i++)
+    if(objc_getAssociatedObject(self, OCMRetainedObjectArgumentsKey) == nil)
     {
-        const char *argType = OCMTypeWithoutQualifiers([signature getArgumentTypeAtIndex:i]);
-        if(strcmp(argType, "*") == 0)
-            return YES;
+        NSMutableArray *retainedArguments = [[NSMutableArray alloc] init];
+        NSMethodSignature *methodSignature = self.methodSignature;
+
+        id target = self.target;
+        if((target != nil) && (target != objectToExclude) && !object_isClass(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
+            // is a block though (the argument type at index 0 is always "@" even if
+            // the target is a Class or block), and in practice it's OK since you
+            // can't mock a block.
+            [retainedArguments addObject:target];
+        }
+
+        NSUInteger numberOfArguments = methodSignature.numberOfArguments;
+        for(NSUInteger index = 2; index < numberOfArguments; index++)
+        {
+            const char *argumentType = [methodSignature getArgumentTypeAtIndex:index];
+            if(OCMIsObjectType(argumentType) && !OCMIsClassType(argumentType))
+            {
+                id argument;
+                [self getArgument:&argument atIndex:index];
+                if((argument != nil) && (argument != objectToExclude))
+                {
+                    if(OCMIsBlockType(argumentType))
+                    {
+                        // block types need to be copied in case they're stack blocks
+                        id blockArgument = [argument copy];
+                        [retainedArguments addObject:blockArgument];
+                        [blockArgument release];
+                    }
+                    else
+                    {
+                        [retainedArguments addObject:argument];
+                    }
+                }
+            }
+        }
+
+        const char *returnType = methodSignature.methodReturnType;
+        if(OCMIsObjectType(returnType) && !OCMIsClassType(returnType))
+        {
+            id returnValue;
+            [self getReturnValue:&returnValue];
+            if((returnValue != nil) && (returnValue != objectToExclude))
+            {
+                if(OCMIsBlockType(returnType))
+                {
+                    id blockReturnValue = [returnValue copy];
+                    [retainedArguments addObject:blockReturnValue];
+                    [blockReturnValue release];
+                }
+                else
+                {
+                    [retainedArguments addObject:returnValue];
+                }
+            }
+        }
+
+        objc_setAssociatedObject(self, OCMRetainedObjectArgumentsKey, retainedArguments, OBJC_ASSOCIATION_RETAIN);
+        [retainedArguments release];
     }
-    return NO;
 }
 
 

Source/OCMock/OCMFunctions.m

@@ -35,15 +35,36 @@ - (void)failWithException:(NSException *)exception;
 
 #pragma mark  Functions related to ObjC type system
 
+static BOOL OCMIsUnQualifiedClassType(const char *unqualifiedObjCType)
+{
+    return (strcmp(unqualifiedObjCType, @encode(Class)) == 0);
+}
+
+
+static BOOL OCMIsUnQualifiedBlockType(const char *unqualifiedObjCType)
+{
+    char blockType[] = @encode(void(^)());
+    if(strcmp(unqualifiedObjCType, blockType) == 0)
+        return YES;
+
+    // sometimes block argument/return types are tacked onto the type, in angle brackets
+    if(strncmp(unqualifiedObjCType, blockType, sizeof(blockType) - 1) == 0 && unqualifiedObjCType[sizeof(blockType) - 1] == '<')
+        return YES;
+
+    return NO;
+}
+
+
 BOOL OCMIsObjectType(const char *objCType)
 {
     objCType = OCMTypeWithoutQualifiers(objCType);
 
-    if(strcmp(objCType, @encode(id)) == 0 || strcmp(objCType, @encode(Class)) == 0)
+    char objectType[] = @encode(id);
+    if(strcmp(objCType, objectType) == 0 || OCMIsUnQualifiedClassType(objCType))
         return YES;
 
     // sometimes the name of an object's class is tacked onto the type, in double quotes
-    if(strncmp(objCType, @encode(id), 1) == 0 && objCType[1] == '\"')
+    if(strncmp(objCType, objectType, sizeof(objectType) - 1) == 0 && objCType[sizeof(objectType) - 1] == '"')
         return YES;
 
     // if the returnType is a typedef to an object, it has the form ^{OriginClass=#}
@@ -54,11 +75,19 @@ BOOL OCMIsObjectType(const char *objCType)
         return YES;
 
     // if the return type is a block we treat it like an object
-    // TODO: if the runtime were to encode the block's argument and/or return types, this test would not be sufficient
-    if(strcmp(objCType, @encode(void(^)())) == 0)
-        return YES;
+    return OCMIsUnQualifiedBlockType(objCType);
+}
 
-    return NO;
+
+BOOL OCMIsClassType(const char *objCType)
+{
+    return OCMIsUnQualifiedClassType(OCMTypeWithoutQualifiers(objCType));
+}
+
+
+BOOL OCMIsBlockType(const char *objCType)
+{
+    return OCMIsUnQualifiedBlockType(OCMTypeWithoutQualifiers(objCType));
 }
 
 

Source/OCMock/OCMFunctionsPrivate.h

@@ -28,7 +28,8 @@
 @class OCPartialMockObject;
 
 
-OCMOCK_EXTERN BOOL OCMIsObjectType(const char *objCType);
+OCMOCK_EXTERN BOOL OCMIsClassType(const char *objCType);
+OCMOCK_EXTERN BOOL OCMIsBlockType(const char *objCType);
 const char *OCMTypeWithoutQualifiers(const char *objCType);
 BOOL OCMEqualTypesAllowingOpaqueStructs(const char *type1, const char *type2);
 CFNumberType OCMNumberTypeForObjCType(const char *objcType);

Source/OCMock/OCMInvocationMatcher.m

@@ -21,6 +21,7 @@
 #import "NSInvocation+OCMAdditions.h"
 #import "OCMInvocationMatcher.h"
 #import "OCClassMockObject.h"
+#import "OCMFunctions.h"
 #import "OCMFunctionsPrivate.h"
 #import "OCMBlockArgCaller.h"
 
@@ -41,11 +42,11 @@ - (void)dealloc
 - (void)setInvocation:(NSInvocation *)anInvocation
 {
     [recordedInvocation release];
-    // When the method has a char* argument we do not retain the arguments. This makes it possible
-    // to match char* args literally and with anyPointer. Not retaining the argument means that
-    // in these cases tests that use their own autorelease pools may fail unexpectedly.
-    if(![anInvocation hasCharPointerArgument])
-        [anInvocation retainArguments];
+    // Don't do a regular -retainArguments on the invocation that we use for matching. NSInvocation
+    // effectively does an strcpy on char* arguments which messes up matching them literally and blows
+    // up with anyPointer (in strlen since it's not actually a C string). Also on the off-chance that
+    // anInvocation contains self as an argument, -retainArguments would create a retain cycle.
+    [anInvocation retainObjectArgumentsExcluding:self];
     recordedInvocation = [anInvocation retain];
 }
 

Source/OCMock/OCMMacroState.m

@@ -24,28 +24,34 @@
 
 @implementation OCMMacroState
 
-static OCMMacroState *globalState;
+static NSString *const OCMGlobalStateKey = @"OCMGlobalStateKey";
 
 #pragma mark  Methods to begin/end macros
 
 + (void)beginStubMacro
 {
     OCMStubRecorder *recorder = [[[OCMStubRecorder alloc] init] autorelease];
-    globalState = [[[OCMMacroState alloc] initWithRecorder:recorder] autorelease];
+    OCMMacroState *macroState = [[OCMMacroState alloc] initWithRecorder:recorder];
+    [NSThread currentThread].threadDictionary[OCMGlobalStateKey] = macroState;
+    [macroState release];
 }
 
 + (OCMStubRecorder *)endStubMacro
 {
-    OCMStubRecorder *recorder = (OCMStubRecorder *)[globalState recorder];
-    globalState = nil;
-    return recorder;
+    NSMutableDictionary *threadDictionary = [NSThread currentThread].threadDictionary;
+    OCMMacroState *globalState = threadDictionary[OCMGlobalStateKey];
+    OCMStubRecorder *recorder = [(OCMStubRecorder *)[globalState recorder] retain];
+    [threadDictionary removeObjectForKey:OCMGlobalStateKey];
+    return [recorder autorelease];
 }
 
 
 + (void)beginExpectMacro
 {
     OCMExpectationRecorder *recorder = [[[OCMExpectationRecorder alloc] init] autorelease];
-    globalState = [[[OCMMacroState alloc] initWithRecorder:recorder] autorelease];
+    OCMMacroState *macroState = [[OCMMacroState alloc] initWithRecorder:recorder];
+    [NSThread currentThread].threadDictionary[OCMGlobalStateKey] = macroState;
+    [macroState release];
 }
 
 + (OCMStubRecorder *)endExpectMacro
@@ -58,20 +64,22 @@ + (void)beginVerifyMacroAtLocation:(OCMLocation *)aLocation
 {
     OCMVerifier *recorder = [[[OCMVerifier alloc] init] autorelease];
     [recorder setLocation:aLocation];
-    globalState = [[[OCMMacroState alloc] initWithRecorder:recorder] autorelease];
+    OCMMacroState *macroState = [[OCMMacroState alloc] initWithRecorder:recorder];
+    [NSThread currentThread].threadDictionary[OCMGlobalStateKey] = macroState;
+    [macroState release];
 }
 
 + (void)endVerifyMacro
 {
-    globalState = nil;
+    [[NSThread currentThread].threadDictionary removeObjectForKey:OCMGlobalStateKey];
 }
 
 
 #pragma mark  Accessing global state
 
 + (OCMMacroState *)globalState
 {
-    return globalState;
+    return [NSThread currentThread].threadDictionary[OCMGlobalStateKey];
 }
 
 
@@ -90,8 +98,7 @@ - (id)initWithRecorder:(OCMRecorder *)aRecorder
 - (void)dealloc
 {
     [recorder release];
-    if(globalState == self)
-        globalState = nil;
+    NSAssert([NSThread currentThread].threadDictionary[OCMGlobalStateKey] != self, @"Unexpected dealloc while set as the global state");
     [super dealloc];
 }
 

Source/OCMock/OCMock.h

@@ -42,17 +42,27 @@
 ({ \
     _OCMSilenceWarnings( \
         [OCMMacroState beginStubMacro]; \
-        invocation; \
-        [OCMMacroState endStubMacro]; \
+        OCMStubRecorder *recorder = nil; \
+        @try{ \
+            invocation; \
+        }@finally{ \
+            recorder = [OCMMacroState endStubMacro]; \
+        } \
+        recorder; \
     ); \
 })
 
 #define OCMExpect(invocation) \
 ({ \
     _OCMSilenceWarnings( \
         [OCMMacroState beginExpectMacro]; \
-        invocation; \
-        [OCMMacroState endExpectMacro]; \
+        OCMStubRecorder *recorder = nil; \
+        @try{ \
+            invocation; \
+        }@finally{ \
+            recorder = [OCMMacroState endExpectMacro]; \
+        } \
+        recorder; \
     ); \
 })
 
@@ -71,8 +81,11 @@
 ({ \
     _OCMSilenceWarnings( \
         [OCMMacroState beginVerifyMacroAtLocation:OCMMakeLocation(self, __FILE__, __LINE__)]; \
-        invocation; \
-        [OCMMacroState endVerifyMacro]; \
+        @try{ \
+            invocation; \
+        }@finally{ \
+            [OCMMacroState endVerifyMacro]; \
+        } \
     ); \
 })