Aspects源码解读与runtime

开篇

由于近期在搞优化无侵入埋点，在github上发现了一个基于swizzling method 的轻量级开源框架 Aspects，其总共只有一个类文件。为了加强了解，在这里写下我学习和总结的心得。

基本原理

使用Aspects，能够在一个方法前/后插入一段代码，也能直接替换整个方法的实现。

这是怎么做到的呢？

我们都知道oc是一门动态语言，在执行一个函数的时候，其实是在给对象发消息，根据函数名生成selector，通过selector找到指向具体函数实现的IMP，然后才执行真正的函数逻辑。在这个基础上，如果我们人为的改变selector与IMP的对应关系，那就能让原本的函数做一些其它的事情。我们先来看一下oc中类的结构：

在其中，我们现在需要关注的是objc_method_list **methodLists，这个方法链表内储存的是Method类型：

Method结构体中包含一个SEL和IMP，实际上相当于在SEL和IMP之间作了一个映射。有了SEL，我们便可以找到对应的IMP，从而调用方法的实现代码。

但是函数具体是怎么执行的呢？见下图(转):

通过上图我们可以发现，在函数执行的过程当中，如果selector有对应的IMP，则直接执行，如果没有，那么就会进入消息转发流程，依次有
resolveInstanceMethod
forwardingTargetForSelector
methodSignatureForSelector
forwardInvocation等4个函数。
Aspects选择了在forwardInvocation这个阶段处理是因为：
resolvedInstanceMethod 适合给类/对象动态添加一个相应的实现，
forwardingTargetForSelector 适合将消息转发给其他对象处理,相对而言，forwardInvocation 是里面最灵活，最能符合需求的。

因此 Aspects 的方案就是，对于待 hook 的 selector，将其指向 objc_msgForward / _objc_msgForward_stret ,同时生成一个新的 aliasSelector 指向原来的 IMP，并且 hook 住 forwardInvocation 函数，使他指向自己的实现。按照上面的思路，当被 hook 的 selector 被执行的时候，首先根据 selector 找到了 objc_msgForward / _objc_msgForward_stret ,而这个会触发消息转发，从而进入 forwardInvocation。同时由于 forwardInvocation 的指向也被修改了，因此会转入新的forwardInvocation 函数，在里面执行需要嵌入的附加代码，完成之后，再转回原来的 IMP。

分析源码

在看具体实现之前，我们需要了解Aspects类里面定义的几个重要的数据结构：

AspectOptions

typedef NS_OPTIONS(NSUInteger, AspectOptions) {
    AspectPositionAfter   = 0,            /// Called after the original implementation (default)
    AspectPositionInstead = 1,            /// Will replace the original implementation.
    AspectPositionBefore  = 2,            /// Called before the original implementation.
    
    AspectOptionAutomaticRemoval = 1 << 3 /// Will remove the hook after the first execution.
};

这里表示了block的执行时机，可以是在原函数之前、之后或者完全替换掉原函数。

AspectsContainer

@interface AspectsContainer : NSObject
- (void)addAspect:(AspectIdentifier *)aspect withOptions:(AspectOptions)injectPosition;
- (BOOL)removeAspect:(id)aspect;
- (BOOL)hasAspects;
@property (atomic, copy) NSArray *beforeAspects;
@property (atomic, copy) NSArray *insteadAspects;
@property (atomic, copy) NSArray *afterAspects;
@end

跟踪对象的所有情况

AspectIdentifier

@interface AspectIdentifier : NSObject
+ (instancetype)identifierWithSelector:(SEL)selector object:(id)object options:(AspectOptions)options block:(id)block error:(NSError **)error;
- (BOOL)invokeWithInfo:(id)info;
@property (nonatomic, assign) SEL selector;
@property (nonatomic, strong) id block;
@property (nonatomic, strong) NSMethodSignature *blockSignature;
@property (nonatomic, weak) id object;
@property (nonatomic, assign) AspectOptions options;
@end

