Deal with mocks that do direct referencing of instance variables.

Reallocate class and partial mocks based on the size of the object they are mocking.
For class mocks allow direct referencing of instance variables.
For Partial mocks fills the space with 0xEC and will throw an exception if we detect that an instance variable has been written to. The value 0xEB will intentionally likely cause crashes if the memory is read.

Author: dmaclach

Source/OCMock/OCClassMockObject.h

@@ -25,4 +25,6 @@
 
 - (void)assertClassIsSupported:(Class)aClass;
 
+- (void)scribbleOnMemory:(void *)start ofSize:(size_t)size;
+- (void)verifyScribbleAt:(void *)start ofSize:(size_t)size;
 @end

Source/OCMock/OCClassMockObject.m

@@ -30,6 +30,8 @@ @interface OCClassMockObjectInstanceVars : NSObject
 @property (nonatomic) Class mockedClass;
 @property (nonatomic) Class originalMetaClass;
 @property (nonatomic) Class classCreatedForNewMetaClass;
+@property (nonatomic) void *classScribbleStart;
+@property (nonatomic) size_t classScribbleSize;
 @end
 
 @implementation OCClassMockObjectInstanceVars
@@ -39,22 +41,37 @@ @interface OCClassMockObject ()
 @property (nonatomic) Class mockedClass;
 @property (nonatomic) Class originalMetaClass;
 @property (nonatomic) Class classCreatedForNewMetaClass;
+@property (nonatomic) void *classScribbleStart;
+@property (nonatomic) size_t classScribbleSize;
 @end
 
 static const char *OCClassMockObjectInstanceVarsKey = "OCClassMockObjectInstanceVarsKey";
 
 @implementation OCClassMockObject
 
-#pragma mark Initialisers, description, accessors, etc.
+#pragma mark  Initialisers, description, etc.
 
 - (id)initWithClass:(Class)aClass
 {
     [self assertClassIsSupported:aClass];
+
+    size_t allocedSize = class_getInstanceSize(aClass);
+    Class selfClass = object_getClass(self);
+    size_t selfSize = class_getInstanceSize(selfClass);
+    if(allocedSize > selfSize)
+    {
+        self = realloc(self, allocedSize);
+    }
     self = [super init];
+
     OCClassMockObjectInstanceVars *vars = [[OCClassMockObjectInstanceVars alloc] init];
     objc_setAssociatedObject(self, OCClassMockObjectInstanceVarsKey, vars, OBJC_ASSOCIATION_RETAIN_NONATOMIC);
     [vars release];
 
+    self.classScribbleSize = allocedSize - selfSize;
+    self.classScribbleStart = (void *)self + selfSize;
+    [self scribbleOnMemory:self.classScribbleStart ofSize:self.classScribbleSize];
+
     self.mockedClass = aClass;
     [self prepareClassForClassMethodMocking];
     return self;
@@ -71,9 +88,20 @@ - (NSString *)description
     return [NSString stringWithFormat:@"OCClassMockObject(%@)", NSStringFromClass(self.mockedClass)];
 }
 
+- (void)scribbleOnMemory:(void *)start ofSize:(size_t)size
+{
+  bzero(start, size);
+}
+
+- (void)verifyScribbleAt:(void *)start ofSize:(size_t)size
+{
+  // Default version does no verification
+}
+
 #pragma mark  Setters/Getters
 
-- (OCClassMockObjectInstanceVars *)classMockObjectInstanceVars {
+- (OCClassMockObjectInstanceVars *)classMockObjectInstanceVars
+{
     return objc_getAssociatedObject(self, OCClassMockObjectInstanceVarsKey);
 }
 
@@ -92,6 +120,16 @@ - (Class)originalMetaClass
     return self.classMockObjectInstanceVars.originalMetaClass;
 }
 
+- (void *)classScribbleStart
+{
+  return self.classMockObjectInstanceVars.classScribbleStart;
+}
+
+- (size_t)classScribbleSize
+{
+  return self.classMockObjectInstanceVars.classScribbleSize;
+}
+
 - (void)setMockedClass:(Class)mockedClass
 {
     self.classMockObjectInstanceVars.mockedClass = mockedClass;
@@ -120,6 +158,16 @@ - (void)assertClassIsSupported:(Class)aClass
     }
 }
 
