Hook方法的新姿式--Stinger (使用libffi實現AOP )

前言

    衆所周知,使用runtime的提供的接口,咱們能夠設定原方法的IMP,或交換原方法和目標方法的IMP,以徹底代替原方法的實現,或爲原實現先後至關於加一段額外的代碼。html

@interface ClassA: NSObject
- (void)methodA;
+ (void)methodB;
@end

...

@implementation ClassA (Swizzle)

+ (void)load {
    Method originalMethod = class_getInstanceMethod(self, @selector(methodA));
    Method swizzledMethod = class_getInstanceMethod(self, @selector(swizzled_methodA));
    method_exchangeImplementations(originalMethod, swizzledMethod);
}

- (void)swizzled_methodA {
    ...
    [self swizzled_methodA];
    ...
}

@end

複製代碼

    使用知名的AOP庫 Aspects ,能夠更便捷地爲原方法實現先後增長(代替)額外的執行。git

// hook instance method
[ClassA aspect_hookSelector:@selector(methodA)
                      withOptions:AspectPositionAfter
                       usingBlock:^{...}
                            error:nil];
  
// hook class method
[object_getClass(ClassA) aspect_hookSelector:@selector(methodB)
                      withOptions:AspectPositionAfter
                       usingBlock:^{...}
                            error:nil];
複製代碼

    另外,Aspects 支持屢次hook同一個方法,支持從hook返回的id<AspectToken>對象刪除對應的hook。
    IMP即函數指針,Aspects 的大體原理:替換原方法的IMP消息轉發函數指針 _objc_msgForward_objc_msgForward_stret,把原方法IMP添加並對應到SEL aspects_originalSelector,將forwardInvocation:IMP替換爲參數對齊的C函數__ASPECTS_ARE_BEING_CALLED__(NSObject *self, SEL selector, NSInvocation *invocation)的指針。在__ASPECTS_ARE_BEING_CALLED__函數中,替換invocationselectoraspects_originalSelector,至關於要發送調用原始方法實現的消息。對於插入位置在前面,替換,後面的多個block,構建新的blockInvocation,從invocation中提取參數,最後經過invokeWithTarget:block來完成依次調用。有關消息轉發的介紹,能夠參考筆者的另外一篇文章用代碼理解ObjC中的發送消息和消息轉發
    Aspects 實現代碼裏的不少細節處理是很使人稱道的,且支持hook類的單個實例對象的方法(相似於KVO的isa-swizzlling)。但因爲對原方法調用直接進行了消息轉發,到真正的IMP對應的函數被執行前,經歷了對其餘多個消息的處理,invoke block也須要額外的invocation構建開銷。做者也在註釋中寫道,不適合對每秒鐘超過1000次的方法增長切面代碼。此外,使用其餘方式對Aspect hook過的方法進行hook時,如直接替換爲新的IMP,新hook獲得的原始實現是_objc_msgForward,以前的aspect_hook會失效,新的hook也將執行異常。
    那麼不由要思考,有沒有一種方式能夠替換原方法的IMP爲一個和原方法參數相同(type encoding)的方法的函數指針,做爲殼,處理消息時,在這個殼內部拿到全部參數,最後經過函數指針直接執行「前」、「原始/替換」,「後」的多個代碼塊。使人驚喜的是,libffi 能夠幫咱們作到這一切。github

1. 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。固然,在這個函數裏咱們能夠作一些額外的操做。
數組

1.1 ffi_type

    根據參數個數和參數類型生成的各自的ffi_type。安全

int fun1 (int a, int b) {
    return a + b;
}

int fun2 (int a, int b) {
    return 2 * a + b;
}

...

ffi_type **types;  // 參數類型
types = malloc(sizeof(ffi_type *) * 2) ;
types[0] = &ffi_type_sint;
types[1] = &ffi_type_sint;
ffi_type *retType = &ffi_type_sint;
複製代碼

1.2 ffi_call

     根據ffi_type生成特定cif,輸入cif、 函數指針、參數地址動態調用函數。bash

void **args = malloc(sizeof(void *) * 2);
int x = 1, y = 2;
args[0] = &x;
args[1] = &y;

int ret;