一个具体的Aspect，其中包括了block主体、执行时机options、方法SEL、方法参数、签名。将我们传入的block封装成AspectIdentifier。

AspectInfo

@interface AspectInfo : NSObject 
- (id)initWithInstance:(__unsafe_unretained id)instance invocation:(NSInvocation *)invocation;
@property (nonatomic, unsafe_unretained, readonly) id instance;
@property (nonatomic, strong, readonly) NSArray *arguments;
@property (nonatomic, strong, readonly) NSInvocation *originalInvocation;
@end

内部主要是NSInvocation信息，将其封装了一层。

AspectTracker

@interface AspectTracker : NSObject
- (id)initWithTrackedClass:(Class)trackedClass;
@property (nonatomic, strong) Class trackedClass;
@property (nonatomic, readonly) NSString *trackedClassName;
@property (nonatomic, strong) NSMutableSet *selectorNames;
@property (nonatomic, strong) NSMutableDictionary *selectorNamesToSubclassTrackers;
- (void)addSubclassTracker:(AspectTracker *)subclassTracker hookingSelectorName:(NSString *)selectorName;
- (void)removeSubclassTracker:(AspectTracker *)subclassTracker hookingSelectorName:(NSString *)selectorName;
- (BOOL)subclassHasHookedSelectorName:(NSString *)selectorName;
- (NSSet *)subclassTrackersHookingSelectorName:(NSString *)selectorName;
@end

用于跟踪类是否已经被hook并给其上标记，防止重复替换方法。

具体实现

我们从入口函数开始：

static id aspect_add(id self, SEL selector, AspectOptions options, id block, NSError **error) {
    __block AspectIdentifier *identifier = nil;
    aspect_performLocked(^{
        if (aspect_isSelectorAllowedAndTrack(self, selector, options, error)) {
            AspectsContainer *aspectContainer = aspect_getContainerForObject(self, selector);
            identifier = [AspectIdentifier identifierWithSelector:selector object:self options:options block:block error:error];
            if (identifier) {
                [aspectContainer addAspect:identifier withOptions:options];

                // Modify the class to allow message interception.
                aspect_prepareClassAndHookSelector(self, selector, error);
            }
        }
    });
    return identifier;
}

这段函数总共做了哪些事情呢？
1.aspect_isSelectorAllowedAndTrack()对传进来的类和参数进行判断。利用黑名单来判断传进来的类和参数能否被hook,比如对于某些方法（retain、release、autorelease、forwardInvocation）是不能hook的，同时对于dealloc来说，只能把代码插入到dealloc之前，还需要检测类是否能响应方法。另外对于类对象而言，还需要对类的层级结构(父类、子类)进行判断，防止同一方法被hook多次。

2.使用aspect_getContainerForObject()创建一个AspectsContainer容器，其中用到了动态添加属性。

static AspectsContainer *aspect_getContainerForObject(NSObject *self, SEL selector) {
    NSCParameterAssert(self);
    SEL aliasSelector = aspect_aliasForSelector(selector);
    AspectsContainer *aspectContainer = objc_getAssociatedObject(self, aliasSelector);
    if (!aspectContainer) {
        aspectContainer = [AspectsContainer new];
        objc_setAssociatedObject(self, aliasSelector, aspectContainer, OBJC_ASSOCIATION_RETAIN);
    }
    return aspectContainer;
}

3.使用+ (instancetype)identifierWithSelector:(SEL)selector object:(id)object options:(AspectOptions)options block:(id)block error:(NSError **)error;创建一个AspectsIdentifier并把其加入之前创建的AspectsContainer中去，由于我们使用Aspects是用block来替换方法的，所以就需要将我们传入的block替换成具体的方法，所以AspectsIdentifier中就有一个NSMethodSignature *blockSignature;方法签名属性。具体的转换方法如下：

