Use libffi to implement AOP

preface

As we all know, using the interface provided by The Runtime, we can set the IMPs of the original method, or swap the IMPs of the original method and the target method to completely replace the original method implementation, or add an extra piece of code to the original implementation.

@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]; . } @endCopy the code

Using well-known AOP library Aspects makes it easier to add (replace) additional implementations before and after the original method implementation.

// 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];Copy the code

In addition, Aspects supports multiple hook methods, allowing the corresponding hook to be deleted from the ID

object returned from the hook. IMP is the general principle of function Pointers and Aspects: Replace the IMP of the original method with the message forwarding function pointer _objc_msgForward or _objc_msgForward_stret, add the original method IMP and match it to SEL aspects_originalSelector, Replace the IMP from the forwardInvocation with the argument alignment C function __ASPECTS_ARE_BEING_CALLED__(NSObject *self, SEL selector, NSInvocation * Invocation). In the __ASPECTS_ARE_BEING_CALLED__ function, the replacement selector is aspects_originalSelector, which is equivalent to sending the message that calls the original method implementation. Build a new Block Invocation for blocks inserted first, then replaced, then later, extracting the parameters from the Invocation and then invokeWithTarget: Block. For an introduction to message forwarding, see my other article understanding Send messages and message forwarding in ObjC in code. Much of the detail in the Aspects implementation code is commendable and supports the method of single instance objects of hook classes (similar to ISa-Swizzlling of KVO). However, the Invoke Block also requires additional invocation overhead because the message is forwarded directly to the original method call and the corresponding function to the real IMP goes through multiple other messages before being executed. The author also writes in the comment that it is not appropriate to add aspect code to methods that are more than 1000 times per second. In addition, when using other methods to hook the Aspect hook method, such as directly replacing it with a new IMP, the original implementation of the new hook is _objc_msgForward, the previous Aspect_hook will fail, and the new hook will execute an exception. Is there a way to replace the IMP of the original method with a function pointer that has the same type encoding as the original method (type encoding)? Multiple code blocks for “after”. Surprisingly, libffi can do all this for us.

1. Libffi profile

Libffi implements runtime on C. In simple terms, libffi generates a template (ffi_cif) based on the number of arguments (ffi_type). You can enter the template, function pointer, and parameter address to complete the function call directly (ffi_call); The template can also generate a so-called closure (FFi_closure) and get a pointer to that address, The custom void function(ffi_cif *cif, void *ret, void **args, void * userData) is implemented, where we can get the addresses of all parameters (including the return values), as well as the custom data userData. Of course, we can do some extra things in this function.

1.1 ffi_type

The respective ffi_type generated based on the number and type of arguments.

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

int fun2 (int a, int b) {
    return2 * a + b; }... ffi_type **types; // Parameter types = malloc(sizeof(ffi_type *) * 2); types[0] = &ffi_type_sint; types[1] = &ffi_type_sint; ffi_type *retType = &ffi_type_sint;Copy the code

1.2 ffi_call

Generates a specific CIF based on ffi_type, and dynamically calls the function by entering CIF, function pointer, and parameter address.

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); // dynamically call fun1 ffi_call(&cif, fun1, &ret, args); . // Output: ret = 3;Copy the code

1.3 ffi_prep_closure_loc

Closure is generated and a function pointer imp is generated. When executed to IMP, all input parameters are taken, followed by ffi_function.

