Skip to content

Commit

Permalink
Removing now works via calling remove on the aspect token.
Browse files Browse the repository at this point in the history
  • Loading branch information
@steipete
steipete committed May 4, 2014
1 parent 1b94a30 commit 4de65c8
Show file tree
Hide file tree
Showing 4 changed files with 88 additions and 77 deletions.
19 changes: 11 additions & 8 deletions Aspects.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,15 @@ typedef NS_ENUM(NSUInteger, AspectPosition) {
AspectPositionAfter
};

/// Opaque Aspect Token that allows to deregister the hook.
@protocol Aspect <NSObject>

/// Deegister an aspect.
/// @return YES if deregistration is successful, otherwise NO.
- (BOOL)remove;

@end

/**
Aspects uses Objective-C message forwarding to hook into messages. This will create some overhead. Don't add aspects to methods that are called a lot. Aspects is meant for view/controller code that is not called a 1000 times per second.
Expand All @@ -24,15 +33,9 @@ typedef NS_ENUM(NSUInteger, AspectPosition) {
/// Adds a block of code before/instead/after the current `selector` for a specific object.
/// If you choose `AspectPositionInstead`, `arguments` contains an additional argument which is the original invocation.
/// @return A token which allows to later deregister the aspect.
- (id)aspect_hookSelector:(SEL)selector atPosition:(AspectPosition)position withBlock:(void (^)(__unsafe_unretained id object, NSArray *arguments))block;
- (id<Aspect>)aspect_hookSelector:(SEL)selector atPosition:(AspectPosition)position withBlock:(void (^)(__unsafe_unretained id object, NSArray *arguments))block;

/// Hooks a selector class-wide.
+ (id)aspect_hookSelector:(SEL)selector atPosition:(AspectPosition)position withBlock:(void (^)(__unsafe_unretained id object, NSArray *arguments))block;

/// Unregister an aspect.
/// @note This is a single method that works for both object and class-based hooks, as all required state is preserved in the opaque aspect token.
/// Can also be called via `[NSObject aspect_remove:]`. The class context is not used.
/// @return YES if deregistration is successful, otherwise NO.
+ (BOOL)aspect_remove:(id)aspect;
+ (id<Aspect>)aspect_hookSelector:(SEL)selector atPosition:(AspectPosition)position withBlock:(void (^)(__unsafe_unretained id object, NSArray *arguments))block;

@end
90 changes: 48 additions & 42 deletions Aspects.m
View file
Original file line number Diff line number Diff line change
Expand Up @@ -53,10 +53,6 @@ - (id)aspect_hookSelector:(SEL)selector atPosition:(AspectPosition)position with
return aspect_add(self, selector, position, block);
}

+ (BOOL)aspect_remove:(id)aspect {
return aspect_remove(aspect);
}

///////////////////////////////////////////////////////////////////////////////////////////
#pragma mark - Private Helper

Expand Down Expand Up @@ -110,34 +106,8 @@ static SEL aspect_aliasForSelector(SEL selector) {
return NSSelectorFromString([AspectMessagePrefix stringByAppendingFormat:@"_%@", NSStringFromSelector(selector)]);
}

// Loads or creates the aspect container.
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;
}

static AspectsContainer *aspect_getContainerForClass(Class klass, SEL selector) {
NSCParameterAssert(klass);
AspectsContainer *classContainer = nil;
do {
classContainer = objc_getAssociatedObject(klass, selector);
if (classContainer.hasAspects) break;
}while ((klass = class_getSuperclass(klass)));

return classContainer;
}

static void aspect_destroyContainerForObject(id<NSObject> self, SEL selector) {
NSCParameterAssert(self);
SEL aliasSelector = aspect_aliasForSelector(selector);
objc_setAssociatedObject(self, aliasSelector, nil, OBJC_ASSOCIATION_RETAIN);
}
///////////////////////////////////////////////////////////////////////////////////////////
#pragma mark - Class + Selector Preparation

static void aspect_prepareClassAndHookSelector(NSObject *self, SEL selector) {
NSCParameterAssert(selector);
Expand Down Expand Up @@ -403,19 +373,51 @@ static void __ASPECTS_ARE_BEING_CALLED__(__unsafe_unretained NSObject *self, SEL
}
#undef aspect_invoke

///////////////////////////////////////////////////////////////////////////////////////////
#pragma mark - Aspect Container Management

// Loads or creates the aspect container.
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;
}

static AspectsContainer *aspect_getContainerForClass(Class klass, SEL selector) {
NSCParameterAssert(klass);
AspectsContainer *classContainer = nil;
do {
classContainer = objc_getAssociatedObject(klass, selector);
if (classContainer.hasAspects) break;
}while ((klass = class_getSuperclass(klass)));

return classContainer;
}

static void aspect_destroyContainerForObject(id<NSObject> self, SEL selector) {
NSCParameterAssert(self);
SEL aliasSelector = aspect_aliasForSelector(selector);
objc_setAssociatedObject(self, aliasSelector, nil, OBJC_ASSOCIATION_RETAIN);
}

@end

///////////////////////////////////////////////////////////////////////////////////////////
#pragma mark - Selector Blacklist Checking

@interface AspectSelectorTracker : NSObject
- (id)initWithTrackedClass:(Class)trackedClass parent:(AspectSelectorTracker *)parent;
@interface AspectTracker : NSObject
- (id)initWithTrackedClass:(Class)trackedClass parent:(AspectTracker *)parent;
@property (nonatomic, strong) Class trackedClass;
@property (nonatomic, strong) NSMutableSet *selectorNames;
@property (nonatomic, weak) AspectSelectorTracker *parentEntry;
@property (nonatomic, weak) AspectTracker *parentEntry;
@end
@implementation AspectSelectorTracker
- (id)initWithTrackedClass:(Class)trackedClass parent:(AspectSelectorTracker *)parent {
@implementation AspectTracker
- (id)initWithTrackedClass:(Class)trackedClass parent:(AspectTracker *)parent {
if (self = [super init]) {
_trackedClass = trackedClass;
_parentEntry = parent;
Expand Down Expand Up @@ -454,12 +456,12 @@ BOOL aspect_isSelectorAllowed(Class class, SEL selector, AspectPosition position
// Search for the current class and the class hierarchy.
Class currentClass = class;
do {
AspectSelectorTracker *tracker = swizzledClassesDict[currentClass];
AspectTracker *tracker = swizzledClassesDict[currentClass];
if ([tracker.selectorNames containsObject:selectorName]) {

// Find the topmost class for the log.
if (tracker.parentEntry) {
AspectSelectorTracker *topmostEntry = tracker.parentEntry;
AspectTracker *topmostEntry = tracker.parentEntry;
while (topmostEntry.parentEntry) {
topmostEntry = topmostEntry.parentEntry;
}
Expand All @@ -474,11 +476,11 @@ BOOL aspect_isSelectorAllowed(Class class, SEL selector, AspectPosition position

// Add the selector as being modified.
currentClass = class;
AspectSelectorTracker *parentTracker = nil;
AspectTracker *parentTracker = nil;
do {
AspectSelectorTracker *tracker = swizzledClassesDict[currentClass];
AspectTracker *tracker = swizzledClassesDict[currentClass];
if (!tracker) {
tracker = [[AspectSelectorTracker alloc] initWithTrackedClass:currentClass parent:parentTracker];
tracker = [[AspectTracker alloc] initWithTrackedClass:currentClass parent:parentTracker];
swizzledClassesDict[(id<NSCopying>)currentClass] = tracker;
}
[tracker.selectorNames addObject:selectorName];
Expand Down Expand Up @@ -591,6 +593,10 @@ - (NSString *)description {
return [NSString stringWithFormat:@"<%@: %p, SEL:%@ object:%@ block:%@>", self.class, self, NSStringFromSelector(self.selector), self.object, self.block];
}

- (BOOL)remove {
return aspect_remove(self);
}

@end

///////////////////////////////////////////////////////////////////////////////////////////
Expand Down
46 changes: 23 additions & 23 deletions AspectsDemo/AspectsDemoTests/AspectsDemoTests.m
View file
Original file line number Diff line number Diff line change
Expand Up @@ -103,16 +103,16 @@ - (void)testInsteadHook {
XCTAssertEqualObjects(testLabel3.text, @"Global", @"Must match");
testLabel3.text = @"Test";
XCTAssertEqualObjects(testLabel3.text, @"Global", @"Must match");
XCTAssertTrue([UILabel aspect_remove:globalAspect], @"Must work");
XCTAssertTrue([globalAspect remove], @"Must work");
XCTAssertEqualObjects(testLabel3.text, @"Test", @"Must match");

// Test that removing an aspect returns the original.
XCTAssertEqualObjects(testLabel.text, @"Custom Text", @"Must match");
XCTAssertTrue([UILabel aspect_remove:aspect], @"Must return YES");
XCTAssertTrue([aspect remove], @"Must return YES");
XCTAssertEqualObjects(testLabel.text, @"Default text", @"Must match");
XCTAssertFalse([UILabel aspect_remove:aspect], @"Must return NO");
XCTAssertFalse([aspect remove], @"Must return NO");

XCTAssertTrue([UILabel aspect_remove:aspect2], @"Must be able to deregister");
XCTAssertTrue([aspect2 remove], @"Must be able to deregister");
}

- (void)testAspectsCalledPerObject {
Expand All @@ -130,7 +130,7 @@ - (void)testAspectsCalledPerObject {
[testClass2 testCall];
XCTAssertFalse(called, @"Flag must have been NOT set.");

XCTAssertTrue([NSObject aspect_remove:aspect], @"Must be able to deregister");
XCTAssertTrue([aspect remove], @"Must be able to deregister");
}

- (void)testExecutionOrderAndMultipleRegistation {
Expand Down Expand Up @@ -158,12 +158,12 @@ - (void)testExecutionOrderAndMultipleRegistation {
XCTAssertTrue(called_after, @"Flag must have been set.");
XCTAssertTrue(called_after2, @"Flag must have been set.");

XCTAssertTrue([NSObject aspect_remove:aspect_after], @"Must be able to deregister");
XCTAssertTrue([NSObject aspect_remove:aspect_before], @"Must be able to deregister");
XCTAssertTrue([NSObject aspect_remove:aspect_after2], @"Must be able to deregister");
XCTAssertFalse([NSObject aspect_remove:aspect_after], @"Must not be able to deregister twice");
XCTAssertFalse([NSObject aspect_remove:aspect_before], @"Must not be able to deregister twice");
XCTAssertFalse([NSObject aspect_remove:aspect_after2], @"Must not be able to deregister twice");
XCTAssertTrue([aspect_after remove], @"Must be able to deregister");
XCTAssertTrue([aspect_before remove], @"Must be able to deregister");
XCTAssertTrue([aspect_after2 remove], @"Must be able to deregister");
XCTAssertFalse([aspect_after remove], @"Must not be able to deregister twice");
XCTAssertFalse([aspect_before remove], @"Must not be able to deregister twice");
XCTAssertFalse([aspect_after2 remove], @"Must not be able to deregister twice");
}

- (void)testExample {
Expand All @@ -179,7 +179,7 @@ - (void)testExample {
[testClass testCall];
}];
XCTAssertTrue(testCallCalled, @"Calling testCallAndExecuteBlock must call testCall");
XCTAssertTrue([NSObject aspect_remove:aspectToken], @"Must be able to deregister");
XCTAssertTrue([aspectToken remove], @"Must be able to deregister");
}

- (void)testStructReturn {
Expand All @@ -190,7 +190,7 @@ - (void)testStructReturn {

CGRect rectHooked = [testClass testThatReturnsAStruct];
XCTAssertTrue(CGRectEqualToRect(rect, rectHooked), @"Must be equal");
XCTAssertTrue([NSObject aspect_remove:aspect], @"Must be able to deregister");
XCTAssertTrue([aspect remove], @"Must be able to deregister");
}

- (void)testHookReleaseIsNotAllowed {
Expand Down Expand Up @@ -256,14 +256,14 @@ - (void)testInstanceTokenDeregistration {

XCTAssertNotEqualObjects(testClass.class, object_getClass(testClass), @"Object must have a custom subclass.");

XCTAssertTrue([TestClass aspect_remove:aspectToken], @"Deregistration must work");
XCTAssertTrue([aspectToken remove], @"Deregistration must work");
XCTAssertEqualObjects(testClass.class, object_getClass(testClass), @"Object must not have a custom subclass.");

testCallCalled = NO;
[testClass testCall];
XCTAssertFalse(testCallCalled, @"Hook must no longer work.");

XCTAssertFalse([TestClass aspect_remove:aspectToken], @"Deregistration must not work twice");
XCTAssertFalse([aspectToken remove], @"Deregistration must not work twice");
}

- (void)testGlobalTokenDeregistrationWithCustomForwardInvocation {
Expand Down Expand Up @@ -294,7 +294,7 @@ - (void)testGlobalTokenDeregistrationWithCustomForwardInvocation {
XCTAssertNotEqual(method_getImplementation(forwardInvocationMethod), originalForwardInvocationIMP, @"Implementations must not be equal");
}

XCTAssertTrue([TestClass aspect_remove:token], @"Deregistration must work");
XCTAssertTrue([token remove], @"Deregistration must work");

// Test that forwardInvocation (again) points to NSObject and thus is correctly restored.
{
Expand All @@ -306,7 +306,7 @@ - (void)testGlobalTokenDeregistrationWithCustomForwardInvocation {
[testClass test];
XCTAssertFalse(testCalled, @"Hook must no longer work.");

XCTAssertFalse([TestWithCustomForwardInvocation aspect_remove:token], @"Deregistration must not work twice");
XCTAssertFalse([token remove], @"Deregistration must not work twice");
}

- (void)testGlobalTokenDeregistration {
Expand Down Expand Up @@ -337,7 +337,7 @@ - (void)testGlobalTokenDeregistration {
XCTAssertNotEqual(method_getImplementation(forwardInvocationMethod), method_getImplementation(objectMethod), @"Implementations must not be equal");
}

XCTAssertTrue([TestClass aspect_remove:token], @"Deregistration must work");
XCTAssertTrue([token remove], @"Deregistration must work");

// Test that forwardInvocation (again) points to NSObject and thus is correctly restored.
{
Expand All @@ -350,7 +350,7 @@ - (void)testGlobalTokenDeregistration {
[testClass testCall];
XCTAssertFalse(testCallCalled, @"Hook must no longer work.");

XCTAssertFalse([TestClass aspect_remove:token], @"Deregistration must not work twice");
XCTAssertFalse([token remove], @"Deregistration must not work twice");
}

- (void)testSimpleDeregistration {
Expand All @@ -364,7 +364,7 @@ - (void)testSimpleDeregistration {
XCTAssertTrue(called, @"Flag must have been set.");

called = NO;
[TestClass aspect_remove:aspectToken];
XCTAssertTrue([aspectToken remove], @"Must allow deregistration");
[testClass testCall];
XCTAssertFalse(called, @"Flag must have been NOT set.");
}
Expand Down Expand Up @@ -393,7 +393,7 @@ - (void)testKVOCoexistance {
XCTAssertTrue(hookCalled, @"Hook must be called");
XCTAssertFalse(testClass.kvoTestCalled, @"KVO must no longer work");

XCTAssertTrue([NSObject aspect_remove:aspectToken], @"Must be able to deregister");
XCTAssertTrue([aspectToken remove], @"Must be able to deregister");
}

// TODO: Pre-registeded KVO is currently not working.
Expand Down Expand Up @@ -448,7 +448,7 @@ - (void)testEnsureForwardInvocationIsCalled {
#pragma clang diagnostic pop
XCTAssertTrue(testClass.forwardInvocationCalled, @"Must have called custom forwardInvocation:");

XCTAssertTrue([NSObject aspect_remove:aspectToken], @"Must be able to deregister");
XCTAssertTrue([aspectToken remove], @"Must be able to deregister");
}

@end
Expand Down Expand Up @@ -520,7 +520,7 @@ - (void)testSelectorMangling2 {
XCTAssertTrue(A_aspect_called, @"A aspect should be called");
XCTAssertFalse(B_aspect_called, @"B aspect should not be called");

XCTAssertTrue([NSObject aspect_remove:aspectToken1], @"Must be able to deregister");
XCTAssertTrue([aspectToken1 remove], @"Must be able to deregister");
}

@end
10 changes: 6 additions & 4 deletions README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,18 +11,19 @@ Aspects extends `NSObject` with the following methods:
/// Adds a block of code before/instead/after the current `selector` for a specific object.
/// If you choose `AspectPositionInstead`, `arguments` contains an additional argument which is the original invocation.
/// @return A token which allows to later deregister the aspect.
- (id)aspect_hookSelector:(SEL)selector
- (id<Aspect>)aspect_hookSelector:(SEL)selector
atPosition:(AspectPosition)position
withBlock:(void (^)(id object, NSArray *arguments))block;

/// Hooks a selector class-wide.
+ (id)aspect_hookSelector:(SEL)selector
+ (id<Aspect>)aspect_hookSelector:(SEL)selector
atPosition:(AspectPosition)position
withBlock:(void (^)(id object, NSArray *arguments))block;

/// Unregister an aspect.
/// Deregister an aspect.
/// @return YES if deregistration is successful, otherwise NO.
+ (BOOL)aspect_remove:(id)aspect;
id<Aspect> aspect = ...;
[aspect remove];
```

Adding aspects returns an opaque token which can be used to deregister again. All calls are thread-safe.
Expand Down Expand Up @@ -151,6 +152,7 @@ Release Notes
Version 1.1.0 (Upcoming)

- Renamed the files from NSObject+Aspects.m/h to just Aspects.m/h.
- Removing now works via calling `remove` on the aspect token.
- Allow hooking dealloc.
- Fixes infinite loop if the same method is hooked for multiple classes. Hooking will only work for one class in the hierarchy.
- Additional checks to prevent things like hooking retain/release/autorelease or forwardInvocation:
Expand Down

0 comments on commit 4de65c8

Please sign in to comment.