static NSMethodSignature *aspect_blockMethodSignature(id block, NSError **error) {
    AspectBlockRef layout = (__bridge void *)block;
    if (!(layout->flags & AspectBlockFlagsHasSignature)) {
        NSString *description = [NSString stringWithFormat:@"The block %@ doesn't contain a type signature.", block];
        AspectError(AspectErrorMissingBlockSignature, description);
        return nil;
    }
    void *desc = layout->descriptor;
    desc += 2 * sizeof(unsigned long int);
    if (layout->flags & AspectBlockFlagsHasCopyDisposeHelpers) {
        desc += 2 * sizeof(void *);
    }
    if (!desc) {
        NSString *description = [NSString stringWithFormat:@"The block %@ doesn't has a type signature.", block];
        AspectError(AspectErrorMissingBlockSignature, description);
        return nil;
    }
    const char *signature = (*(const char **)desc);
    return [NSMethodSignature signatureWithObjCTypes:signature];
}

其中AspectBlockRef是Aspects内部自定义的一个block结构：

typedef struct _AspectBlock {
    __unused Class isa;
    AspectBlockFlags flags;
    __unused int reserved;
    void (__unused *invoke)(struct _AspectBlock *block, ...);
    struct {
        unsigned long int reserved;
        unsigned long int size;
        // requires AspectBlockFlagsHasCopyDisposeHelpers
        void (*copy)(void *dst, const void *src);
        void (*dispose)(const void *);
        // requires AspectBlockFlagsHasSignature
        const char *signature;
        const char *layout;
    } *descriptor;
    // imported variables
} *AspectBlockRef;

在将block转换成方法签名之后，使用aspect_isCompatibleBlockSignature来判断转换出的方法签名和替换方法的方法签名参数等是否一致：

static BOOL aspect_isCompatibleBlockSignature(NSMethodSignature *blockSignature, id object, SEL selector, NSError **error) {
    BOOL signaturesMatch = YES;
    NSMethodSignature *methodSignature = [[object class] instanceMethodSignatureForSelector:selector];
    if (blockSignature.numberOfArguments > methodSignature.numberOfArguments) {
        signaturesMatch = NO;
    }else {
        if (blockSignature.numberOfArguments > 1) {
            const char *blockType = [blockSignature getArgumentTypeAtIndex:1];
            if (blockType[0] != '@') {
                signaturesMatch = NO;
            }
        }
        // Argument 0 is self/block, argument 1 is SEL or id. We start comparing at argument 2.
        // The block can have less arguments than the method, that's ok.
        if (signaturesMatch) {
            for (NSUInteger idx = 2; idx < blockSignature.numberOfArguments; idx++) {
                const char *methodType = [methodSignature getArgumentTypeAtIndex:idx];
                const char *blockType = [blockSignature getArgumentTypeAtIndex:idx];
                // Only compare parameter, not the optional type data.
                if (!methodType || !blockType || methodType[0] != blockType[0]) {
                    signaturesMatch = NO; break;
                }
            }
        }
    }
    if (!signaturesMatch) {
        NSString *description = [NSString stringWithFormat:@"Block signature %@ doesn't match %@.", blockSignature, methodSignature];
        AspectError(AspectErrorIncompatibleBlockSignature, description);
        return NO;
    }
    return YES;
}

一般的方法都有两个隐藏的参数，第一个是self，第二个是SEL，所以要从之后开始比较。

4.真正核心的交换方法aspect_prepareClassAndHookSelector。当forwardInvocation被hook之后，将对传入的selector进行hook，将selector指向了转发IMP，同时生成一个aliasSelector，指向原来的IMP。