ffi_cif cif; 
// 生成模板
ffi_prep_cif(&cif, FFI_DEFAULT_ABI, 2, retType, types); 
// 動態調用fun1
ffi_call(&cif, fun1,  &ret, args);
...

// 輸出: ret = 3;
複製代碼

1.3 ffi_prep_closure_loc

    生成closure,併產生一個函數指針imp,當執行到imp時,得到全部輸入參數, 後續將執行ffi_function。數據結構

void ffi_function(ffi_cif *cif, void *ret, void **args, void *userdata) {
    ...
    // args爲全部參數的內存地址
}

ffi_cif cif; 
// 生成模板
ffi_prep_cif(&cif, FFI_DEFAULT_ABI, 2, returnType, types);
ffi_prep_closure_loc(_closure, &_cif, ffi_function, (__bridge void *)(self), imp);

void *imp = NULL;
ffi_closure *closure = ffi_closure_alloc(sizeof(ffi_closure), (void **)&imp);

//生成ffi_closure
ffi_prep_closure_loc(closure, &cif, ffi_function, (__bridge void *)(self), stingerIMP);
複製代碼

    libffi 能調用任意 C 函數的原理與objc_msgSend的原理相似,其底層是用匯編實現的,ffi_call根據模板cif和參數值,把參數都按規則塞到棧/寄存器裏,調用的函數能夠按規則獲得參數,調用完再獲取返回值,清理數據。經過其餘方式調用上文中的imp,ffi_closure可根據棧/寄存器、模板cif拿到全部的參數,接着執行自定義函數ffi_function裏的代碼。JPBlock的實現正是利用了後一種方式,更多細節介紹能夠參考 bang: 如何動態調用 C 函數
    到這裏,對於如何hook ObjC方法和實現AOP,想必你們已經有了一些思路,咱們能夠將ffi_closure關聯的指針替換原方法的IMP,當對象收到該方法的消息時objc_msgSend(id self, SEL sel, ...),將最終執行自定義函數void ffi_function(ffi_cif *cif, void *ret, void **args, void *userdata)。而實現這一切的主要工做是:設計可行的結構,存儲類的多個hook信息;根據包含不一樣參數的方法和切面block,生成包含匹配ffi_type的cif;替換類某個方法的實現爲ffi_closure關聯的imp,記錄hook;在ffi_function裏,根據得到的參數,動態調用原始imp和block。
閉包

#import <Foundation/Foundation.h>
#import "StingerParams.h"

typedef NSString *STIdentifier;

typedef NS_ENUM(NSInteger, STOption) {
  STOptionAfter = 0,
  STOptionInstead = 1,
  STOptionBefore = 2,
};

@interface NSObject (Stinger)

+ (BOOL)st_hookInstanceMethod:(SEL)sel option:(STOption)option usingIdentifier:(STIdentifier)identifier withBlock:(id)block;

+ (BOOL)st_hookClassMethod:(SEL)sel option:(STOption)option usingIdentifier:(STIdentifier)identifier withBlock:(id)block;

+ (NSArray<STIdentifier> *)st_allIdentifiersForKey:(SEL)key;

+ (BOOL)st_removeHookWithIdentifier:(STIdentifier)identifier forKey:(SEL)key;

@end
複製代碼

    下文將圍繞一些重要的點來介紹下筆者的實現。Stingerapp

2 方法簽名 & ffi_type

2.1 方法簽名 -> ffi_type

    對於方法的簽名和type encoding,筆者在 用代碼理解ObjC中的發送消息和消息轉發 一文中已經有了很多介紹。簡而言之,type encoding 字符串與方法的返回類型及參數類型是一一對應的。例如:- (void)print1:(NSString *)s;的type encoding爲v24@0:8@16v對應void@對應id(這裏是self),:對應SEL@對應id(這裏是NSString *),另外一方面,每一種參數類型都對應一種ffi_type,如v對應ffi_type_void, @對應ffi_type_pointer。能夠用type encoding生成一個NSMethodSignature實例對象,利用numberOfArguments- (const char *)getArgumentTypeAtIndex:(NSUInteger)idx;方法獲取每個位置上的參數類型。固然,也能夠過濾掉數字來分隔字符串v24@0:8@16(@?爲block),獲得參數類型數組(JSPatch中使用了這一方式)。接着,咱們對字符和ffi_type作一一對應便可完成從方法簽名到ffi_type的轉換。ide