void ffi_function(ffi_cif *cif, void *ret, void **args, void *userdata) { ... // args is the memory address of all parameters} ffi_cif cif; // Generate template 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); // Generate ffi_closure FFi_prep_closure_loc (closure, &cif, ffi_function, (__bridge void *)(self), stingerIMP);Copy the code

Libffi can call any C function in the same way that objc_msgSend does in assembly. Ffi_call pushes parameters into the stack/register according to the template CIF and parameter values. By calling imp in other ways, ffi_closure can get all the parameters based on the stack/register, template CIF, and then execute the code in the custom ffi_function. The JPBlock implementation takes advantage of the latter approach, and for more details see BANG: How to dynamically call C functions. We can replace the IMP of ffi_closure with the pointer associated with ffi_closure. When the object receives a message from this method objc_msgSend(id self, SEL SEL…) , will eventually execute the custom function void ffi_function(ffi_cif *cif, void *ret, void **args, void * userData). The main tasks to achieve all these are: designing feasible structures and storing multiple hook information of the class; Generate cif containing matching FFi_type according to method and section block containing different parameters; The implementation of a method of the replacement class is the IMP associated with FFi_closure, which records the hook; In ffi_function, the original IMP and block are dynamically called based on the parameters obtained.

#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
Copy the code

The following sections describe my implementation around some important points. Stinger

2 Method signature & ffi_type

2.1 Method signature -> ffi_type

For method signatures and Type encoding, I’ve covered a lot in understanding sending messages and message forwarding in ObjC in code. In short, the Type Encoding string corresponds one-to-one to the return type and parameter type of the method. For example: – (void)print1:(NSString *)s; The type encoding of v24@0:8@16. V corresponds to void, @ to ID (in this case self), : to SEL, @ to ID (in this case NSString *). On the other hand, each argument type corresponds to an FFi_type, such as v for FFi_type_void and @ for FFi_type_pointer. You can generate an NSMethodSignature instance object with type encoding using numberOfArguments and – (const char *)getArgumentTypeAtIndex:(NSUInteger)idx; Method to get the parameter type at each location. Of course, you can also filter out numbers to separate strings v24@0:8@16(@? Block) to get an array of parameter types (used in JSPatch). We then convert the method signature to ffi_type by mapping the character to ffi_type.

_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;
}
Copy the code

2.2 introduction to block

2.2.1 Signature & Function pointer

void (^block)(id<StingerParams> params, NSString *s) = ^(id<StingerParams> params, NSString *s) {
    NSLog(@"---after2 print1: %@", s);
}
Copy the code

A block is an ObjC object, and several block types can be thought of as inheriting from NSBlock. Blocks are special in that they ostensibly hold data and objects (variable capture aside), and have executable code that is called in a manner similar to calling C functions, the equivalent of data plus function. Block types are mysterious, but we can see the complete data structure of blocks from OpenSource-Apple/objC4 and oclang/docs/ 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 1struct 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 };Copy the code

Many people have probably seen the code of BlocksKit, and know that Block objects can be strongly converted to Block_layout type, and the signature of the Block can be obtained through identifiers and memory address offsets.

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];
}
Copy the code

NSString *signature = signatureForBlock(block) // Output NSString: @"v24@? 0@\"
      
       \"8@\"NSString\"16"
      

Copy the code

For the simplest signature of a Block object, we can still build NSMethodSignature to get one by one, or we can get an array of characters by filtering out numbers and ‘\”‘.

_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@\"
      
       \"8@\"NSString\"16"
      */ -> v,@? , @ @Copy the code

As you can see, the first sign is “@?” , meaning that the first parameter is blCOk itself, followed by the blCOk parameter type. In the same way, we can still match ffi_type with type encoding. In addition, we can obtain the function pointer of the Block object directly.

BlockIMP impForBlock(id block) {
  struct Block_layout *layout = (__bridge void *)block;
  return layout->invoke;
}
Copy the code

A simple attempt is to call the contained function of the Block object directly.

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"); // Output: after2print1: tt
Copy the code

In addition, the parameter corresponding to the function pointer obtained by IMP _Nonnull imp_implementationWithBlock(ID _Nonnull block) does not contain the block itself, which means that the signature has changed.

* Adds available methods to block objects

There are several ways in which we can feel that a Block object has a new instance method.

NSString *signature = [block signature];
void *blockIMP = [block blockIMP];
Copy the code

The way to do that is to add instance methods to the NSBlock class in STBlock.

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
Copy the code

#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)); }... @endCopy the code

By doing so, you add a message that can be processed to the Block object. However, if you try to call it from a load method of another class, you may encounter the problem that the load method of the STBlock class is not loaded.

