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

@@ -23,4 +23,6 @@
 - (Class)mockedClass;
 - (Class)mockObjectClass;  // since -class returns the mockedClass
 
+- (void)scribbleOnMemory:(void *)start ofSize:(size_t)size;
+- (void)verifyScribbleAt:(void *)start ofSize:(size_t)size;
 @end

Source/OCMock/OCClassMockObject.m

@@ -25,6 +25,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
@@ -34,24 +36,38 @@ @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
 {
 	if(aClass == Nil)
 		[NSException raise:NSInvalidArgumentException format:@"Class cannot be Nil."];
 
+  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;
@@ -68,9 +84,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);
 }
 
@@ -89,6 +116,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;
@@ -104,6 +141,15 @@ - (void)setOriginalMetaClass:(Class)originalMetaClass
   self.classMockObjectInstanceVars.originalMetaClass = originalMetaClass;
 }
 
+- (void)setClassScribbleSize:(size_t)classScribbleSize
+{
+  self.classMockObjectInstanceVars.classScribbleSize = classScribbleSize;
+}
+
+- (void)setClassScribbleStart:(void *)classScribbleStart
+{
+  self.classMockObjectInstanceVars.classScribbleStart = classScribbleStart;
+}
 
 #pragma mark  Extending/overriding superclass behaviour
 
@@ -119,6 +165,7 @@ - (void)stopMocking
         self.classCreatedForNewMetaClass = nil;
     }
     [super stopMocking];
+  	[self verifyScribbleAt:self.classScribbleStart ofSize:self.classScribbleSize];
 }
 
 

Source/OCMock/OCMockObject.m

@@ -106,7 +106,7 @@ + (id)observerMock
 }
 
 
-#pragma mark  Initialisers, description, accessors, etc.
+#pragma mark  Initialisers, description, etc.
 
 - (instancetype)init
 {

Source/OCMock/OCPartialMockObject.m

@@ -33,14 +33,19 @@ @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;
 @end
 
 @implementation OCPartialMockObject
 
-#pragma mark  Initialisers, description, accessors, etc.
+#pragma mark  Initialisers, description, etc.
 
 - (id)initWithObject:(NSObject *)anObject
 {
@@ -63,9 +68,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

@@ -193,6 +193,22 @@ - (NSString *)stringValue;
 
 @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";
 
@@ -1108,6 +1124,28 @@ - (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