_args = malloc(sizeof(ffi_type *) * argumentCount) ;
for (int i = 0; i < argumentCount; i++) {
  ffi_type* current_ffi_type = ffiTypeWithType(self.signature.argumentTypes[i]);
  NSAssert(current_ffi_type, @"can't find a ffi_type of %@", self.signature.argumentTypes[i]);
  _args[i] = current_ffi_type;
}
複製代碼

2.2 淺談block

2.2.1 簽名 & 函數指針

void (^block)(id<StingerParams> params, NSString *s) = ^(id<StingerParams> params, NSString *s) {
    NSLog(@"---after2 print1: %@", s);
}
複製代碼

    block是一個ObjC對象,能夠認爲幾種block類型都繼承於NSBlock。block很特殊,從表面來看包含了持有了數據和對象(暫不討論變量捕獲),並擁有可執行的代碼,調用方式相似於調用C函數,等同於數據加函數。Block類型很神祕,但咱們從 opensource-apple/objc4oclang/docs/block 中看到Block 完整的數據結構。

enum {
  BLOCK_DEALLOCATING =      (0x0001),  // runtime
  BLOCK_REFCOUNT_MASK =     (0xfffe),  // runtime
  BLOCK_NEEDS_FREE =        (1 << 24), // runtime
  BLOCK_HAS_COPY_DISPOSE =  (1 << 25), // compiler
  BLOCK_HAS_CTOR =          (1 << 26), // compiler: helpers have C++ code
  BLOCK_IS_GC =             (1 << 27), // runtime
  BLOCK_IS_GLOBAL =         (1 << 28), // compiler
  BLOCK_USE_STRET =         (1 << 29), // compiler: undefined if !BLOCK_HAS_SIGNATURE
  BLOCK_HAS_SIGNATURE  =    (1 << 30)  // compiler
};

// revised new layout

#define BLOCK_DESCRIPTOR_1 1
struct Block_descriptor_1 {
  unsigned long int reserved;
  unsigned long int size;
};

#define BLOCK_DESCRIPTOR_2 1
struct Block_descriptor_2 {
  // requires BLOCK_HAS_COPY_DISPOSE
  void (*copy)(void *dst, const void *src);
  void (*dispose)(const void *);
};

#define BLOCK_DESCRIPTOR_3 1
struct Block_descriptor_3 {
  // requires BLOCK_HAS_SIGNATURE
  const char *signature;
  const char *layout;
};

struct Block_layout {
  void *isa;
  volatile int flags; // contains ref count
  int reserved;
  void (*invoke)(void *, ...);
  struct Block_descriptor_1 *descriptor;
  // imported variables
};
複製代碼

    不少人大概已經看過BlocksKit的代碼,瞭解到Block對象能夠強轉爲Block_layout類型,經過標識符和內存地址偏移獲取block的簽名signature

NSString *signatureForBlock(id block) {
  struct Block_layout *layout = (__bridge void *)block;
  if (!(layout->flags & BLOCK_HAS_SIGNATURE))
    return nil;
  
  void *descRef = layout->descriptor;
  descRef += 2 * sizeof(unsigned long int);
  
  if (layout->flags & BLOCK_HAS_COPY_DISPOSE)
    descRef += 2 * sizeof(void *);
  
  if (!descRef)
    return nil;
  
  const char *signature = (*(const char **)descRef);
  return [NSString stringWithUTF8String:signature];
}

複製代碼
NSString *signature = signatureForBlock(block)
// 輸出 NSString:@"v24@?0@\"<StingerParams>\"8@\"NSString\"16"
複製代碼

    對於Block對象的的最簡簽名,咱們仍然能夠構建NSMethodSignature來逐一獲取,也能夠經過過濾掉數字及'\"'來得到字符數組。

_argumentTypes = [[NSMutableArray alloc] init];
NSInteger descNum = 0; // num of '\"' in block signature type encoding
for (int i = 0; i < _types.length; i ++) {
    unichar c = [_types characterAtIndex:i];
    NSString *arg;
    if (c == '\"') ++descNum;
    if ((descNum % 2) != 0 || (c == '\"' || isdigit(c))) {
      continue;
 }
    ...
}  
/*@"v24@?0@\"<StingerParams>\"8@\"NSString\"16"
*/ -> v,@?,@,@
複製代碼

    能夠看到,簽名的第一位是"@?",意味着第一個參數爲blcok本身,後面的纔是blcok的參數類型。同理,咱們依然能夠經過type encoding匹配到對應的ffi_type
    此外,咱們能夠直接獲取到Block對象的函數指針。