Store hook information & generate two FFI_CIF objects

3.1 StingerInfo

Simple objects are used here to store individual hook information.

@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
Copy the code

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
Copy the code

3.2.1 management StingerInfo

Here, three arrays are used to store the ID

object of a class hook location before, after, and after the original implementation, and to save the original IMP. Adding and removing objects with ID

are thread-safe.

3.2.2 Generating method call template CIF

Based on the type encoding provided by the original method, ffi_type is generated for each parameter, then the CIF object is generated, and finally ffi_prep_CLOSURE_LOc is equivalent to generating the empty shell function StingerIMP. Calling StingerIMP will ultimately execute to custom static void ffi_function(ffi_cif *cif, void *ret, void **args, void * userData) functions, This function gets all the arguments obtained when StingerIMP is called.

- (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;
}
Copy the code

3.2.3 Generating Block To invoke template blockCif

This is similar to the previous generation of the method call template CIF, except the shell FFi_closure is not generated. It is worth noting that bits 0 (@self) and 1 (: SEL) of the original method type encoing are replaced with bits (@? Block) and (@id

). This means that the signature type of the slice Block object is qualified.

- (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"); }}Copy the code

In the non-instead position, the return value of a block can be arbitrary; When writing a block, the zeroth bit of the block (regardless of the block itself) should be of type ID, followed by the argument corresponding to the original method.

3.2.4 ffi_function!!!!!!

In this function, we get the memory address of all the input parameters for the original method call. We first generate a new parameter set, innerArgs, from the block_CIF template, with bit 0 reserved for the Block object, bit 1 reserved for the StingerParams object, and copy the original parameters from bit 2. Here is how to complete the section code and the original IMP execution: 1. Use ffi_call(&(self->_blockCif), impForBlock(block), NULL, innerArgs); Completes all calls to the front block where the cut position is. Use block templates blockCif and innerArgs. 2. Use ffi_call(cif, (void (*)(void))self.originalIMP/impForBlock(block), ret, args); Complete the call to the original IMP or the replacement location Block IMP. The original template CIF and the original parameter ARgs are used, and the return value may be generated. Use ffi_call(&(self->_blockCif), impForBlock(block), NULL, innerArgs); Completes all calls to the block behind the cut position. Use block templates blockCif and 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);
}
Copy the code

Note: The StingerParams object contains the message receiver SLF, the selector sel for the current message, and the Invocation of the original method (using invokeUsingIMP: complete the invocation), This invocation is only appropriate when the method is replaced and the original return value is required as a parameter. Other hooks use optionBefore or after without paying attention to the 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
Copy the code

4 Replacement method implementation & record HOOK

The first hook is created to try to replace IMP ffi_prep_closure_LOc, which is implemented as ffi_prep_closure_LOc. Hook Info is added directly to the associated ID

object. In terms of conditions, the most important are two points. The first point is that for a SEL in a class (parent class), SEL should be able to find the corresponding Method M and IMP IMP; The second point is that the signature of the block on the section isMatched with that of the original method, and the signature of the block on the section meets the requirements (isMatched method).

#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
Copy the code

* Usage Examples

import UIKit;

@interface ASViewController : UIViewController

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

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

@end
Copy the code

#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

Copy the code

The Stinger usage is similar to Aspects, but since the block and the original IMP are called directly from the function pointer when the message is received, no extra messages are handled and no NSInvocation objects are instantiated, the two lib_CIF objects are hook ready, compared to Aspects, The measured time from receiving the message to executing the section information (e.g., before block) is reduced by an order of magnitude. The validity of ST_hook can still be guaranteed when other hook methods are used.

Thank you for watching, the level is limited, if there is an error, please correct.

Github.com/Assuner-Lee…

The resources

Github.com/opensource-… Blog.cnbang.net/tech/3219/ juejin. Cn/post / 684490… Github.com/mikeash/MAB…

Please correct any intellectual property rights, copyright issues or theoretical errors.

Please indicate the original author and above information.