Aspects.m 38 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945
  1. //
  2. // Aspects.m
  3. // Aspects - A delightful, simple library for aspect oriented programming.
  4. //
  5. // Copyright (c) 2014 Peter Steinberger. Licensed under the MIT license.
  6. //
  7. #import "Aspects.h"
  8. #import <libkern/OSAtomic.h>
  9. #import <objc/runtime.h>
  10. #import <objc/message.h>
  11. #define AspectLog(...)
  12. //#define AspectLog(...) do { NSLog(__VA_ARGS__); }while(0)
  13. #define AspectLogError(...) do { NSLog(__VA_ARGS__); }while(0)
  14. // Block internals.
  15. typedef NS_OPTIONS(int, AspectBlockFlags) {
  16. AspectBlockFlagsHasCopyDisposeHelpers = (1 << 25),
  17. AspectBlockFlagsHasSignature = (1 << 30)
  18. };
  19. typedef struct _AspectBlock {
  20. __unused Class isa;
  21. AspectBlockFlags flags;
  22. __unused int reserved;
  23. void (__unused *invoke)(struct _AspectBlock *block, ...);
  24. struct {
  25. unsigned long int reserved;
  26. unsigned long int size;
  27. // requires AspectBlockFlagsHasCopyDisposeHelpers
  28. void (*copy)(void *dst, const void *src);
  29. void (*dispose)(const void *);
  30. // requires AspectBlockFlagsHasSignature
  31. const char *signature;
  32. const char *layout;
  33. } *descriptor;
  34. // imported variables
  35. } *AspectBlockRef;
  36. @interface AspectInfo : NSObject <AspectInfo>
  37. - (id)initWithInstance:(__unsafe_unretained id)instance invocation:(NSInvocation *)invocation;
  38. @property (nonatomic, unsafe_unretained, readonly) id instance;
  39. @property (nonatomic, strong, readonly) NSArray *arguments;
  40. @property (nonatomic, strong, readonly) NSInvocation *originalInvocation;
  41. @end
  42. // Tracks a single aspect.
  43. @interface AspectIdentifier : NSObject
  44. + (instancetype)identifierWithSelector:(SEL)selector object:(id)object options:(AspectOptions)options block:(id)block error:(NSError **)error;
  45. - (BOOL)invokeWithInfo:(id<AspectInfo>)info;
  46. @property (nonatomic, assign) SEL selector;
  47. @property (nonatomic, strong) id block;
  48. @property (nonatomic, strong) NSMethodSignature *blockSignature;
  49. @property (nonatomic, weak) id object;
  50. @property (nonatomic, assign) AspectOptions options;
  51. @end
  52. // Tracks all aspects for an object/class.
  53. @interface AspectsContainer : NSObject
  54. - (void)addAspect:(AspectIdentifier *)aspect withOptions:(AspectOptions)injectPosition;
  55. - (BOOL)removeAspect:(id)aspect;
  56. - (BOOL)hasAspects;
  57. @property (atomic, copy) NSArray *beforeAspects;
  58. @property (atomic, copy) NSArray *insteadAspects;
  59. @property (atomic, copy) NSArray *afterAspects;
  60. @end
  61. @interface AspectTracker : NSObject
  62. - (id)initWithTrackedClass:(Class)trackedClass;
  63. @property (nonatomic, strong) Class trackedClass;
  64. @property (nonatomic, readonly) NSString *trackedClassName;
  65. @property (nonatomic, strong) NSMutableSet *selectorNames;
  66. @property (nonatomic, strong) NSMutableDictionary *selectorNamesToSubclassTrackers;
  67. - (void)addSubclassTracker:(AspectTracker *)subclassTracker hookingSelectorName:(NSString *)selectorName;
  68. - (void)removeSubclassTracker:(AspectTracker *)subclassTracker hookingSelectorName:(NSString *)selectorName;
  69. - (BOOL)subclassHasHookedSelectorName:(NSString *)selectorName;
  70. - (NSSet *)subclassTrackersHookingSelectorName:(NSString *)selectorName;
  71. @end
  72. @interface NSInvocation (Aspects)
  73. - (NSArray *)aspects_arguments;
  74. @end
  75. #define AspectPositionFilter 0x07
  76. #define AspectError(errorCode, errorDescription) do { \
  77. AspectLogError(@"Aspects: %@", errorDescription); \
  78. if (error) { *error = [NSError errorWithDomain:AspectErrorDomain code:errorCode userInfo:@{NSLocalizedDescriptionKey: errorDescription}]; }}while(0)
  79. NSString *const AspectErrorDomain = @"AspectErrorDomain";
  80. static NSString *const AspectsSubclassSuffix = @"_Aspects_";
  81. static NSString *const AspectsMessagePrefix = @"aspects_";
  82. @implementation NSObject (Aspects)
  83. ///////////////////////////////////////////////////////////////////////////////////////////
  84. #pragma mark - Public Aspects API
  85. + (id<AspectToken>)aspect_hookSelector:(SEL)selector
  86. withOptions:(AspectOptions)options
  87. usingBlock:(id)block
  88. error:(NSError **)error {
  89. return aspect_add((id)self, selector, options, block, error);
  90. }
  91. /// @return A token which allows to later deregister the aspect.
  92. - (id<AspectToken>)aspect_hookSelector:(SEL)selector
  93. withOptions:(AspectOptions)options
  94. usingBlock:(id)block
  95. error:(NSError **)error {
  96. return aspect_add(self, selector, options, block, error);
  97. }
  98. ///////////////////////////////////////////////////////////////////////////////////////////
  99. #pragma mark - Private Helper
  100. static id aspect_add(id self, SEL selector, AspectOptions options, id block, NSError **error) {
  101. NSCParameterAssert(self);
  102. NSCParameterAssert(selector);
  103. NSCParameterAssert(block);
  104. __block AspectIdentifier *identifier = nil;
  105. aspect_performLocked(^{
  106. if (aspect_isSelectorAllowedAndTrack(self, selector, options, error)) {
  107. AspectsContainer *aspectContainer = aspect_getContainerForObject(self, selector);
  108. identifier = [AspectIdentifier identifierWithSelector:selector object:self options:options block:block error:error];
  109. if (identifier) {
  110. [aspectContainer addAspect:identifier withOptions:options];
  111. // Modify the class to allow message interception.
  112. aspect_prepareClassAndHookSelector(self, selector, error);
  113. }
  114. }
  115. });
  116. return identifier;
  117. }
  118. static BOOL aspect_remove(AspectIdentifier *aspect, NSError **error) {
  119. NSCAssert([aspect isKindOfClass:AspectIdentifier.class], @"Must have correct type.");
  120. __block BOOL success = NO;
  121. aspect_performLocked(^{
  122. id self = aspect.object; // strongify
  123. if (self) {
  124. AspectsContainer *aspectContainer = aspect_getContainerForObject(self, aspect.selector);
  125. success = [aspectContainer removeAspect:aspect];
  126. aspect_cleanupHookedClassAndSelector(self, aspect.selector);
  127. // destroy token
  128. aspect.object = nil;
  129. aspect.block = nil;
  130. aspect.selector = NULL;
  131. }else {
  132. NSString *errrorDesc = [NSString stringWithFormat:@"Unable to deregister hook. Object already deallocated: %@", aspect];
  133. AspectError(AspectErrorRemoveObjectAlreadyDeallocated, errrorDesc);
  134. }
  135. });
  136. return success;
  137. }
  138. static void aspect_performLocked(dispatch_block_t block) {
  139. static OSSpinLock aspect_lock = OS_SPINLOCK_INIT;
  140. OSSpinLockLock(&aspect_lock);
  141. block();
  142. OSSpinLockUnlock(&aspect_lock);
  143. }
  144. static SEL aspect_aliasForSelector(SEL selector) {
  145. NSCParameterAssert(selector);
  146. return NSSelectorFromString([AspectsMessagePrefix stringByAppendingFormat:@"_%@", NSStringFromSelector(selector)]);
  147. }
  148. static NSMethodSignature *aspect_blockMethodSignature(id block, NSError **error) {
  149. AspectBlockRef layout = (__bridge void *)block;
  150. if (!(layout->flags & AspectBlockFlagsHasSignature)) {
  151. NSString *description = [NSString stringWithFormat:@"The block %@ doesn't contain a type signature.", block];
  152. AspectError(AspectErrorMissingBlockSignature, description);
  153. return nil;
  154. }
  155. void *desc = layout->descriptor;
  156. desc += 2 * sizeof(unsigned long int);
  157. if (layout->flags & AspectBlockFlagsHasCopyDisposeHelpers) {
  158. desc += 2 * sizeof(void *);
  159. }
  160. if (!desc) {
  161. NSString *description = [NSString stringWithFormat:@"The block %@ doesn't has a type signature.", block];
  162. AspectError(AspectErrorMissingBlockSignature, description);
  163. return nil;
  164. }
  165. const char *signature = (*(const char **)desc);
  166. return [NSMethodSignature signatureWithObjCTypes:signature];
  167. }
  168. static BOOL aspect_isCompatibleBlockSignature(NSMethodSignature *blockSignature, id object, SEL selector, NSError **error) {
  169. NSCParameterAssert(blockSignature);
  170. NSCParameterAssert(object);
  171. NSCParameterAssert(selector);
  172. BOOL signaturesMatch = YES;
  173. NSMethodSignature *methodSignature = [[object class] instanceMethodSignatureForSelector:selector];
  174. if (blockSignature.numberOfArguments > methodSignature.numberOfArguments) {
  175. signaturesMatch = NO;
  176. }else {
  177. if (blockSignature.numberOfArguments > 1) {
  178. const char *blockType = [blockSignature getArgumentTypeAtIndex:1];
  179. if (blockType[0] != '@') {
  180. signaturesMatch = NO;
  181. }
  182. }
  183. // Argument 0 is self/block, argument 1 is SEL or id<AspectInfo>. We start comparing at argument 2.
  184. // The block can have less arguments than the method, that's ok.
  185. if (signaturesMatch) {
  186. for (NSUInteger idx = 2; idx < blockSignature.numberOfArguments; idx++) {
  187. const char *methodType = [methodSignature getArgumentTypeAtIndex:idx];
  188. const char *blockType = [blockSignature getArgumentTypeAtIndex:idx];
  189. // Only compare parameter, not the optional type data.
  190. if (!methodType || !blockType || methodType[0] != blockType[0]) {
  191. signaturesMatch = NO; break;
  192. }
  193. }
  194. }
  195. }
  196. if (!signaturesMatch) {
  197. NSString *description = [NSString stringWithFormat:@"Block signature %@ doesn't match %@.", blockSignature, methodSignature];
  198. AspectError(AspectErrorIncompatibleBlockSignature, description);
  199. return NO;
  200. }
  201. return YES;
  202. }
  203. ///////////////////////////////////////////////////////////////////////////////////////////
  204. #pragma mark - Class + Selector Preparation
  205. static BOOL aspect_isMsgForwardIMP(IMP impl) {
  206. return impl == _objc_msgForward
  207. #if !defined(__arm64__)
  208. || impl == (IMP)_objc_msgForward_stret
  209. #endif
  210. ;
  211. }
  212. static IMP aspect_getMsgForwardIMP(NSObject *self, SEL selector) {
  213. IMP msgForwardIMP = _objc_msgForward;
  214. #if !defined(__arm64__)
  215. // As an ugly internal runtime implementation detail in the 32bit runtime, we need to determine of the method we hook returns a struct or anything larger than id.
  216. // https://developer.apple.com/library/mac/documentation/DeveloperTools/Conceptual/LowLevelABI/000-Introduction/introduction.html
  217. // https://github.com/ReactiveCocoa/ReactiveCocoa/issues/783
  218. // http://infocenter.arm.com/help/topic/com.arm.doc.ihi0042e/IHI0042E_aapcs.pdf (Section 5.4)
  219. Method method = class_getInstanceMethod(self.class, selector);
  220. const char *encoding = method_getTypeEncoding(method);
  221. BOOL methodReturnsStructValue = encoding[0] == _C_STRUCT_B;
  222. if (methodReturnsStructValue) {
  223. @try {
  224. NSUInteger valueSize = 0;
  225. NSGetSizeAndAlignment(encoding, &valueSize, NULL);
  226. if (valueSize == 1 || valueSize == 2 || valueSize == 4 || valueSize == 8) {
  227. methodReturnsStructValue = NO;
  228. }
  229. } @catch (__unused NSException *e) {}
  230. }
  231. if (methodReturnsStructValue) {
  232. msgForwardIMP = (IMP)_objc_msgForward_stret;
  233. }
  234. #endif
  235. return msgForwardIMP;
  236. }
  237. static void aspect_prepareClassAndHookSelector(NSObject *self, SEL selector, NSError **error) {
  238. NSCParameterAssert(selector);
  239. Class klass = aspect_hookClass(self, error);
  240. Method targetMethod = class_getInstanceMethod(klass, selector);
  241. IMP targetMethodIMP = method_getImplementation(targetMethod);
  242. if (!aspect_isMsgForwardIMP(targetMethodIMP)) {
  243. // Make a method alias for the existing method implementation, it not already copied.
  244. const char *typeEncoding = method_getTypeEncoding(targetMethod);
  245. SEL aliasSelector = aspect_aliasForSelector(selector);
  246. if (![klass instancesRespondToSelector:aliasSelector]) {
  247. __unused BOOL addedAlias = class_addMethod(klass, aliasSelector, method_getImplementation(targetMethod), typeEncoding);
  248. NSCAssert(addedAlias, @"Original implementation for %@ is already copied to %@ on %@", NSStringFromSelector(selector), NSStringFromSelector(aliasSelector), klass);
  249. }
  250. // We use forwardInvocation to hook in.
  251. class_replaceMethod(klass, selector, aspect_getMsgForwardIMP(self, selector), typeEncoding);
  252. AspectLog(@"Aspects: Installed hook for -[%@ %@].", klass, NSStringFromSelector(selector));
  253. }
  254. }
  255. // Will undo the runtime changes made.
  256. static void aspect_cleanupHookedClassAndSelector(NSObject *self, SEL selector) {
  257. NSCParameterAssert(self);
  258. NSCParameterAssert(selector);
  259. Class klass = object_getClass(self);
  260. BOOL isMetaClass = class_isMetaClass(klass);
  261. if (isMetaClass) {
  262. klass = (Class)self;
  263. }
  264. // Check if the method is marked as forwarded and undo that.
  265. Method targetMethod = class_getInstanceMethod(klass, selector);
  266. IMP targetMethodIMP = method_getImplementation(targetMethod);
  267. if (aspect_isMsgForwardIMP(targetMethodIMP)) {
  268. // Restore the original method implementation.
  269. const char *typeEncoding = method_getTypeEncoding(targetMethod);
  270. SEL aliasSelector = aspect_aliasForSelector(selector);
  271. Method originalMethod = class_getInstanceMethod(klass, aliasSelector);
  272. IMP originalIMP = method_getImplementation(originalMethod);
  273. NSCAssert(originalMethod, @"Original implementation for %@ not found %@ on %@", NSStringFromSelector(selector), NSStringFromSelector(aliasSelector), klass);
  274. class_replaceMethod(klass, selector, originalIMP, typeEncoding);
  275. AspectLog(@"Aspects: Removed hook for -[%@ %@].", klass, NSStringFromSelector(selector));
  276. }
  277. // Deregister global tracked selector
  278. aspect_deregisterTrackedSelector(self, selector);
  279. // Get the aspect container and check if there are any hooks remaining. Clean up if there are not.
  280. AspectsContainer *container = aspect_getContainerForObject(self, selector);
  281. if (!container.hasAspects) {
  282. // Destroy the container
  283. aspect_destroyContainerForObject(self, selector);
  284. // Figure out how the class was modified to undo the changes.
  285. NSString *className = NSStringFromClass(klass);
  286. if ([className hasSuffix:AspectsSubclassSuffix]) {
  287. Class originalClass = NSClassFromString([className stringByReplacingOccurrencesOfString:AspectsSubclassSuffix withString:@""]);
  288. NSCAssert(originalClass != nil, @"Original class must exist");
  289. object_setClass(self, originalClass);
  290. AspectLog(@"Aspects: %@ has been restored.", NSStringFromClass(originalClass));
  291. // We can only dispose the class pair if we can ensure that no instances exist using our subclass.
  292. // Since we don't globally track this, we can't ensure this - but there's also not much overhead in keeping it around.
  293. //objc_disposeClassPair(object.class);
  294. }else {
  295. // Class is most likely swizzled in place. Undo that.
  296. if (isMetaClass) {
  297. aspect_undoSwizzleClassInPlace((Class)self);
  298. }else if (self.class != klass) {
  299. aspect_undoSwizzleClassInPlace(klass);
  300. }
  301. }
  302. }
  303. }
  304. ///////////////////////////////////////////////////////////////////////////////////////////
  305. #pragma mark - Hook Class
  306. static Class aspect_hookClass(NSObject *self, NSError **error) {
  307. NSCParameterAssert(self);
  308. Class statedClass = self.class;
  309. Class baseClass = object_getClass(self);
  310. NSString *className = NSStringFromClass(baseClass);
  311. // Already subclassed
  312. if ([className hasSuffix:AspectsSubclassSuffix]) {
  313. return baseClass;
  314. // We swizzle a class object, not a single object.
  315. }else if (class_isMetaClass(baseClass)) {
  316. return aspect_swizzleClassInPlace((Class)self);
  317. // Probably a KVO'ed class. Swizzle in place. Also swizzle meta classes in place.
  318. }else if (statedClass != baseClass) {
  319. return aspect_swizzleClassInPlace(baseClass);
  320. }
  321. // Default case. Create dynamic subclass.
  322. const char *subclassName = [className stringByAppendingString:AspectsSubclassSuffix].UTF8String;
  323. Class subclass = objc_getClass(subclassName);
  324. if (subclass == nil) {
  325. subclass = objc_allocateClassPair(baseClass, subclassName, 0);
  326. if (subclass == nil) {
  327. NSString *errrorDesc = [NSString stringWithFormat:@"objc_allocateClassPair failed to allocate class %s.", subclassName];
  328. AspectError(AspectErrorFailedToAllocateClassPair, errrorDesc);
  329. return nil;
  330. }
  331. aspect_swizzleForwardInvocation(subclass);
  332. aspect_hookedGetClass(subclass, statedClass);
  333. aspect_hookedGetClass(object_getClass(subclass), statedClass);
  334. objc_registerClassPair(subclass);
  335. }
  336. object_setClass(self, subclass);
  337. return subclass;
  338. }
  339. static NSString *const AspectsForwardInvocationSelectorName = @"__aspects_forwardInvocation:";
  340. static void aspect_swizzleForwardInvocation(Class klass) {
  341. NSCParameterAssert(klass);
  342. // If there is no method, replace will act like class_addMethod.
  343. IMP originalImplementation = class_replaceMethod(klass, @selector(forwardInvocation:), (IMP)__ASPECTS_ARE_BEING_CALLED__, "v@:@");
  344. if (originalImplementation) {
  345. class_addMethod(klass, NSSelectorFromString(AspectsForwardInvocationSelectorName), originalImplementation, "v@:@");
  346. }
  347. AspectLog(@"Aspects: %@ is now aspect aware.", NSStringFromClass(klass));
  348. }
  349. static void aspect_undoSwizzleForwardInvocation(Class klass) {
  350. NSCParameterAssert(klass);
  351. Method originalMethod = class_getInstanceMethod(klass, NSSelectorFromString(AspectsForwardInvocationSelectorName));
  352. Method objectMethod = class_getInstanceMethod(NSObject.class, @selector(forwardInvocation:));
  353. // There is no class_removeMethod, so the best we can do is to retore the original implementation, or use a dummy.
  354. IMP originalImplementation = method_getImplementation(originalMethod ?: objectMethod);
  355. class_replaceMethod(klass, @selector(forwardInvocation:), originalImplementation, "v@:@");
  356. AspectLog(@"Aspects: %@ has been restored.", NSStringFromClass(klass));
  357. }
  358. static void aspect_hookedGetClass(Class class, Class statedClass) {
  359. NSCParameterAssert(class);
  360. NSCParameterAssert(statedClass);
  361. Method method = class_getInstanceMethod(class, @selector(class));
  362. IMP newIMP = imp_implementationWithBlock(^(id self) {
  363. return statedClass;
  364. });
  365. class_replaceMethod(class, @selector(class), newIMP, method_getTypeEncoding(method));
  366. }
  367. ///////////////////////////////////////////////////////////////////////////////////////////
  368. #pragma mark - Swizzle Class In Place
  369. static void _aspect_modifySwizzledClasses(void (^block)(NSMutableSet *swizzledClasses)) {
  370. static NSMutableSet *swizzledClasses;
  371. static dispatch_once_t pred;
  372. dispatch_once(&pred, ^{
  373. swizzledClasses = [NSMutableSet new];
  374. });
  375. @synchronized(swizzledClasses) {
  376. block(swizzledClasses);
  377. }
  378. }
  379. static Class aspect_swizzleClassInPlace(Class klass) {
  380. NSCParameterAssert(klass);
  381. NSString *className = NSStringFromClass(klass);
  382. _aspect_modifySwizzledClasses(^(NSMutableSet *swizzledClasses) {
  383. if (![swizzledClasses containsObject:className]) {
  384. aspect_swizzleForwardInvocation(klass);
  385. [swizzledClasses addObject:className];
  386. }
  387. });
  388. return klass;
  389. }
  390. static void aspect_undoSwizzleClassInPlace(Class klass) {
  391. NSCParameterAssert(klass);
  392. NSString *className = NSStringFromClass(klass);
  393. _aspect_modifySwizzledClasses(^(NSMutableSet *swizzledClasses) {
  394. if ([swizzledClasses containsObject:className]) {
  395. aspect_undoSwizzleForwardInvocation(klass);
  396. [swizzledClasses removeObject:className];
  397. }
  398. });
  399. }
  400. ///////////////////////////////////////////////////////////////////////////////////////////
  401. #pragma mark - Aspect Invoke Point
  402. // This is a macro so we get a cleaner stack trace.
  403. #define aspect_invoke(aspects, info) \
  404. for (AspectIdentifier *aspect in aspects) {\
  405. [aspect invokeWithInfo:info];\
  406. if (aspect.options & AspectOptionAutomaticRemoval) { \
  407. aspectsToRemove = [aspectsToRemove?:@[] arrayByAddingObject:aspect]; \
  408. } \
  409. }
  410. // This is the swizzled forwardInvocation: method.
  411. static void __ASPECTS_ARE_BEING_CALLED__(__unsafe_unretained NSObject *self, SEL selector, NSInvocation *invocation) {
  412. NSCParameterAssert(self);
  413. NSCParameterAssert(invocation);
  414. SEL originalSelector = invocation.selector;
  415. SEL aliasSelector = aspect_aliasForSelector(invocation.selector);
  416. invocation.selector = aliasSelector;
  417. AspectsContainer *objectContainer = objc_getAssociatedObject(self, aliasSelector);
  418. AspectsContainer *classContainer = aspect_getContainerForClass(object_getClass(self), aliasSelector);
  419. AspectInfo *info = [[AspectInfo alloc] initWithInstance:self invocation:invocation];
  420. NSArray *aspectsToRemove = nil;
  421. // Before hooks.
  422. aspect_invoke(classContainer.beforeAspects, info);
  423. aspect_invoke(objectContainer.beforeAspects, info);
  424. // Instead hooks.
  425. BOOL respondsToAlias = YES;
  426. if (objectContainer.insteadAspects.count || classContainer.insteadAspects.count) {
  427. aspect_invoke(classContainer.insteadAspects, info);
  428. aspect_invoke(objectContainer.insteadAspects, info);
  429. }else {
  430. Class klass = object_getClass(invocation.target);
  431. do {
  432. if ((respondsToAlias = [klass instancesRespondToSelector:aliasSelector])) {
  433. [invocation invoke];
  434. break;
  435. }
  436. }while (!respondsToAlias && (klass = class_getSuperclass(klass)));
  437. }
  438. // After hooks.
  439. aspect_invoke(classContainer.afterAspects, info);
  440. aspect_invoke(objectContainer.afterAspects, info);
  441. // If no hooks are installed, call original implementation (usually to throw an exception)
  442. if (!respondsToAlias) {
  443. invocation.selector = originalSelector;
  444. SEL originalForwardInvocationSEL = NSSelectorFromString(AspectsForwardInvocationSelectorName);
  445. if ([self respondsToSelector:originalForwardInvocationSEL]) {
  446. ((void( *)(id, SEL, NSInvocation *))objc_msgSend)(self, originalForwardInvocationSEL, invocation);
  447. }else {
  448. [self doesNotRecognizeSelector:invocation.selector];
  449. }
  450. }
  451. // Remove any hooks that are queued for deregistration.
  452. [aspectsToRemove makeObjectsPerformSelector:@selector(remove)];
  453. }
  454. #undef aspect_invoke
  455. ///////////////////////////////////////////////////////////////////////////////////////////
  456. #pragma mark - Aspect Container Management
  457. // Loads or creates the aspect container.
  458. static AspectsContainer *aspect_getContainerForObject(NSObject *self, SEL selector) {
  459. NSCParameterAssert(self);
  460. SEL aliasSelector = aspect_aliasForSelector(selector);
  461. AspectsContainer *aspectContainer = objc_getAssociatedObject(self, aliasSelector);
  462. if (!aspectContainer) {
  463. aspectContainer = [AspectsContainer new];
  464. objc_setAssociatedObject(self, aliasSelector, aspectContainer, OBJC_ASSOCIATION_RETAIN);
  465. }
  466. return aspectContainer;
  467. }
  468. static AspectsContainer *aspect_getContainerForClass(Class klass, SEL selector) {
  469. NSCParameterAssert(klass);
  470. AspectsContainer *classContainer = nil;
  471. do {
  472. classContainer = objc_getAssociatedObject(klass, selector);
  473. if (classContainer.hasAspects) break;
  474. }while ((klass = class_getSuperclass(klass)));
  475. return classContainer;
  476. }
  477. static void aspect_destroyContainerForObject(id<NSObject> self, SEL selector) {
  478. NSCParameterAssert(self);
  479. SEL aliasSelector = aspect_aliasForSelector(selector);
  480. objc_setAssociatedObject(self, aliasSelector, nil, OBJC_ASSOCIATION_RETAIN);
  481. }
  482. ///////////////////////////////////////////////////////////////////////////////////////////
  483. #pragma mark - Selector Blacklist Checking
  484. static NSMutableDictionary *aspect_getSwizzledClassesDict() {
  485. static NSMutableDictionary *swizzledClassesDict;
  486. static dispatch_once_t pred;
  487. dispatch_once(&pred, ^{
  488. swizzledClassesDict = [NSMutableDictionary new];
  489. });
  490. return swizzledClassesDict;
  491. }
  492. static BOOL aspect_isSelectorAllowedAndTrack(NSObject *self, SEL selector, AspectOptions options, NSError **error) {
  493. static NSSet *disallowedSelectorList;
  494. static dispatch_once_t pred;
  495. dispatch_once(&pred, ^{
  496. disallowedSelectorList = [NSSet setWithObjects:@"retain", @"release", @"autorelease", @"forwardInvocation:", nil];
  497. });
  498. // Check against the blacklist.
  499. NSString *selectorName = NSStringFromSelector(selector);
  500. if ([disallowedSelectorList containsObject:selectorName]) {
  501. NSString *errorDescription = [NSString stringWithFormat:@"Selector %@ is blacklisted.", selectorName];
  502. AspectError(AspectErrorSelectorBlacklisted, errorDescription);
  503. return NO;
  504. }
  505. // Additional checks.
  506. AspectOptions position = options&AspectPositionFilter;
  507. if ([selectorName isEqualToString:@"dealloc"] && position != AspectPositionBefore) {
  508. NSString *errorDesc = @"AspectPositionBefore is the only valid position when hooking dealloc.";
  509. AspectError(AspectErrorSelectorDeallocPosition, errorDesc);
  510. return NO;
  511. }
  512. if (![self respondsToSelector:selector] && ![self.class instancesRespondToSelector:selector]) {
  513. NSString *errorDesc = [NSString stringWithFormat:@"Unable to find selector -[%@ %@].", NSStringFromClass(self.class), selectorName];
  514. AspectError(AspectErrorDoesNotRespondToSelector, errorDesc);
  515. return NO;
  516. }
  517. // Search for the current class and the class hierarchy IF we are modifying a class object
  518. if (class_isMetaClass(object_getClass(self))) {
  519. Class klass = [self class];
  520. NSMutableDictionary *swizzledClassesDict = aspect_getSwizzledClassesDict();
  521. Class currentClass = [self class];
  522. AspectTracker *tracker = swizzledClassesDict[currentClass];
  523. if ([tracker subclassHasHookedSelectorName:selectorName]) {
  524. NSSet *subclassTracker = [tracker subclassTrackersHookingSelectorName:selectorName];
  525. NSSet *subclassNames = [subclassTracker valueForKey:@"trackedClassName"];
  526. NSString *errorDescription = [NSString stringWithFormat:@"Error: %@ already hooked subclasses: %@. A method can only be hooked once per class hierarchy.", selectorName, subclassNames];
  527. AspectError(AspectErrorSelectorAlreadyHookedInClassHierarchy, errorDescription);
  528. return NO;
  529. }
  530. do {
  531. tracker = swizzledClassesDict[currentClass];
  532. if ([tracker.selectorNames containsObject:selectorName]) {
  533. if (klass == currentClass) {
  534. // Already modified and topmost!
  535. return YES;
  536. }
  537. NSString *errorDescription = [NSString stringWithFormat:@"Error: %@ already hooked in %@. A method can only be hooked once per class hierarchy.", selectorName, NSStringFromClass(currentClass)];
  538. AspectError(AspectErrorSelectorAlreadyHookedInClassHierarchy, errorDescription);
  539. return NO;
  540. }
  541. } while ((currentClass = class_getSuperclass(currentClass)));
  542. // Add the selector as being modified.
  543. currentClass = klass;
  544. AspectTracker *subclassTracker = nil;
  545. do {
  546. tracker = swizzledClassesDict[currentClass];
  547. if (!tracker) {
  548. tracker = [[AspectTracker alloc] initWithTrackedClass:currentClass];
  549. swizzledClassesDict[(id<NSCopying>)currentClass] = tracker;
  550. }
  551. if (subclassTracker) {
  552. [tracker addSubclassTracker:subclassTracker hookingSelectorName:selectorName];
  553. } else {
  554. [tracker.selectorNames addObject:selectorName];
  555. }
  556. // All superclasses get marked as having a subclass that is modified.
  557. subclassTracker = tracker;
  558. }while ((currentClass = class_getSuperclass(currentClass)));
  559. } else {
  560. return YES;
  561. }
  562. return YES;
  563. }
  564. static void aspect_deregisterTrackedSelector(id self, SEL selector) {
  565. if (!class_isMetaClass(object_getClass(self))) return;
  566. NSMutableDictionary *swizzledClassesDict = aspect_getSwizzledClassesDict();
  567. NSString *selectorName = NSStringFromSelector(selector);
  568. Class currentClass = [self class];
  569. AspectTracker *subclassTracker = nil;
  570. do {
  571. AspectTracker *tracker = swizzledClassesDict[currentClass];
  572. if (subclassTracker) {
  573. [tracker removeSubclassTracker:subclassTracker hookingSelectorName:selectorName];
  574. } else {
  575. [tracker.selectorNames removeObject:selectorName];
  576. }
  577. if (tracker.selectorNames.count == 0 && tracker.selectorNamesToSubclassTrackers) {
  578. [swizzledClassesDict removeObjectForKey:currentClass];
  579. }
  580. subclassTracker = tracker;
  581. }while ((currentClass = class_getSuperclass(currentClass)));
  582. }
  583. @end
  584. @implementation AspectTracker
  585. - (id)initWithTrackedClass:(Class)trackedClass {
  586. if (self = [super init]) {
  587. _trackedClass = trackedClass;
  588. _selectorNames = [NSMutableSet new];
  589. _selectorNamesToSubclassTrackers = [NSMutableDictionary new];
  590. }
  591. return self;
  592. }
  593. - (BOOL)subclassHasHookedSelectorName:(NSString *)selectorName {
  594. return self.selectorNamesToSubclassTrackers[selectorName] != nil;
  595. }
  596. - (void)addSubclassTracker:(AspectTracker *)subclassTracker hookingSelectorName:(NSString *)selectorName {
  597. NSMutableSet *trackerSet = self.selectorNamesToSubclassTrackers[selectorName];
  598. if (!trackerSet) {
  599. trackerSet = [NSMutableSet new];
  600. self.selectorNamesToSubclassTrackers[selectorName] = trackerSet;
  601. }
  602. [trackerSet addObject:subclassTracker];
  603. }
  604. - (void)removeSubclassTracker:(AspectTracker *)subclassTracker hookingSelectorName:(NSString *)selectorName {
  605. NSMutableSet *trackerSet = self.selectorNamesToSubclassTrackers[selectorName];
  606. [trackerSet removeObject:subclassTracker];
  607. if (trackerSet.count == 0) {
  608. [self.selectorNamesToSubclassTrackers removeObjectForKey:selectorName];
  609. }
  610. }
  611. - (NSSet *)subclassTrackersHookingSelectorName:(NSString *)selectorName {
  612. NSMutableSet *hookingSubclassTrackers = [NSMutableSet new];
  613. for (AspectTracker *tracker in self.selectorNamesToSubclassTrackers[selectorName]) {
  614. if ([tracker.selectorNames containsObject:selectorName]) {
  615. [hookingSubclassTrackers addObject:tracker];
  616. }
  617. [hookingSubclassTrackers unionSet:[tracker subclassTrackersHookingSelectorName:selectorName]];
  618. }
  619. return hookingSubclassTrackers;
  620. }
  621. - (NSString *)trackedClassName {
  622. return NSStringFromClass(self.trackedClass);
  623. }
  624. - (NSString *)description {
  625. return [NSString stringWithFormat:@"<%@: %@, trackedClass: %@, selectorNames:%@, subclass selector names: %@>", self.class, self, NSStringFromClass(self.trackedClass), self.selectorNames, self.selectorNamesToSubclassTrackers.allKeys];
  626. }
  627. @end
  628. ///////////////////////////////////////////////////////////////////////////////////////////
  629. #pragma mark - NSInvocation (Aspects)
  630. @implementation NSInvocation (Aspects)
  631. // Thanks to the ReactiveCocoa team for providing a generic solution for this.
  632. - (id)aspect_argumentAtIndex:(NSUInteger)index {
  633. const char *argType = [self.methodSignature getArgumentTypeAtIndex:index];
  634. // Skip const type qualifier.
  635. if (argType[0] == _C_CONST) argType++;
  636. #define WRAP_AND_RETURN(type) do { type val = 0; [self getArgument:&val atIndex:(NSInteger)index]; return @(val); } while (0)
  637. if (strcmp(argType, @encode(id)) == 0 || strcmp(argType, @encode(Class)) == 0) {
  638. __autoreleasing id returnObj;
  639. [self getArgument:&returnObj atIndex:(NSInteger)index];
  640. return returnObj;
  641. } else if (strcmp(argType, @encode(SEL)) == 0) {
  642. SEL selector = 0;
  643. [self getArgument:&selector atIndex:(NSInteger)index];
  644. return NSStringFromSelector(selector);
  645. } else if (strcmp(argType, @encode(Class)) == 0) {
  646. __autoreleasing Class theClass = Nil;
  647. [self getArgument:&theClass atIndex:(NSInteger)index];
  648. return theClass;
  649. // Using this list will box the number with the appropriate constructor, instead of the generic NSValue.
  650. } else if (strcmp(argType, @encode(char)) == 0) {
  651. WRAP_AND_RETURN(char);
  652. } else if (strcmp(argType, @encode(int)) == 0) {
  653. WRAP_AND_RETURN(int);
  654. } else if (strcmp(argType, @encode(short)) == 0) {
  655. WRAP_AND_RETURN(short);
  656. } else if (strcmp(argType, @encode(long)) == 0) {
  657. WRAP_AND_RETURN(long);
  658. } else if (strcmp(argType, @encode(long long)) == 0) {
  659. WRAP_AND_RETURN(long long);
  660. } else if (strcmp(argType, @encode(unsigned char)) == 0) {
  661. WRAP_AND_RETURN(unsigned char);
  662. } else if (strcmp(argType, @encode(unsigned int)) == 0) {
  663. WRAP_AND_RETURN(unsigned int);
  664. } else if (strcmp(argType, @encode(unsigned short)) == 0) {
  665. WRAP_AND_RETURN(unsigned short);
  666. } else if (strcmp(argType, @encode(unsigned long)) == 0) {
  667. WRAP_AND_RETURN(unsigned long);
  668. } else if (strcmp(argType, @encode(unsigned long long)) == 0) {
  669. WRAP_AND_RETURN(unsigned long long);
  670. } else if (strcmp(argType, @encode(float)) == 0) {
  671. WRAP_AND_RETURN(float);
  672. } else if (strcmp(argType, @encode(double)) == 0) {
  673. WRAP_AND_RETURN(double);
  674. } else if (strcmp(argType, @encode(BOOL)) == 0) {
  675. WRAP_AND_RETURN(BOOL);
  676. } else if (strcmp(argType, @encode(bool)) == 0) {
  677. WRAP_AND_RETURN(BOOL);
  678. } else if (strcmp(argType, @encode(char *)) == 0) {
  679. WRAP_AND_RETURN(const char *);
  680. } else if (strcmp(argType, @encode(void (^)(void))) == 0) {
  681. __unsafe_unretained id block = nil;
  682. [self getArgument:&block atIndex:(NSInteger)index];
  683. return [block copy];
  684. } else {
  685. NSUInteger valueSize = 0;
  686. NSGetSizeAndAlignment(argType, &valueSize, NULL);
  687. unsigned char valueBytes[valueSize];
  688. [self getArgument:valueBytes atIndex:(NSInteger)index];
  689. return [NSValue valueWithBytes:valueBytes objCType:argType];
  690. }
  691. return nil;
  692. #undef WRAP_AND_RETURN
  693. }
  694. - (NSArray *)aspects_arguments {
  695. NSMutableArray *argumentsArray = [NSMutableArray array];
  696. for (NSUInteger idx = 2; idx < self.methodSignature.numberOfArguments; idx++) {
  697. [argumentsArray addObject:[self aspect_argumentAtIndex:idx] ?: NSNull.null];
  698. }
  699. return [argumentsArray copy];
  700. }
  701. @end
  702. ///////////////////////////////////////////////////////////////////////////////////////////
  703. #pragma mark - AspectIdentifier
  704. @implementation AspectIdentifier
  705. + (instancetype)identifierWithSelector:(SEL)selector object:(id)object options:(AspectOptions)options block:(id)block error:(NSError **)error {
  706. NSCParameterAssert(block);
  707. NSCParameterAssert(selector);
  708. NSMethodSignature *blockSignature = aspect_blockMethodSignature(block, error); // TODO: check signature compatibility, etc.
  709. if (!aspect_isCompatibleBlockSignature(blockSignature, object, selector, error)) {
  710. return nil;
  711. }
  712. AspectIdentifier *identifier = nil;
  713. if (blockSignature) {
  714. identifier = [AspectIdentifier new];
  715. identifier.selector = selector;
  716. identifier.block = block;
  717. identifier.blockSignature = blockSignature;
  718. identifier.options = options;
  719. identifier.object = object; // weak
  720. }
  721. return identifier;
  722. }
  723. - (BOOL)invokeWithInfo:(id<AspectInfo>)info {
  724. NSInvocation *blockInvocation = [NSInvocation invocationWithMethodSignature:self.blockSignature];
  725. NSInvocation *originalInvocation = info.originalInvocation;
  726. NSUInteger numberOfArguments = self.blockSignature.numberOfArguments;
  727. // Be extra paranoid. We already check that on hook registration.
  728. if (numberOfArguments > originalInvocation.methodSignature.numberOfArguments) {
  729. AspectLogError(@"Block has too many arguments. Not calling %@", info);
  730. return NO;
  731. }
  732. // The `self` of the block will be the AspectInfo. Optional.
  733. if (numberOfArguments > 1) {
  734. [blockInvocation setArgument:&info atIndex:1];
  735. }
  736. void *argBuf = NULL;
  737. for (NSUInteger idx = 2; idx < numberOfArguments; idx++) {
  738. const char *type = [originalInvocation.methodSignature getArgumentTypeAtIndex:idx];
  739. NSUInteger argSize;
  740. NSGetSizeAndAlignment(type, &argSize, NULL);
  741. if (!(argBuf = reallocf(argBuf, argSize))) {
  742. AspectLogError(@"Failed to allocate memory for block invocation.");
  743. return NO;
  744. }
  745. [originalInvocation getArgument:argBuf atIndex:idx];
  746. [blockInvocation setArgument:argBuf atIndex:idx];
  747. }
  748. [blockInvocation invokeWithTarget:self.block];
  749. if (argBuf != NULL) {
  750. free(argBuf);
  751. }
  752. return YES;
  753. }
  754. - (NSString *)description {
  755. return [NSString stringWithFormat:@"<%@: %p, SEL:%@ object:%@ options:%tu block:%@ (#%tu args)>", self.class, self, NSStringFromSelector(self.selector), self.object, self.options, self.block, self.blockSignature.numberOfArguments];
  756. }
  757. - (BOOL)remove {
  758. return aspect_remove(self, NULL);
  759. }
  760. @end
  761. ///////////////////////////////////////////////////////////////////////////////////////////
  762. #pragma mark - AspectsContainer
  763. @implementation AspectsContainer
  764. - (BOOL)hasAspects {
  765. return self.beforeAspects.count > 0 || self.insteadAspects.count > 0 || self.afterAspects.count > 0;
  766. }
  767. - (void)addAspect:(AspectIdentifier *)aspect withOptions:(AspectOptions)options {
  768. NSParameterAssert(aspect);
  769. NSUInteger position = options&AspectPositionFilter;
  770. switch (position) {
  771. case AspectPositionBefore: self.beforeAspects = [(self.beforeAspects ?:@[]) arrayByAddingObject:aspect]; break;
  772. case AspectPositionInstead: self.insteadAspects = [(self.insteadAspects?:@[]) arrayByAddingObject:aspect]; break;
  773. case AspectPositionAfter: self.afterAspects = [(self.afterAspects ?:@[]) arrayByAddingObject:aspect]; break;
  774. }
  775. }
  776. - (BOOL)removeAspect:(id)aspect {
  777. for (NSString *aspectArrayName in @[NSStringFromSelector(@selector(beforeAspects)),
  778. NSStringFromSelector(@selector(insteadAspects)),
  779. NSStringFromSelector(@selector(afterAspects))]) {
  780. NSArray *array = [self valueForKey:aspectArrayName];
  781. NSUInteger index = [array indexOfObjectIdenticalTo:aspect];
  782. if (array && index != NSNotFound) {
  783. NSMutableArray *newArray = [NSMutableArray arrayWithArray:array];
  784. [newArray removeObjectAtIndex:index];
  785. [self setValue:newArray forKey:aspectArrayName];
  786. return YES;
  787. }
  788. }
  789. return NO;
  790. }
  791. - (NSString *)description {
  792. return [NSString stringWithFormat:@"<%@: %p, before:%@, instead:%@, after:%@>", self.class, self, self.beforeAspects, self.insteadAspects, self.afterAspects];
  793. }
  794. @end
  795. ///////////////////////////////////////////////////////////////////////////////////////////
  796. #pragma mark - AspectInfo
  797. @implementation AspectInfo
  798. @synthesize arguments = _arguments;
  799. - (id)initWithInstance:(__unsafe_unretained id)instance invocation:(NSInvocation *)invocation {
  800. NSCParameterAssert(instance);
  801. NSCParameterAssert(invocation);
  802. if (self = [super init]) {
  803. _instance = instance;
  804. _originalInvocation = invocation;
  805. }
  806. return self;
  807. }
  808. - (NSArray *)arguments {
  809. // Lazily evaluate arguments, boxing is expensive.
  810. if (!_arguments) {
  811. _arguments = self.originalInvocation.aspects_arguments;
  812. }
  813. return _arguments;
  814. }
  815. @end