BlockIMP impForBlock(id block) {
  struct Block_layout *layout = (__bridge void *)block;
  return layout->invoke;
}
複製代碼

    作一個簡單的嘗試,直接調用Block對象的包含的函數。

void (^block2)(NSString *s) = ^(NSString *s) {
    NSLog(@"---after2 print1: %@", s);
  };
  void (*blockIMP) (id block, NSString *s) = (void (*) (id block, NSString *s))impForBlock(block2);
  blockIMP(block2, @"tt");
  
  // 輸出:---after2 print1: tt
複製代碼

此外,實測經過IMP _Nonnull imp_implementationWithBlock(id _Nonnull block)得到的函數指針對應的參數並不包含Block對象自身,意味着簽名發生了變化。

* 爲block對象增長可用方法

    經過一些方式,咱們能夠以爲Block對象擁有了新的實例方法。

NSString *signature = [block signature];
void *blockIMP = [block blockIMP];
複製代碼

    作法是在STBlock裏爲NSBlock類增長實例方法。

typedef void *BlockIMP;

@interface STBlock : NSObject
+ (instancetype)new NS_UNAVAILABLE;
- (instancetype)init NS_UNAVAILABLE;

- (NSString *)signature;
- (BlockIMP)blockIMP;

NSString *signatureForBlock(id block);
BlockIMP impForBlock(id block);
@end
複製代碼
#define NSBlock NSClassFromString(@"NSBlock")

void addInstanceMethodForBlock(SEL sel) {
  Method m = class_getInstanceMethod(STBlock.class, sel);
  if (!m) return;
  IMP imp = method_getImplementation(m);
  const char *typeEncoding = method_getTypeEncoding(m);
  class_addMethod(NSBlock, sel, imp, typeEncoding);
}

@implementation STBlock

+ (void)load {
  addInstanceMethodForBlock(@selector(signature));
  addInstanceMethodForBlock(@selector(blockIMP));
}

...

@end
複製代碼

這樣作,能夠爲Block對象增長可處理的消息。但若是在其餘類的load方法裏嘗試調用,可能會遇到STBlock類裏load方法未加載的問題。

3 存儲hook信息 & 生成兩個ffi_cif對象

3.1 StingerInfo

    這裏使用簡單的對象來存儲單個hook信息。

@protocol StingerInfo <NSObject>
@required
@property (nonatomic, copy) id block;
@property (nonatomic, assign) STOption option;
@property (nonatomic, copy) STIdentifier identifier;

@optional
+ (instancetype)infoWithOption:(STOption)option withIdentifier:(STIdentifier)identifier withBlock:(id)block;
@end

@interface StingerInfo : NSObject <StingerInfo>
@end
複製代碼

3.2 StingerInfoPool

typedef void *StingerIMP;

@protocol StingerInfoPool <NSObject>

@required
@property (nonatomic, strong, readonly) NSMutableArray<id<StingerInfo>> *beforeInfos;
@property (nonatomic, strong, readonly) NSMutableArray<id<StingerInfo>> *insteadInfos;
@property (nonatomic, strong, readonly) NSMutableArray<id<StingerInfo>> *afterInfos;
@property (nonatomic, strong, readonly) NSMutableArray<NSString *> *identifiers;

@property (nonatomic, copy) NSString *typeEncoding;
@property (nonatomic) IMP originalIMP;
@property (nonatomic) SEL sel;

- (StingerIMP)stingerIMP;
- (BOOL)addInfo:(id<StingerInfo>)info;
- (BOOL)removeInfoForIdentifier:(STIdentifier)identifier;

@optional
@property (nonatomic, weak) Class cls;
+ (instancetype)poolWithTypeEncoding:(NSString *)typeEncoding originalIMP:(IMP)imp selector:(SEL)sel;
@end

@interface StingerInfoPool : NSObject <StingerInfoPool>
@end
複製代碼

