main
  1//
  2// Licensed under the terms in License.txt
  3//
  4// Copyright 2010 Allen Ding. All rights reserved.
  5//
  6
  7#import "KWIntercept.h"
  8#import "KWMessagePattern.h"
  9#import "KWMessageSpying.h"
 10#import "KWStub.h"
 11
 12static const char * const KWInterceptClassSuffix = "_KWIntercept";
 13static NSMutableDictionary *KWObjectStubs = nil;
 14static NSMutableDictionary *KWMessageSpies = nil;
 15static NSMutableArray *KWRestoredObjects = nil;
 16
 17#pragma mark -
 18#pragma mark Intercept Enabled Method Implementations
 19
 20Class KWRestoreOriginalClass(id anObject);
 21void KWInterceptedForwardInvocation(id anObject, SEL aSelector, NSInvocation* anInvocation);
 22void KWInterceptedDealloc(id anObject, SEL aSelector);
 23Class KWInterceptedClass(id anObject, SEL aSelector);
 24Class KWInterceptedSuperclass(id anObject, SEL aSelector);
 25
 26#pragma mark -
 27#pragma mark Getting Forwarding Implementations
 28
 29#pragma clang diagnostic push
 30#pragma clang diagnostic ignored "-Wundeclared-selector"
 31
 32IMP KWRegularForwardingImplementation(void) {
 33    return class_getMethodImplementation([NSObject class], @selector(KWNonExistantSelector));
 34}
 35
 36IMP KWStretForwardingImplementation(void) {
 37    return class_getMethodImplementation_stret([NSObject class], @selector(KWNonExistantSelector));
 38}
 39
 40#pragma clang diagnostic pop
 41
 42IMP KWForwardingImplementationForMethodEncoding(const char* encoding) {
 43#if TARGET_CPU_ARM
 44    const NSUInteger stretLengthThreshold = 4;
 45#elif TARGET_CPU_X86
 46    const NSUInteger stretLengthThreshold = 8;
 47#else
 48    // TODO: This just makes an assumption right now. Expand to support all
 49    // official architectures correctly.
 50    const NSUInteger stretLengthThreshold = 8;
 51#endif // #if TARGET_CPU_ARM
 52
 53    NSMethodSignature *signature = [NSMethodSignature signatureWithObjCTypes:encoding];
 54
 55    if (*[signature methodReturnType] == '{' && [signature methodReturnLength] > stretLengthThreshold) {
 56        NSLog(@"Warning: The Objective-C runtime appears to have bugs when forwarding messages with certain struct layouts as return types, so if a crash occurs this could be the culprit");
 57        return KWStretForwardingImplementation();
 58    } else {
 59        return KWRegularForwardingImplementation();
 60    }
 61}
 62
 63#pragma mark -
 64#pragma mark Getting Intercept Class Information
 65
 66BOOL KWObjectIsClass(id anObject) {
 67    return class_isMetaClass(object_getClass(anObject));
 68}
 69
 70BOOL KWClassIsInterceptClass(Class aClass) {
 71    const char *name = class_getName(aClass);
 72    char *result = strstr(name, KWInterceptClassSuffix);
 73    return result != nil;
 74}
 75
 76int interceptCount = 0;
 77
 78NSString *KWInterceptClassNameForClass(Class aClass) {
 79    const char *className = class_getName(aClass);
 80    interceptCount++;
 81    return [NSString stringWithFormat:@"%s%s%d", className, KWInterceptClassSuffix, interceptCount];
 82}
 83
 84Class KWInterceptClassForCanonicalClass(Class canonicalClass) {
 85    NSString *interceptClassName = KWInterceptClassNameForClass(canonicalClass);
 86    Class interceptClass = NSClassFromString(interceptClassName);
 87
 88    if (interceptClass != nil)
 89        return interceptClass;
 90
 91    interceptClass = objc_allocateClassPair(canonicalClass, [interceptClassName UTF8String], 0);
 92    objc_registerClassPair(interceptClass);
 93
 94    class_addMethod(interceptClass, @selector(forwardInvocation:), (IMP)KWInterceptedForwardInvocation, "v@:@");
 95    class_addMethod(interceptClass, @selector(dealloc), (IMP)KWInterceptedDealloc, "v@:");
 96    class_addMethod(interceptClass, @selector(class), (IMP)KWInterceptedClass, "#@:");
 97    class_addMethod(interceptClass, @selector(superclass), (IMP)KWInterceptedSuperclass, "#@:");
 98
 99    Class interceptMetaClass = object_getClass(interceptClass);
100    class_addMethod(interceptMetaClass, @selector(forwardInvocation:), (IMP)KWInterceptedForwardInvocation, "v@:@");
101
102    return interceptClass;
103}
104
105Class KWRealClassForClass(Class aClass) {
106    if (KWClassIsInterceptClass(aClass))
107        return [aClass superclass];
108
109    return aClass;
110}
111
112#pragma mark -
113#pragma mark Enabling Intercepting
114
115static BOOL IsTollFreeBridged(Class class, id obj)
116{
117    // this is a naive check, but good enough for the purposes of failing fast
118    return [NSStringFromClass(class) hasPrefix:@"NSCF"];
119}
120
121// Canonical class is the non-intercept, non-metaclass, class for an object.
122//
123// (e.g. [Animal class] would be canonical, not
124// object_getClass([Animal class]), if the Animal class has not been touched
125// by the intercept mechanism.
126
127Class KWSetupObjectInterceptSupport(id anObject) {
128    Class objectClass = object_getClass(anObject);
129
130    if (IsTollFreeBridged(objectClass, anObject)) {
131        [NSException raise:@"KWTollFreeBridgingInterceptException" format:@"Attempted to stub object of class %@. Kiwi does not support setting expectation or stubbing methods on toll-free bridged objects.", NSStringFromClass(objectClass)];
132    }
133
134    if (KWClassIsInterceptClass(objectClass))
135        return objectClass;
136
137    BOOL objectIsClass = KWObjectIsClass(anObject);
138    Class canonicalClass =  objectIsClass ? anObject : objectClass;
139    Class canonicalInterceptClass = KWInterceptClassForCanonicalClass(canonicalClass);
140    Class interceptClass = objectIsClass ? object_getClass(canonicalInterceptClass) : canonicalInterceptClass;
141
142    object_setClass(anObject, interceptClass);
143
144    return interceptClass;
145}
146
147void KWSetupMethodInterceptSupport(Class interceptClass, SEL aSelector) {
148    BOOL isMetaClass = class_isMetaClass(interceptClass);
149    Method method = isMetaClass ? class_getClassMethod(interceptClass, aSelector)
150                                : class_getInstanceMethod(interceptClass, aSelector);
151
152    if (method == nil) {
153        [NSException raise:NSInvalidArgumentException format:@"cannot setup intercept support for -%@ because no such method exists",
154                                                             NSStringFromSelector(aSelector)];
155    }
156
157    const char *encoding = method_getTypeEncoding(method);
158    IMP forwardingImplementation = KWForwardingImplementationForMethodEncoding(encoding);
159    class_addMethod(interceptClass, aSelector, forwardingImplementation, encoding);
160}
161
162#pragma mark -
163#pragma mark Intercept Enabled Method Implementations
164
165Class KWRestoreOriginalClass(id anObject) {
166    Class interceptClass = object_getClass(anObject);
167    if (KWClassIsInterceptClass(interceptClass))
168    {
169        Class originalClass = class_getSuperclass(interceptClass);
170        // anObject->isa = originalClass;
171        object_setClass(anObject, originalClass);
172    }
173    return interceptClass;
174}
175
176void KWInterceptedForwardInvocation(id anObject, SEL aSelector, NSInvocation* anInvocation) {
177    NSValue *key = [NSValue valueWithNonretainedObject:anObject];
178    NSMutableDictionary *spyArrayDictionary = KWMessageSpies[key];
179
180    for (KWMessagePattern *messagePattern in spyArrayDictionary) {
181        if ([messagePattern matchesInvocation:anInvocation]) {
182            NSArray *spies = spyArrayDictionary[messagePattern];
183
184            for (NSValue *spyWrapper in spies) {
185                id<KWMessageSpying> spy = [spyWrapper nonretainedObjectValue];
186                [spy object:anObject didReceiveInvocation:anInvocation];
187            }
188        }
189    }
190
191    NSMutableArray *stubs = KWObjectStubs[key];
192
193    for (KWStub *stub in stubs) {
194        if ([stub processInvocation:anInvocation])
195            return;
196    }
197
198    Class interceptClass = KWRestoreOriginalClass(anObject);
199    [anInvocation invoke];
200    // anObject->isa = interceptClass;
201    object_setClass(anObject, interceptClass);
202}
203
204void KWInterceptedDealloc(id anObject, SEL aSelector) {
205    NSValue *key = [NSValue valueWithNonretainedObject:anObject];
206    [KWMessageSpies removeObjectForKey:key];
207    [KWObjectStubs removeObjectForKey:key];
208
209    KWRestoreOriginalClass(anObject);
210    [anObject dealloc];
211}
212
213Class KWInterceptedClass(id anObject, SEL aSelector) {
214    Class interceptClass = object_getClass(anObject);
215    Class originalClass = class_getSuperclass(interceptClass);
216    return originalClass;
217}
218
219Class KWInterceptedSuperclass(id anObject, SEL aSelector) {
220    Class interceptClass = object_getClass(anObject);
221    Class originalClass = class_getSuperclass(interceptClass);
222    Class originalSuperclass = class_getSuperclass(originalClass);
223    return originalSuperclass;
224}
225
226#pragma mark - Managing Stubs & Spies
227
228void KWClearStubsAndSpies(void) {
229    KWRestoredObjects = [NSMutableArray array];
230    KWClearAllMessageSpies();
231    KWClearAllObjectStubs();
232    KWRestoredObjects = nil;
233}
234
235#pragma mark -
236#pragma mark Managing Objects Stubs
237
238void KWAssociateObjectStub(id anObject, KWStub *aStub, BOOL overrideExisting) {
239    if (KWObjectStubs == nil)
240        KWObjectStubs = [[NSMutableDictionary alloc] init];
241
242    NSValue *key = [NSValue valueWithNonretainedObject:anObject];
243    NSMutableArray *stubs = KWObjectStubs[key];
244
245    if (stubs == nil) {
246        stubs = [[NSMutableArray alloc] init];
247        KWObjectStubs[key] = stubs;
248        [stubs release];
249    }
250
251    NSUInteger stubCount = [stubs count];
252
253    for (NSUInteger i = 0; i < stubCount; ++i) {
254        KWStub *existingStub = stubs[i];
255
256        if ([aStub.messagePattern isEqualToMessagePattern:existingStub.messagePattern]) {
257            if (overrideExisting) {
258                [stubs removeObjectAtIndex:i];
259                break;
260            } else {
261                return;
262            }
263        }
264    }
265
266    [stubs addObject:aStub];
267}
268
269void KWClearObjectStubs(id anObject) {
270    NSValue *key = [NSValue valueWithNonretainedObject:anObject];
271    [KWObjectStubs removeObjectForKey:key];
272}
273
274void KWClearAllObjectStubs(void) {
275    for (NSValue *objectKey in KWObjectStubs) {
276        id stubbedObject = [objectKey nonretainedObjectValue];
277        if ([KWRestoredObjects containsObject:stubbedObject]) {
278            continue;
279        }
280        KWRestoreOriginalClass(stubbedObject);
281        [KWRestoredObjects addObject:stubbedObject];
282    }
283    [KWObjectStubs removeAllObjects];
284}
285
286#pragma mark -
287#pragma mark Managing Message Spies
288
289void KWAssociateMessageSpy(id anObject, id aSpy, KWMessagePattern *aMessagePattern) {
290    if (KWMessageSpies == nil)
291        KWMessageSpies = [[NSMutableDictionary alloc] init];
292
293    NSValue *key = [NSValue valueWithNonretainedObject:anObject];
294    NSMutableDictionary *spies = KWMessageSpies[key];
295
296    if (spies == nil) {
297        spies = [[NSMutableDictionary alloc] init];
298        KWMessageSpies[key] = spies;
299        [spies release];
300    }
301
302    NSMutableArray *messagePatternSpies = spies[aMessagePattern];
303
304    if (messagePatternSpies == nil) {
305        messagePatternSpies = [[NSMutableArray alloc] init];
306        spies[aMessagePattern] = messagePatternSpies;
307        [messagePatternSpies release];
308    }
309
310    NSValue *spyWrapper = [NSValue valueWithNonretainedObject:aSpy];
311
312    if ([messagePatternSpies containsObject:spyWrapper])
313        return;
314
315    [messagePatternSpies addObject:spyWrapper];
316}
317
318void KWClearObjectSpy(id anObject, id aSpy, KWMessagePattern *aMessagePattern) {
319    NSValue *key = [NSValue valueWithNonretainedObject:anObject];
320    NSMutableDictionary *spyArrayDictionary = KWMessageSpies[key];
321    NSMutableArray *spies = spyArrayDictionary[aMessagePattern];
322    NSValue *spyWrapper = [NSValue valueWithNonretainedObject:aSpy];
323    [spies removeObject:spyWrapper];
324}
325
326void KWClearAllMessageSpies(void) {
327    for (NSValue *objectKey in KWMessageSpies) {
328        id spiedObject = [objectKey nonretainedObjectValue];
329        if ([KWRestoredObjects containsObject:spiedObject]) {
330            continue;
331        }
332        KWRestoreOriginalClass(spiedObject);
333        [KWRestoredObjects addObject:spiedObject];
334    }
335    [KWMessageSpies removeAllObjects];
336}