+- (void)setClassScribbleSize:(size_t)classScribbleSize
+{
+    self.classMockObjectInstanceVars.classScribbleSize = classScribbleSize;
+}
+
+- (void)setClassScribbleStart:(void *)classScribbleStart
+{
+    self.classMockObjectInstanceVars.classScribbleStart = classScribbleStart;
+}
+
 #pragma mark Extending/overriding superclass behaviour
 
 - (void)stopMocking
@@ -134,6 +182,7 @@ - (void)stopMocking
         self.classCreatedForNewMetaClass = nil;
     }
     [super stopMocking];
+  	[self verifyScribbleAt:self.classScribbleStart ofSize:self.classScribbleSize];
 }
 
 

Source/OCMock/OCPartialMockObject.m

@@ -32,6 +32,11 @@ @implementation OCPartialMockObjectInstanceVars
 
 static const char *OCPartialMockObjectInstanceVarsKey = "OCPartialMockObjectInstanceVarsKey";
 
+// 0xEB chosen intentionally to try and force crashes.
+// It has both the high and low bit set, and 0xEBEBEBEBEB..etc
+// should be recognizable in a debugger as a bad value.
+static uint8_t OCScribbleByte = 0xEB;
+
 @interface OCPartialMockObject ()
 @property (nonatomic) NSObject *realObject;
 @property (nonatomic) NSInvocation *invocationFromMock;
@@ -61,9 +66,29 @@ - (NSString *)description
     return [NSString stringWithFormat:@"OCPartialMockObject(%@)", NSStringFromClass(self.mockedClass)];
 }
 
+- (void)scribbleOnMemory:(void *)start ofSize:(size_t)size;
+{
+  for(size_t i = 0; i < size; ++i)
+  {
+    ((uint8_t*)start)[i] = OCScribbleByte;
+  }
+}
+
+- (void)verifyScribbleAt:(void *)start ofSize:(size_t)size;
+{
+  for(size_t i = 0; i < size; ++i)
+  {
+    if(((uint8_t*)start)[i] != OCScribbleByte)
+    {
+      [NSException raise:NSInternalInconsistencyException format:@"The class that partial mock `%@` does internal direct ivar accesses. You must use the real object instead of the mock for all uses other than setting/verifying stubs/expectations etc.", self];
+    }
+  }
+}
+
 #pragma mark  Setters/Getters
 
-- (OCPartialMockObjectInstanceVars *)partialMockObjectInstanceVars {
+- (OCPartialMockObjectInstanceVars *)partialMockObjectInstanceVars
+{
     return objc_getAssociatedObject(self, OCPartialMockObjectInstanceVarsKey);
 }
 

Source/OCMockTests/OCMockObjectTests.m

@@ -216,6 +216,23 @@ + (BOOL)supportsMocking:(NSString **)reasonPtr
 
 @end
 
+@interface TestClassLargeClass : NSObject
+{
+  int foo[4096];
+}
+@end
+
+@implementation TestClassLargeClass
+
+- (void)dirtyInstanceVariables:(TestClassLargeClass *)cls
+{
+  for(int i = 0; i < 4096; ++i) {
+    cls->foo[i] = i;
+  }
+}
+
+@end
+
 static NSString *TestNotification = @"TestNotification";
 
 
@@ -1166,4 +1183,26 @@ - (void)testMockObjectsHaveNoInstanceVariables
   XCTAssertEqual(class_getInstanceSize([NSProxy class]), class_getInstanceSize([OCClassMockObject class]));
 }
 
+- (void)testClassMockAllowsDirectMemoryAccess
+{
+  TestClassLargeClass *one = [[TestClassLargeClass alloc] init];
+  id mockOne = OCMClassMock([TestClassLargeClass class]);
+  [one dirtyInstanceVariables:mockOne];
+}
+
+- (void)performDirectMemoryAccess
+{
+  @autoreleasepool {
+    TestClassLargeClass *one = [[TestClassLargeClass alloc] init];
+    TestClassLargeClass *two = [[TestClassLargeClass alloc] init];
+    id mockTwo = OCMPartialMock(two);
+    [one dirtyInstanceVariables:mockTwo];
+  }
+}
+
+- (void)testPartialClassMockDoesNotAllowDirectMemoryAccess
+{
+  XCTAssertThrowsSpecificNamed([self performDirectMemoryAccess], NSException, NSInternalInconsistencyException);
+}
+
 @end