3.2.1 管理StingerInfo

    這裏利用三個數組來存儲某個類hook位置在原實現前、替換、實現後的id<StingerInfo>對象,並保存了原始imp。添加和刪除id<StingerInfo>對象的操做是線程安全的。

3.2.2 生成方法調用模板 cif

    根據原始方法提供的type encoding,生成各個參數對應的ffi_type,繼而生成cif對象,最後調用ffi_prep_closure_loc至關於生成空殼函數StingerIMP。調用StingerIMP將最終執行到自定義的static void ffi_function(ffi_cif *cif, void *ret, void **args, void *userdata)函數,此函數可得到調用StingerIMP時得到的全部參數。

- (StingerIMP)stingerIMP {
  ffi_type *returnType = ffiTypeWithType(self.signature.returnType);
  NSAssert(returnType, @"can't find a ffi_type of %@", self.signature.returnType);
  
  NSUInteger argumentCount = self.signature.argumentTypes.count;
  StingerIMP stingerIMP = NULL;
  _args = malloc(sizeof(ffi_type *) * argumentCount) ;
  
  for (int i = 0; i < argumentCount; i++) {
    ffi_type* current_ffi_type = ffiTypeWithType(self.signature.argumentTypes[i]);
    NSAssert(current_ffi_type, @"can't find a ffi_type of %@", self.signature.argumentTypes[i]);
    _args[i] = current_ffi_type;
  }
  
  _closure = ffi_closure_alloc(sizeof(ffi_closure), (void **)&stingerIMP);
  
  if(ffi_prep_cif(&_cif, FFI_DEFAULT_ABI, (unsigned int)argumentCount, returnType, _args) == FFI_OK) {
    if (ffi_prep_closure_loc(_closure, &_cif, ffi_function, (__bridge void *)(self), stingerIMP) != FFI_OK) {
      NSAssert(NO, @"genarate IMP failed");
    }
  } else {
    NSAssert(NO, @"FUCK");
  }
  
  [self _genarateBlockCif];
  return stingerIMP;
}
複製代碼

3.2.3 生成block調用模板 blockCif

    與前面生成方法調用模板cif相似,只不過這裏沒有生成殼子ffi_closure。值得注意的是,這裏把原始方法type encoing的第0位(@ self)和第1位(: SEL)替換爲(@?block)和(@ id<StingerParams>)。意味着,限定了切面Block對象的簽名類型。

- (void)_genarateBlockCif {
  ffi_type *returnType = ffiTypeWithType(self.signature.returnType);
  
  NSUInteger argumentCount = self.signature.argumentTypes.count;
  _blockArgs = malloc(sizeof(ffi_type *) *argumentCount);
  
  ffi_type *current_ffi_type_0 = ffiTypeWithType(@"@?");
  _blockArgs[0] = current_ffi_type_0;
  ffi_type *current_ffi_type_1 = ffiTypeWithType(@"@");
  _blockArgs[1] = current_ffi_type_1;
  
  for (int i = 2; i < argumentCount; i++){
    ffi_type* current_ffi_type = ffiTypeWithType(self.signature.argumentTypes[i]);
    _blockArgs[i] = current_ffi_type;
  }
  
  if(ffi_prep_cif(&_blockCif, FFI_DEFAULT_ABI, (unsigned int)argumentCount, returnType, _blockArgs) != FFI_OK) {
    NSAssert(NO, @"FUCK");
  }
}
複製代碼

在非instead位置,block的返回值能夠爲任意;寫block時,block的第0位(不考慮block自身)參數類型應該爲id,後面接的是與原方法對應的參數。

3.2.4 ffi_function !!!

    在這個函數裏,獲取到了調用原始方法時的全部入參的內存地址,先是根據block_cif模板生成新的參數集innerArgs,第0位留給Block對象,第1位留給StingerParams對象,從第2位開始複製原始的參數。
    如下是完成切面代碼和原始imp執行的過程:
    1. 利用ffi_call(&(self->_blockCif), impForBlock(block), NULL, innerArgs);完成全部切面位置在前block的調用。使用block模板blockCif和innerArgs。
    2. 利用ffi_call(cif, (void (*)(void))self.originalIMP / impForBlock(block), ret, args);完成對原始IMP或替換位置block imp的調用。使用原始模板cif和原始參數args,並可能產生返回值。
    3. 利用ffi_call(&(self->_blockCif), impForBlock(block), NULL, innerArgs);完成全部切面位置在後的block的調用。使用block模板blockCif和innerArgs。