static void aspect_prepareClassAndHookSelector(NSObject *self, SEL selector, NSError **error) {
    Class klass = aspect_hookClass(self, error);
    Method targetMethod = class_getInstanceMethod(klass, selector);
    IMP targetMethodIMP = method_getImplementation(targetMethod);
    //如果发现 selector 已经指向了转发 IMP ,那就就不需要进行交换了
    if (!aspect_isMsgForwardIMP(targetMethodIMP)) {
        //* Make a method alias for the existing method implementation, it not already copied.
        const char *typeEncoding = method_getTypeEncoding(targetMethod);
        SEL aliasSelector = aspect_aliasForSelector(selector);
        if (![klass instancesRespondToSelector:aliasSelector]) {
            __unused BOOL addedAlias = class_addMethod(klass, aliasSelector, method_getImplementation(targetMethod), typeEncoding);
            NSCAssert(addedAlias, @"Original implementation for %@ is already copied to %@ on %@", NSStringFromSelector(selector), NSStringFromSelector(aliasSelector), klass);
        }
        // We use forwardInvocation to hook in.
        class_replaceMethod(klass, selector, aspect_getMsgForwardIMP(self, selector), typeEncoding);
        AspectLog(@"Aspects: Installed hook for -[%@ %@].", klass, NSStringFromSelector(selector));
    }
}

使用aspect_hookClass方法来创建hook子类，在aspect_swizzleClassInPlace内通过aspect_swizzleForwardInvocation方法将本类中的forwardInvocation替换成我们自定义的__ASPECTS_ARE_BEING_CALLED__方法，当类没找到消息处理时，通过forwardInvocation转到我们自定义的方法内，就可以统一处理了。

static Class aspect_swizzleClassInPlace(Class klass) {
    NSCParameterAssert(klass);
    NSString *className = NSStringFromClass(klass);
    //保证线程、数据安全
    _aspect_modifySwizzledClasses(^(NSMutableSet *swizzledClasses) {
        if (![swizzledClasses containsObject:className]) {
            // 将原IMP指向forwardInvocation
            aspect_swizzleForwardInvocation(klass);
            [swizzledClasses addObject:className];
        }
    });
    return klass;
}

static NSString *const AspectsForwardInvocationSelectorName = @"__aspects_forwardInvocation:";
static void aspect_swizzleForwardInvocation(Class klass) {
    NSCParameterAssert(klass);
    // 将__aspects_forwardInvocation:指向originalImplementation,
    // 将originalImplementation添加到Method，以便下次调用，直接就可以了    
    IMP originalImplementation = class_replaceMethod(klass, @selector(forwardInvocation:), (IMP)__ASPECTS_ARE_BEING_CALLED__, "v@:@");
    if (originalImplementation) {
        class_addMethod(klass, NSSelectorFromString(AspectsForwardInvocationSelectorName), originalImplementation, "v@:@");
    }
    AspectLog(@"Aspects: %@ is now aspect aware.", NSStringFromClass(klass));
}

5.转发最终的逻辑代码最终转入 ASPECTS_ARE_BEING_CALLED 函数的处理中。

// This is the swizzled forwardInvocation: method.
static void __ASPECTS_ARE_BEING_CALLED__(__unsafe_unretained NSObject *self, SEL selector, NSInvocation *invocation) {
    SEL originalSelector = invocation.selector;
    SEL aliasSelector = aspect_aliasForSelector(invocation.selector);
    invocation.selector = aliasSelector;
    // 本对象的AspectsContainer,添加到对象的aspect
    AspectsContainer *objectContainer = objc_getAssociatedObject(self, aliasSelector);
    // 这个类的AspectsContainer，添加类上面的aspect
    AspectsContainer *classContainer = aspect_getContainerForClass(object_getClass(self), aliasSelector);
    AspectInfo *info = [[AspectInfo alloc] initWithInstance:self invocation:invocation];
    NSArray *aspectsToRemove = nil;

    // Before hooks.
    aspect_invoke(classContainer.beforeAspects, info);
    aspect_invoke(objectContainer.beforeAspects, info);

    // Instead hooks.
    BOOL respondsToAlias = YES;
    if (objectContainer.insteadAspects.count || classContainer.insteadAspects.count) {
        aspect_invoke(classContainer.insteadAspects, info);
        aspect_invoke(objectContainer.insteadAspects, info);
    }else {
        Class klass = object_getClass(invocation.target);
        do {
            if ((respondsToAlias = [klass instancesRespondToSelector:aliasSelector])) {
                [invocation invoke];
                break;
            }
        }while (!respondsToAlias && (klass = class_getSuperclass(klass)));
    }

    // After hooks.
    aspect_invoke(classContainer.afterAspects, info);
    aspect_invoke(objectContainer.afterAspects, info);

    // If no hooks are installed, call original implementation (usually to throw an exception)
    if (!respondsToAlias) {
        invocation.selector = originalSelector;
        SEL originalForwardInvocationSEL = NSSelectorFromString(AspectsForwardInvocationSelectorName);
        if ([self respondsToSelector:originalForwardInvocationSEL]) {
            ((void( *)(id, SEL, NSInvocation *))objc_msgSend)(self, originalForwardInvocationSEL, invocation);
        }else {
            [self doesNotRecognizeSelector:invocation.selector];
        }
    }

    // Remove any hooks that are queued for deregistration.
    [aspectsToRemove makeObjectsPerformSelector:@selector(remove)];
}

