开源库 Block Tracker 学习

修改 selector IMP 映射来 hook 方法在开发中很常见,但是 hook 一个 block 实现以及使用场景都较为稀有。最近,腾讯星开源了一个 hook block 的方案 Block Tracker , 使用上看起来如下:

- (void)viewDidLoad {
    [super viewDidLoad];
    // Begin Track
    BTTracker *tracker = [self bt_trackBlockArgOfSelector:@selector(performBlock:) callback:^(id  _Nullable block, BlockTrackerCallbackType type, NSInteger invokeCount, void * _Nullable * _Null_unspecified args, void * _Nullable result, NSArray<NSString *> * _Nonnull callStackSymbols) {
        NSLog(@"%@ invoke count = %ld", BlockTrackerCallbackTypeInvoke == type ? @"BlockTrackerCallbackTypeInvoke" : @"BlockTrackerCallbackTypeDead", (long)invokeCount);
    }];
    // invoke blocks
    __block NSString *word = @"I'm a block";
    [self performBlock:^{
        NSLog(@"add '!!!' to word");
        word = [word stringByAppendingString:@"!!!"];
    }];
    [self performBlock:^{
        NSLog(@"%@", word);
    }];
}

- (void)performBlock:(void(^)(void))block {
    block();
}

Console log:

add '!!!' to word
BlockTrackerCallbackTypeInvoke invoke count = 1
I'm a block!!!
BlockTrackerCallbackTypeInvoke invoke count = 1
BlockTrackerCallbackTypeDead invoke count = 1
BlockTrackerCallbackTypeDead invoke count = 1

一开始对于实现上有几点疑问:

  1. block 做为方法的一个参数,如何跟 BTTracker 对象进行关联以及管理其生命周期的?
  2. block 的内存结构已经很熟悉,但是 fake 指定 block 时如何构造与原 block 一样的参数以及返回值的 invoke 函数指针呢?

BTTracker

上文展示过,当我们在 hook 方法中的 block 时,调用的是下面这个方法

- (nullable BTTracker *)bt_trackBlockArgOfSelector:(SEL)selector callback:(BlockTrackerCallbackBlock)callback;

它的实现:

- (nullable BTTracker *)bt_trackBlockArgOfSelector:(SEL)selector callback:(BlockTrackerCallbackBlock)callback
{
    Class cls = bt_classOfTarget(self);
    
    Method originMethod = class_getInstanceMethod(cls, selector);
    if (!originMethod) {
        return nil;
    }
    const char *originType = (char *)method_getTypeEncoding(originMethod);
    if (![[NSString stringWithUTF8String:originType] containsString:@"@?"]) {
        return nil;
    }
    NSMutableArray *blockArgIndex = [NSMutableArray array];
    int argIndex = 0; // return type is the first one
    while(originType && *originType)
    {
        originType = BHSizeAndAlignment(originType, NULL, NULL, NULL);
        if ([[NSString stringWithUTF8String:originType] hasPrefix:@"@?"]) {
            [blockArgIndex addObject:@(argIndex)];
        }
        argIndex++;
    }

    BTTracker *tracker = BTEngine.defaultEngine.trackers[bt_methodDescription(self, selector)];
    if (!tracker) {
        tracker = [[BTTracker alloc] initWithTarget:self selector:selector];
        tracker.callback = callback;
        tracker.blockArgIndex = [blockArgIndex copy];
    }
    return [tracker apply] ? tracker : nil;
}

先拿到这个方法的 type encoding,判断参数中是否有 block 如果没有直接返回 nil 。
然后记录参数列表中 block 所在的 index , 并查询 BTEngine 这个 tracker 的管理类中是否有该方法的缓存 tracker 对象。
如果没有则创建新的 tracker 对象,并关联 target, selector, callback 以及刚才记录的 index 。
最后调用 apply 方法并返回 tracker 对象。

于是我们大概能猜测到, apply 方法大概是让 BTEngine 对该 tracker 对象进行管理及其他处理, 看下其实现:

- (BOOL)applyTracker:(BTTracker *)tracker
{
    pthread_mutex_lock(&mutex);
    __block BOOL shouldApply = YES;
    if (bt_checkTrackerValid(tracker)) {
        [self.trackers enumerateKeysAndObjectsUsingBlock:^(NSString * _Nonnull key, BTTracker * _Nonnull obj, BOOL * _Nonnull stop) {
            if (sel_isEqual(tracker.selector, obj.selector)) {
                
                Class clsA = bt_classOfTarget(tracker.target);
                Class clsB = bt_classOfTarget(obj.target);
                
                shouldApply = !([clsA isSubclassOfClass:clsB] || [clsB isSubclassOfClass:clsA]);
                *stop = shouldApply;
                NSCAssert(shouldApply, @"Error: %@ already apply tracker in %@. A message can only have one tracker per class hierarchy.", NSStringFromSelector(obj.selector), NSStringFromClass(clsB));
            }
        }];
        
        if (shouldApply) {
            self.trackers[bt_methodDescription(tracker.target, tracker.selector)] = tracker;
            bt_overrideMethod(tracker.target, tracker.selector);
            bt_configureTargetDealloc(tracker);
        }
    }
    else {
        shouldApply = NO;
    }
    pthread_mutex_unlock(&mutex);
    return shouldApply;
}

全局的互斥锁 mutex 保证线程安全。
bt_checkTrackerValid 方法对提交的 checker 进行了验证,如果关联的 selector 是消息转发或者 target 是该库 BTTracker 类或者 BTEngine 类则验证不通过。
接下来对已经生成的 tracker 进行比对,如果有 track 的对象跟新对象处于继承关系则返回 NO,如果同时 hook 了父子的相同方法,子类调用父类的实现,就会死循环。
最后如果满足 apply 条件,则重写指定 selector 并配置析构注入。

重写时使用了消息转发,将该方法调用消息指定到内部的 handle 函数中进行处理:

static void bt_handleInvocation(NSInvocation *invocation, SEL fixedSelector)
{
    NSString *methodDescriptionForInstance = bt_methodDescription(invocation.target, invocation.selector);
    NSString *methodDescriptionForClass = bt_methodDescription(object_getClass(invocation.target), invocation.selector);
    
    BTTracker *tracker = BTEngine.defaultEngine.trackers[methodDescriptionForInstance];
    if (!tracker) {
        tracker = BTEngine.defaultEngine.trackers[methodDescriptionForClass];
    }
    
    [invocation retainArguments];
    
    for (NSNumber *index in tracker.blockArgIndex) {
        if (index.integerValue < invocation.methodSignature.numberOfArguments) {
            __unsafe_unretained id block;
            [invocation getArgument:&block atIndex:index.integerValue];
            __weak typeof(block) weakBlock = block;
            __weak typeof(tracker) weakTracker = tracker;
            BHToken *tokenAfter = [block block_hookWithMode:BlockHookModeAfter usingBlock:^(BHToken *token) {
                __strong typeof(weakBlock) strongBlock = weakBlock;
                __strong typeof(weakTracker) strongTracker = weakTracker;
                NSNumber *invokeCount = objc_getAssociatedObject(token, NSSelectorFromString(@"invokeCount"));
                if (!invokeCount) {
                    invokeCount = @(1);
                }
                else {
                    invokeCount = [NSNumber numberWithInt:invokeCount.intValue + 1];
                }
                objc_setAssociatedObject(token, NSSelectorFromString(@"invokeCount"), invokeCount, OBJC_ASSOCIATION_RETAIN);
                if (strongTracker.callback) {
                    strongTracker.callback(strongBlock, BlockTrackerCallbackTypeInvoke, invokeCount.intValue, token.args, token.retValue, [NSThread callStackSymbols]);
                }
            }];

            [block block_hookWithMode:BlockHookModeDead usingBlock:^(BHToken *token) {
                __strong typeof(weakTracker) strongTracker = weakTracker;
                NSNumber *invokeCount = objc_getAssociatedObject(tokenAfter, NSSelectorFromString(@"invokeCount"));
                if (strongTracker.callback) {
                    strongTracker.callback(nil, BlockTrackerCallbackTypeDead, invokeCount.intValue, nil, nil, [NSThread callStackSymbols]);
                }
            }];
        }
    }
    invocation.selector = fixedSelector;
    [invocation invoke];
}