static void ffi_function(ffi_cif *cif, void *ret, void **args, void *userdata) {
  StingerInfoPool *self = (__bridge StingerInfoPool *)userdata;
  NSUInteger count = self.signature.argumentTypes.count;
  void **innerArgs = malloc(count * sizeof(*innerArgs));
  
  StingerParams *params = [[StingerParams alloc] init];
  void **slf = args[0];
  params.slf = (__bridge id)(*slf);
  params.sel = self.sel;
  [params addOriginalIMP:self.originalIMP];
  NSInvocation *originalInvocation = [NSInvocation invocationWithMethodSignature:self.ns_signature];
  for (int i = 0; i < count; i ++) {
    [originalInvocation setArgument:args[i] atIndex:i];
  }
  [params addOriginalInvocation:originalInvocation];
  
  innerArgs[1] = &params;
  memcpy(innerArgs + 2, args + 2, (count - 2) * sizeof(*args));
  
  #define ffi_call_infos(infos) \
    for (id<StingerInfo> info in infos) { \
    id block = info.block; \
    innerArgs[0] = &block; \
    ffi_call(&(self->_blockCif), impForBlock(block), NULL, innerArgs); \
  }  \
  // before hooks
  ffi_call_infos(self.beforeInfos);
  // instead hooks
  if (self.insteadInfos.count) {
    id <StingerInfo> info = self.insteadInfos[0];
    id block = info.block;
    innerArgs[0] = &block;
    ffi_call(&(self->_blockCif), impForBlock(block), ret, innerArgs);
  } else {
    // original IMP
    ffi_call(cif, (void (*)(void))self.originalIMP, ret, args);
  }
  // after hooks
  ffi_call_infos(self.afterInfos);
  
  free(innerArgs);
}
複製代碼

注:StingerParams 對象包含了消息接收者slf,當前消息的selector sel, 還包含了可調用原始方法的invocation(使用invokeUsingIMP:完成調用),該invocation僅適合在替換方法且須要原始返回值做參數時調用。其餘hook直接使用optionBefore或after便可, 不用關注該invocation。

#import <Foundation/Foundation.h>

#define ST_NO_RET NULL

@protocol StingerParams
@required
@property (nonatomic, unsafe_unretained) id slf;
@property (nonatomic) SEL sel;

- (void)invokeAndGetOriginalRetValue:(void *)retLoc;
@end

@interface StingerParams : NSObject <StingerParams>
- (void)addOriginalInvocation:(NSInvocation *)invocation;
- (void)addOriginalIMP:(IMP)imp;
@end

複製代碼

4 替換方法實現 & 記錄HOOK

    思路是對某個類以SEL sel爲鍵關聯一個id<StingerInfoPool>對象,第一次hook,新建該對象,嘗試替換原方法實現爲ffi_prep_closure_loc關聯的IMP,後續hook時,將直接添加hook info到關聯的id<StingerInfoPool>對象中。
    關於條件,最主要的就是兩點,第一點就是對於某個類中(父類)的某個SEL sel要能找到對應Method m及IMP imp;第二點即切面block與原方法的簽名是匹配的,且切面block的簽名是符合要求的(isMatched方法)。

#import "Stinger.h"
#import <objc/runtime.h>
#import "StingerInfo.h"
#import "StingerInfoPool.h"
#import "STBlock.h"
#import "STMethodSignature.h"

@implementation NSObject (Stinger)

#pragma - public

+ (BOOL)st_hookInstanceMethod:(SEL)sel option:(STOption)option usingIdentifier:(STIdentifier)identifier withBlock:(id)block {
  return hook(self, sel, option, identifier, block);
}

+ (BOOL)st_hookClassMethod:(SEL)sel option:(STOption)option usingIdentifier:(STIdentifier)identifier withBlock:(id)block {
  return hook(object_getClass(self), sel, option, identifier, block);
}