首先，使用aliasSelector替换掉传进来invocation的selector，之后将invocation转换成AspectInfo，AspectInfo主要包含了一些参数数组等invocation的信息，为了获取参数数组，Aspect为NSInvocation写了一个分类。之后就转到最终调用的地方了：

- (BOOL)invokeWithInfo:(id)info {
    NSInvocation *blockInvocation = [NSInvocation invocationWithMethodSignature:self.blockSignature];
    NSInvocation *originalInvocation = info.originalInvocation;
    NSUInteger numberOfArguments = self.blockSignature.numberOfArguments;
    //参数错误判断
    if (numberOfArguments > originalInvocation.methodSignature.numberOfArguments) {
        AspectLogError(@"Block has too many arguments. Not calling %@", info);
        return NO;
    }
    // The `self` of the block will be the AspectInfo. Optional.
    if (numberOfArguments > 1) {
        [blockInvocation setArgument:&info atIndex:1];
    }
    
    void *argBuf = NULL;
    for (NSUInteger idx = 2; idx < numberOfArguments; idx++) {
        const char *type = [originalInvocation.methodSignature getArgumentTypeAtIndex:idx];
        NSUInteger argSize;
        NSGetSizeAndAlignment(type, &argSize, NULL);
        
        if (!(argBuf = reallocf(argBuf, argSize))) {
            AspectLogError(@"Failed to allocate memory for block invocation.");
            return NO;
        }
        
        [originalInvocation getArgument:argBuf atIndex:idx];
        [blockInvocation setArgument:argBuf atIndex:idx];
    }
    
    [blockInvocation invokeWithTarget:self.block];
    
    if (argBuf != NULL) {
        free(argBuf);
    }
    return YES;
}

取到原invocation和待替换的block的invocation，判断block的参数如果大于原参数，输出错误信息。之后逐个替换参数，然后执行blockInvocation。到这整个aspect工作就全部执行完了。

疑问

1、为什么aspect_invoke()要用宏定义？
2、

    // 本对象的AspectsContainer,添加到对象的aspect
    AspectsContainer *objectContainer = objc_getAssociatedObject(self, aliasSelector);
    // 这个类的AspectsContainer，添加类上面的aspect
    AspectsContainer *classContainer = aspect_getContainerForClass(object_getClass(self), aliasSelector);

这两句代码是怎么执行的？为什么会返回AspectsContainer？

3、最后invoke函数内，为什么只判断block参数个数大于原函数参数是错误的，bolck参数个数小于原函数参数个数时，忽略原函数内多出来的参数，这样做是为什么？

最后

Aspects的源码陆陆续续一星期才看完，深深感觉到看的时候有些力不从心，大概是runtime基础太薄弱了吧。同时感谢http://wereadteam.github.io/2016/06/30/Aspects/提供的帮助