在这一步,才真正对要执行的 block 进行了 hook ,关注 block 执行后以及销毁两个节点,然后才执行 block 。

BlockHook

hook block 的逻辑他单独封装在了 BlockHook 这个库中,那么他是如何做到 fake 一个相同的 block 呢,来看下 _BHBlock 以及 _BHBlockDescriptor 的构造:

struct _BHBlockDescriptor
{
    unsigned long reserved;
    unsigned long size;
    void *rest[1];
};

struct _BHBlock
{
    void *isa;
    int flags;
    int reserved;
    void *invoke;
    struct _BHBlockDescriptor *descriptor;
};

这点上,跟苹果对 block 的构造是一样的,接下来对于替换 invoke 则使用到了 libffi 库,来看下 BHToken 构造时做的事情:

- (id)initWithBlock:(id)block
{
    if((self = [self init]))
    {
        _allocations = [[NSMutableArray alloc] init];
        _block = block;
        _closure = ffi_closure_alloc(sizeof(ffi_closure), &_replacementInvoke);
        _numberOfArguments = [self _prepCIF:&_cif withEncodeString:BHBlockTypeEncodeString(_block)];
        BHDealloc *bhDealloc = [BHDealloc new];
        bhDealloc.token = self;
        objc_setAssociatedObject(block, NSSelectorFromString([NSString stringWithFormat:@"%p", self]), bhDealloc, OBJC_ASSOCIATION_RETAIN);
        [self _prepClosure];
    }
    return self;
}

关于 libffi:

libffi 可以认为是实现了C语言上的 runtime,简单来说,libffi 可根据 参数类型 (ffi_type) ,参数个数生成一个模板 (ffi_cif) ;可以输入 模板、函数指针和参数地址来直接完成函数调用 (ffi_call) ;模板也可以生成一个所谓的闭包 (ffi_closure) ,并得到指针,当执行到这个地址时,会执行到自定义的 void function(ffi_cif *cif, void *ret, void **args, void *userdata) 函数,在这里,我们可以获得所有参数的地址(包括返回值),以及自定义数据 userdata。当然,在这个函数里我们可以做一些额外的操作。

则这里较为关键的两个步骤: 构建 invoke 函数的模板 cif, 以及替换 invoke 。

- (int)_prepCIF:(ffi_cif *)cif withEncodeString:(const char *)str
{
    int argCount;
    ffi_type **argTypes = [self _argsWithEncodeString:str getCount:&argCount];
    
    ffi_status status = ffi_prep_cif(cif, FFI_DEFAULT_ABI, argCount, [self _ffiArgForEncode: str], argTypes);
    if(status != FFI_OK)
    {
        NSLog(@"Got result %ld from ffi_prep_cif", (long)status);
        abort();
    }
    return argCount;
}

- (void)_prepClosure
{
    ffi_status status = ffi_prep_closure_loc(_closure, &_cif, BHFFIClosureFunc, (__bridge void *)(self), _replacementInvoke);
    if(status != FFI_OK)
    {
        NSLog(@"ffi_prep_closure returned %d", (int)status);
        abort();
    }
    // exchange invoke func imp
    _originInvoke = ((__bridge struct _BHBlock *)self.block)->invoke;
    ((__bridge struct _BHBlock *)self.block)->invoke = _replacementInvoke;
}

- (void)invokeOriginalBlock
{
    if (_originInvoke) {
        ffi_call(&_cif, _originInvoke, self.retValue, self.args);
    }
    else {
        NSLog(@"You had lost your originInvoke! Please check the order of removing tokens!");
    }
}

其实,看到这,整个 track block 的过程从上层 api 实现到底层使用 libffi 动态构建函数以及调用的过程都已经浏览了一遍,其他的一些流程就无需过多阐述了。

其他通过代码,也能学到一些有意思的细节, 比如对代码签名,type encode 的使用及处理,比如对消息转发, invocation 的使用等等。

-- EOF --
以上就是这篇文章全部内容,欢迎提出建议和指正,个人联系方式详见关于

Comments
Write a Comment
  • Cocoarannie reply

    TEST--

    • Ran reply

      @Cocoarannie