+ (NSArray<STIdentifier> *)st_allIdentifiersForKey:(SEL)key {
  NSMutableArray *mArray = [[NSMutableArray alloc] init];
  @synchronized(self) {
    [mArray addObjectsFromArray:getAllIdentifiers(self, key)];
    [mArray addObjectsFromArray:getAllIdentifiers(object_getClass(self), key)];
  }
  return [mArray copy];
}

+ (BOOL)st_removeHookWithIdentifier:(STIdentifier)identifier forKey:(SEL)key {
  BOOL hasRemoved = NO;
  @synchronized(self) {
    id<StingerInfoPool> infoPool = getStingerInfoPool(self, key);
    if ([infoPool removeInfoForIdentifier:identifier]) {
      hasRemoved = YES;
    }
    infoPool = getStingerInfoPool(object_getClass(self), key);
    if ([infoPool removeInfoForIdentifier:identifier]) {
      hasRemoved = YES;
    }
  }
  return hasRemoved;
}

#pragma - inline functions

NS_INLINE BOOL hook(Class cls, SEL sel, STOption option, STIdentifier identifier, id block) {
  NSCParameterAssert(cls);
  NSCParameterAssert(sel);
  NSCParameterAssert(option == 0 || option == 1 || option == 2);
  NSCParameterAssert(identifier);
  NSCParameterAssert(block);
  Method m = class_getInstanceMethod(cls, sel);
  NSCAssert(m, @"SEL (%@) doesn't has a imp in Class (%@) originally", NSStringFromSelector(sel), cls);
  if (!m) return NO;
  const char * typeEncoding = method_getTypeEncoding(m);
  STMethodSignature *methodSignature = [[STMethodSignature alloc] initWithObjCTypes:[NSString stringWithUTF8String:typeEncoding]];
  STMethodSignature *blockSignature = [[STMethodSignature alloc] initWithObjCTypes:signatureForBlock(block)];
  if (! isMatched(methodSignature, blockSignature, option, cls, sel, identifier)) {
    return NO;
  }

  IMP originalImp = method_getImplementation(m);
  
  @synchronized(cls) {
    StingerInfo *info = [StingerInfo infoWithOption:option withIdentifier:identifier withBlock:block];
    id<StingerInfoPool> infoPool = getStingerInfoPool(cls, sel);
    
    if (infoPool) {
      return [infoPool addInfo:info];
    }
    
    infoPool = [StingerInfoPool poolWithTypeEncoding:[NSString stringWithUTF8String:typeEncoding] originalIMP:originalImp selector:sel];
    infoPool.cls = cls;
    
    IMP stingerIMP = [infoPool stingerIMP];
    
    if (!(class_addMethod(cls, sel, stingerIMP, typeEncoding))) {
      class_replaceMethod(cls, sel, stingerIMP, typeEncoding);
    }
    const char * st_original_SelName = [[@"st_original_" stringByAppendingString:NSStringFromSelector(sel)] UTF8String];
    class_addMethod(cls, sel_registerName(st_original_SelName), originalImp, typeEncoding);
    
    setStingerInfoPool(cls, sel, infoPool);
    return [infoPool addInfo:info];
  }
}

NS_INLINE id<StingerInfoPool> getStingerInfoPool(Class cls, SEL key) {
  NSCParameterAssert(cls);
  NSCParameterAssert(key);
  return objc_getAssociatedObject(cls, key);
}

NS_INLINE void setStingerInfoPool(Class cls, SEL key, id<StingerInfoPool> infoPool) {
  NSCParameterAssert(cls);
  NSCParameterAssert(key);
  objc_setAssociatedObject(cls, key, infoPool, OBJC_ASSOCIATION_RETAIN);
}

NS_INLINE NSArray<STIdentifier> * getAllIdentifiers(Class cls, SEL key) {
  NSCParameterAssert(cls);
  NSCParameterAssert(key);
  id<StingerInfoPool> infoPool = getStingerInfoPool(cls, key);
  return infoPool.identifiers;
}


