test/customrr-nsobject.m

   1 // This file is used in the customrr-nsobject-*.m tests
   2
   3 #include "test.h"
   4 #include <objc/NSObject.h>
   5 #include <objc/objc-internal.h>
   6
   7 #if __has_feature(ptrauth_calls)
   8 typedef IMP __ptrauth_objc_method_list_imp MethodListIMP;
   9 #else
  10 typedef IMP MethodListIMP;
  11 #endif
  12
  13 static int Retains;
  14 static int Releases;
  15 static int Autoreleases;
  16 static int PlusInitializes;
  17 static int Allocs;
  18 static int AllocWithZones;
  19 static int Inits;
  20 static int PlusNew;
  21 static int Self;
  22 static int PlusSelf;
  23
  24 id (*RealRetain)(id self, SEL _cmd);
  25 void (*RealRelease)(id self, SEL _cmd);
  26 id (*RealAutorelease)(id self, SEL _cmd);
  27 id (*RealAlloc)(id self, SEL _cmd);
  28 id (*RealAllocWithZone)(id self, SEL _cmd, void *zone);
  29 id (*RealPlusNew)(id self, SEL _cmd);
  30 id (*RealSelf)(id self);
  31 id (*RealPlusSelf)(id self);
  32
  33 id HackRetain(id self, SEL _cmd) { Retains++; return RealRetain(self, _cmd); }
  34 void HackRelease(id self, SEL _cmd) { Releases++; return RealRelease(self, _cmd); }
  35 id HackAutorelease(id self, SEL _cmd) { Autoreleases++; return RealAutorelease(self, _cmd); }
  36
  37 id HackAlloc(Class self, SEL _cmd) { Allocs++; return RealAlloc(self, _cmd); }
  38 id HackAllocWithZone(Class self, SEL _cmd, void *zone) { AllocWithZones++; return RealAllocWithZone(self, _cmd, zone); }
  39
  40 void HackPlusInitialize(id self __unused, SEL _cmd __unused) { PlusInitializes++; }
  41
  42 id HackInit(id self, SEL _cmd __unused) { Inits++; return self; }
  43
  44 id HackPlusNew(id self, SEL _cmd __unused) { PlusNew++; return RealPlusNew(self, _cmd); }
  45 id HackSelf(id self) { Self++; return RealSelf(self); }
  46 id HackPlusSelf(id self) { PlusSelf++; return RealPlusSelf(self); }
  47
  48
  49 int main(int argc __unused, char **argv)
  50 {
  51     Class cls = objc_getClass("NSObject");
  52     Method meth;
  53
  54     meth = class_getClassMethod(cls, @selector(initialize));
  55     method_setImplementation(meth, (IMP)HackPlusInitialize);
  56
  57     // We either swizzle the method normally (testing that it properly
  58     // disables optimizations), or we hack the implementation into place
  59     // behind objc's back (so we can see whether it got called with the
  60     // optimizations still enabled).
  61
  62     meth = class_getClassMethod(cls, @selector(allocWithZone:));
  63     RealAllocWithZone = (typeof(RealAllocWithZone))method_getImplementation(meth);
  64 #if SWIZZLE_AWZ
  65     method_setImplementation(meth, (IMP)HackAllocWithZone);
  66 #else
  67     ((MethodListIMP *)meth)[2] = (IMP)HackAllocWithZone;
  68 #endif
  69
  70     meth = class_getClassMethod(cls, @selector(new));
  71     RealPlusNew = (typeof(RealPlusNew))method_getImplementation(meth);
  72 #if SWIZZLE_CORE
  73     method_setImplementation(meth, (IMP)HackPlusNew);
  74 #else
  75     ((MethodListIMP *)meth)[2] = (IMP)HackPlusNew;
  76 #endif
  77
  78     meth = class_getClassMethod(cls, @selector(self));
  79     RealPlusSelf = (typeof(RealPlusSelf))method_getImplementation(meth);
  80 #if SWIZZLE_CORE
  81     method_setImplementation(meth, (IMP)HackPlusSelf);
  82 #else
  83     ((MethodListIMP *)meth)[2] = (IMP)HackPlusSelf;
  84 #endif
  85
  86     meth = class_getInstanceMethod(cls, @selector(self));
  87     RealSelf = (typeof(RealSelf))method_getImplementation(meth);
  88 #if SWIZZLE_CORE
  89     method_setImplementation(meth, (IMP)HackSelf);
  90 #else
  91     ((MethodListIMP *)meth)[2] = (IMP)HackSelf;
  92 #endif
  93
  94     meth = class_getInstanceMethod(cls, @selector(release));
  95     RealRelease = (typeof(RealRelease))method_getImplementation(meth);
  96 #if SWIZZLE_RELEASE
  97     method_setImplementation(meth, (IMP)HackRelease);
  98 #else
  99     ((MethodListIMP *)meth)[2] = (IMP)HackRelease;
 100 #endif
 101
 102     // These other methods get hacked for counting purposes only
 103
 104     meth = class_getInstanceMethod(cls, @selector(retain));
 105     RealRetain = (typeof(RealRetain))method_getImplementation(meth);
 106     ((MethodListIMP *)meth)[2] = (IMP)HackRetain;
 107
 108     meth = class_getInstanceMethod(cls, @selector(autorelease));
 109     RealAutorelease = (typeof(RealAutorelease))method_getImplementation(meth);
 110     ((MethodListIMP *)meth)[2] = (IMP)HackAutorelease;
 111
 112     meth = class_getClassMethod(cls, @selector(alloc));
 113     RealAlloc = (typeof(RealAlloc))method_getImplementation(meth);
 114     ((MethodListIMP *)meth)[2] = (IMP)HackAlloc;
 115
 116     meth = class_getInstanceMethod(cls, @selector(init));
 117     ((MethodListIMP *)meth)[2] = (IMP)HackInit;
 118
 119     // Verify that the swizzles occurred before +initialize by provoking it now
 120     testassert(PlusInitializes == 0);
 121     [NSObject self];
 122     testassert(PlusInitializes == 1);
 123
 124     id obj;
 125     id result;
 126
 127     Allocs = 0;
 128     AllocWithZones = 0;
 129     Inits = 0;
 130     obj = objc_alloc(cls);
 131 #if SWIZZLE_AWZ
 132     testprintf("swizzled AWZ should be called\n");
 133     testassert(Allocs == 1);
 134     testassert(AllocWithZones == 1);
 135     testassert(Inits == 0);
 136 #else
 137     testprintf("unswizzled AWZ should be bypassed\n");
 138     testassert(Allocs == 0);
 139     testassert(AllocWithZones == 0);
 140     testassert(Inits == 0);
 141 #endif
 142     testassert([obj isKindOfClass:[NSObject class]]);
 143
 144     Allocs = 0;
 145     AllocWithZones = 0;
 146     Inits = 0;
 147     obj = [NSObject alloc];
 148 #if SWIZZLE_AWZ
 149     testprintf("swizzled AWZ should be called\n");
 150     testassert(Allocs == 1);
 151     testassert(AllocWithZones == 1);
 152     testassert(Inits == 0);
 153 #else
 154     testprintf("unswizzled AWZ should be bypassed\n");
 155     testassert(Allocs == 1);
 156     testassert(AllocWithZones == 0);
 157     testassert(Inits == 0);
 158 #endif
 159     testassert([obj isKindOfClass:[NSObject class]]);
 160
 161     Allocs = 0;
 162     AllocWithZones = 0;
 163     Inits = 0;
 164     obj = objc_alloc_init(cls);
 165 #if SWIZZLE_AWZ
 166     testprintf("swizzled AWZ should be called\n");
 167     testassert(Allocs == 1);
 168     testassert(AllocWithZones == 1);
 169     testassert(Inits == 1);
 170 #else
 171     testprintf("unswizzled AWZ should be bypassed\n");
 172     testassert(Allocs == 0);
 173     testassert(AllocWithZones == 0);
 174     testassert(Inits == 1);  // swizzled init is still called
 175 #endif
 176     testassert([obj isKindOfClass:[NSObject class]]);
 177
 178     Retains = 0;
 179     result = objc_retain(obj);
 180 #if SWIZZLE_RELEASE
 181     testprintf("swizzled release should force retain\n");
 182     testassert(Retains == 1);
 183 #else
 184     testprintf("unswizzled release should bypass retain\n");
 185     testassert(Retains == 0);
 186 #endif
 187     testassert(result == obj);
 188
 189     Releases = 0;
 190     Autoreleases = 0;
 191     PUSH_POOL {
 192         result = objc_autorelease(obj);
 193 #if SWIZZLE_RELEASE
 194         testprintf("swizzled release should force autorelease\n");
 195         testassert(Autoreleases == 1);
 196 #else
 197         testprintf("unswizzled release should bypass autorelease\n");
 198         testassert(Autoreleases == 0);
 199 #endif
 200         testassert(result == obj);
 201     } POP_POOL
 202
 203 #if SWIZZLE_RELEASE
 204     testprintf("swizzled release should be called\n");
 205     testassert(Releases == 1);
 206 #else
 207     testprintf("unswizzled release should be bypassed\n");
 208     testassert(Releases == 0);
 209 #endif
 210
 211     PlusNew = 0;
 212     Self = 0;
 213     PlusSelf = 0;
 214     Class nso = objc_opt_self([NSObject class]);
 215     obj = objc_opt_new(nso);
 216     obj = objc_opt_self(obj);
 217 #if SWIZZLE_CORE
 218     testprintf("swizzled Core should be called\n");
 219     testassert(PlusNew == 1);
 220     testassert(Self == 1);
 221     testassert(PlusSelf == 1);
 222 #else
 223     testprintf("unswizzled CORE should be bypassed\n");
 224     testassert(PlusNew == 0);
 225     testassert(Self == 0);
 226     testassert(PlusSelf == 0);
 227 #endif
 228     testassert([obj isKindOfClass:nso]);
 229
 230     succeed(basename(argv[0]));
 231 }