From ea57e522c495418a0b4b717dfd5438417a3f7d65 Mon Sep 17 00:00:00 2001
From: maocl023 <maocl023@gmail.com>
Date: Mon, 28 Mar 2016 16:07:16 +0800
Subject: [PATCH] avoid `AspectPositionAfter | AspectPositionInstead` or
 `AspectPositionAfter | AspectPositionBefore` misunderstanding, \ limit the
 AspectOptions value to 0xc(AspectPositionBefore |
 AspectOptionAutomaticRemoval)

---
 Aspects.h |  9 +++++----
 Aspects.m | 28 +++++++++++++++++++++-------
 2 files changed, 26 insertions(+), 11 deletions(-)

diff --git a/Aspects.h b/Aspects.h
index 5508f86..de6e858 100644
--- a/Aspects.h
+++ b/Aspects.h
@@ -8,11 +8,11 @@
 #import <Foundation/Foundation.h>
 
 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.
+    AspectPositionAfter   = 1 << 0,             /// Called after the original implementation (default)
+    AspectPositionInstead = 1 << 1,             /// Will replace the original implementation.
+    AspectPositionBefore  = 1 << 2,             /// Called before the original implementation.
     
-    AspectOptionAutomaticRemoval = 1 << 3 /// Will remove the hook after the first execution.
+    AspectOptionAutomaticRemoval = 1 << 3       /// Will remove the hook after the first execution.
 };
 
 /// Opaque Aspect Token that allows to deregister the hook.
@@ -71,6 +71,7 @@ typedef NS_OPTIONS(NSUInteger, AspectOptions) {
 typedef NS_ENUM(NSUInteger, AspectErrorCode) {
     AspectErrorSelectorBlacklisted,                   /// Selectors like release, retain, autorelease are blacklisted.
     AspectErrorDoesNotRespondToSelector,              /// Selector could not be found.
+    AspectErrorOptionsValue,                          /// Options value should less than 0xc(AspectPositionBefore | AspectOptionAutomaticRemoval).
     AspectErrorSelectorDeallocPosition,               /// When hooking dealloc, only AspectPositionBefore is allowed.
     AspectErrorSelectorAlreadyHookedInClassHierarchy, /// Statically hooking the same method in subclasses is not allowed.
     AspectErrorFailedToAllocateClassPair,             /// The runtime failed creating a class pair.
diff --git a/Aspects.m b/Aspects.m
index c907066..3710b19 100644
--- a/Aspects.m
+++ b/Aspects.m
@@ -81,7 +81,7 @@ @interface NSInvocation (Aspects)
 - (NSArray *)aspects_arguments;
 @end
 
-#define AspectPositionFilter 0x07
+#define AspectPositionFilter 0xc
 
 #define AspectError(errorCode, errorDescription) do { \
 AspectLogError(@"Aspects: %@", errorDescription); \
@@ -465,7 +465,7 @@ static void aspect_undoSwizzleClassInPlace(Class klass) {
 #define aspect_invoke(aspects, info) \
 for (AspectIdentifier *aspect in aspects) {\
     [aspect invokeWithInfo:info];\
-    if (aspect.options & AspectOptionAutomaticRemoval) { \
+    if (aspect.options >= AspectOptionAutomaticRemoval) { \
         aspectsToRemove = [aspectsToRemove?:@[] arrayByAddingObject:aspect]; \
     } \
 }
@@ -581,7 +581,13 @@ static BOOL aspect_isSelectorAllowedAndTrack(NSObject *self, SEL selector, Aspec
     }
 
     // Additional checks.
-    AspectOptions position = options&AspectPositionFilter;
+    if (options > AspectPositionFilter) {
+        NSString *errorDesc = @"AspectOptions max value is AspectPositionBefore | AspectOptionAutomaticRemoval";
+        AspectError(AspectErrorOptionsValue, errorDesc);
+        return NO;
+    }
+
+    AspectOptions position = options % AspectOptionAutomaticRemoval;
     if ([selectorName isEqualToString:@"dealloc"] && position != AspectPositionBefore) {
         NSString *errorDesc = @"AspectPositionBefore is the only valid position when hooking dealloc.";
         AspectError(AspectErrorSelectorDeallocPosition, errorDesc);
@@ -887,11 +893,19 @@ - (BOOL)hasAspects {
 
 - (void)addAspect:(AspectIdentifier *)aspect withOptions:(AspectOptions)options {
     NSParameterAssert(aspect);
-    NSUInteger position = options&AspectPositionFilter;
+    NSUInteger position = options % AspectOptionAutomaticRemoval;
+
     switch (position) {
-        case AspectPositionBefore:  self.beforeAspects  = [(self.beforeAspects ?:@[]) arrayByAddingObject:aspect]; break;
-        case AspectPositionInstead: self.insteadAspects = [(self.insteadAspects?:@[]) arrayByAddingObject:aspect]; break;
-        case AspectPositionAfter:   self.afterAspects   = [(self.afterAspects  ?:@[]) arrayByAddingObject:aspect]; break;
+        case AspectPositionBefore:
+            self.beforeAspects  = [(self.beforeAspects ?:@[]) arrayByAddingObject:aspect];
+            break;
+        case AspectPositionInstead:
+            self.insteadAspects = [(self.insteadAspects?:@[]) arrayByAddingObject:aspect];
+            break;
+        case AspectPositionAfter:
+        case 0:// AspectOptionAutomaticRemoval
+            self.afterAspects   = [(self.afterAspects  ?:@[]) arrayByAddingObject:aspect];
+            break;
     }
 }