NS_INLINE BOOL isMatched(STMethodSignature *methodSignature, STMethodSignature *blockSignature, STOption option, Class cls, SEL sel, NSString *identifier) {
  //argument count
  if (methodSignature.argumentTypes.count != blockSignature.argumentTypes.count) {
    NSCAssert(NO, @"count of arguments isn't equal. Class: (%@), SEL: (%@), Identifier: (%@)", cls, NSStringFromSelector(sel), identifier);
    return NO;
  };
  // loc 1 should be id<StingerParams>.
  if (![blockSignature.argumentTypes[1] isEqualToString:@"@"]) {
     NSCAssert(NO, @"argument 1 should be object type. Class: (%@), SEL: (%@), Identifier: (%@)", cls, NSStringFromSelector(sel), identifier);
    return NO;
  }
  // from loc 2.
  for (NSInteger i = 2; i < methodSignature.argumentTypes.count; i++) {
    if (![blockSignature.argumentTypes[i] isEqualToString:methodSignature.argumentTypes[i]]) {
      NSCAssert(NO, @"argument (%zd) type isn't equal. Class: (%@), SEL: (%@), Identifier: (%@)", i, cls, NSStringFromSelector(sel), identifier);
      return NO;
    }
  }
  // when STOptionInstead, returnType
  if (option == STOptionInstead && ![blockSignature.returnType isEqualToString:methodSignature.returnType]) {
    NSCAssert(NO, @"return type isn't equal. Class: (%@), SEL: (%@), Identifier: (%@)", cls, NSStringFromSelector(sel), identifier);
    return NO;
  }
  
  return YES;
}

@end
複製代碼

* 使用示例

import UIKit;

@interface ASViewController : UIViewController

- (void)print1:(NSString *)s;

- (NSString *)print2:(NSString *)s;

@end
複製代碼
#import "ASViewController+hook.h"

@implementation ASViewController (hook)

+ (void)load {
  /*
   * hook @selector(print1:)
   */
  [self st_hookInstanceMethod:@selector(print1:) option:STOptionBefore usingIdentifier:@"hook_print1_before1" withBlock:^(id<StingerParams> params, NSString *s) {
    NSLog(@"---before1 print1: %@", s);
  }];
  
  [self st_hookInstanceMethod:@selector(print1:) option:STOptionBefore usingIdentifier:@"hook_print1_before2" withBlock:^(id<StingerParams> params, NSString *s) {
    NSLog(@"---before2 print1: %@", s);
  }];
  
  [self st_hookInstanceMethod:@selector(print1:) option:STOptionAfter usingIdentifier:@"hook_print1_after1" withBlock:^(id<StingerParams> params, NSString *s) {
    NSLog(@"---after1 print1: %@", s);
  }];
  
  [self st_hookInstanceMethod:@selector(print1:) option:STOptionAfter usingIdentifier:@"hook_print1_after2" withBlock:^(id<StingerParams> params, NSString *s) {
    NSLog(@"---after2 print1: %@", s);
  }];
  
  /*
   * hook @selector(print2:)
   */
  __block NSString *oldRet, *newRet;
  [self st_hookInstanceMethod:@selector(print2:) option:STOptionInstead usingIdentifier:@"hook_print2_instead" withBlock:^NSString * (id<StingerParams> params, NSString *s) {
    [params invokeAndGetOriginalRetValue:&oldRet];
    newRet = [oldRet stringByAppendingString:@" ++ new-st_instead"];
    NSLog(@"---instead print2 old ret: (%@) / new ret: (%@)", oldRet, newRet);
    return newRet;
  }];
  
  [self st_hookInstanceMethod:@selector(print2:) option:STOptionAfter usingIdentifier:@"hook_print2_after1" withBlock:^(id<StingerParams> params, NSString *s) {
    NSLog(@"---after1 print2 self:%@ SEL: %@ p: %@",[params slf], NSStringFromSelector([params sel]), s);
  }];
}
@end

複製代碼

Stinger用法與Aspects很類似,但收到消息後,因爲block和原始IMP直接使用函數指針進行調用,不處理額外的消息,不用實例化諸多NSInvocation對象,兩個lib_cif對象在hook後也即準備好,相比aspects,實測從收到消息到執行到切面信息(例如before block)這個過程, 花費時間下降一個數量級。使用其餘方式hook時,仍能保證st_hook的有效性。

謝謝觀看,水平有限,若有錯誤,請指正。

github.com/Assuner-Lee…

參考資料

github.com/opensource-…
blog.cnbang.net/tech/3219/
juejin.im/post/5a308f…
github.com/mikeash/MAB…

相關文章
相關標籤/搜索