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