Implemented AddAllocWithZoneMethod to override allocWithZone: on demand.

BUG=None

Review URL: https://codereview.chromium.org/1043263003

Cr-Commit-Position: refs/heads/master@{#325092}

5 files changed, 88 insertions, 0 deletions

diff --git a/ios/web/alloc_with_zone_interceptor.h b/ios/web/alloc_with_zone_interceptor.h
new file mode 100644
index 0000000..0026540
--- /dev/null
+++ b/ios/web/alloc_with_zone_interceptor.h
@@ -0,0 +1,18 @@
+// Copyright 2015 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#ifndef IOS_WEB_ALLOC_WITH_ZONE_INTERCEPTOR_H_
+#define IOS_WEB_ALLOC_WITH_ZONE_INTERCEPTOR_H_
+
+#import <Foundation/Foundation.h>
+
+namespace web {
+// Adds |impl_block| as a custom implementation for |allocWithZone:| method of
+// |target| Class. AddAllocWithZoneMethod can not swizzle existing
+// |allocWithZone:| method of |target| Class. Neither |target| nor any of its
+// superclases must not have overridden |allocWithZone:|.
+void AddAllocWithZoneMethod(Class target, id (^impl_block)(Class, NSZone*));
+}  // namespace web
+
+#endif  // IOS_WEB_ALLOC_WITH_ZONE_INTERCEPTOR_H_
diff --git a/ios/web/alloc_with_zone_interceptor.mm b/ios/web/alloc_with_zone_interceptor.mm
new file mode 100644
index 0000000..ea86016
--- /dev/null
+++ b/ios/web/alloc_with_zone_interceptor.mm
@@ -0,0 +1,27 @@
+// Copyright 2015 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#import "ios/web/alloc_with_zone_interceptor.h"
+
+#import <objc/runtime.h>
+
+#include "base/logging.h"
+
+namespace web {
+
+void AddAllocWithZoneMethod(Class target, id (^impl_block)(Class, NSZone*)) {
+  // Make sure |allocWithZone:| is not already implemented in the target class.
+  Class meta_class = object_getClass(target);
+  DCHECK_EQ(
+      class_getMethodImplementation(meta_class, @selector(allocWithZone:)),
+      class_getMethodImplementation(object_getClass([NSObject class]),
+                                    @selector(allocWithZone:)));
+
+  IMP new_impl = imp_implementationWithBlock(^(id self, NSZone* zone) {
+    return impl_block(self, zone);
+  });
+  class_addMethod(meta_class, @selector(allocWithZone:), new_impl, "v@:@");
+}
+
+}  // namespace web
diff --git a/ios/web/alloc_with_zone_interceptor_unittest.mm b/ios/web/alloc_with_zone_interceptor_unittest.mm
new file mode 100644
index 0000000..6d2a5c9
--- /dev/null
+++ b/ios/web/alloc_with_zone_interceptor_unittest.mm
@@ -0,0 +1,40 @@
+// Copyright 2015 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#import "ios/web/alloc_with_zone_interceptor.h"
+
+#include "base/mac/scoped_nsobject.h"
+#include "testing/gtest/include/gtest/gtest.h"
+#include "testing/gtest_mac.h"
+#include "testing/platform_test.h"
+
+// Testable class to test web::AddAllocWithZoneMethod.
+@interface CRBTestAlloc : NSObject
+@end
+@implementation CRBTestAlloc
+@end
+
+namespace web {
+namespace {
+
+// Test fixture for alloc_with_zone_interceptor functions.
+typedef PlatformTest AllocWithZoneInterceptorTest;
+
+// Tests that AddAllocWithZoneMethod delegates allocation to the caller.
+TEST_F(AllocWithZoneInterceptorTest, AddAllocWithZoneMethod) {
+  NSZone* const kZone = NSDefaultMallocZone();
+  base::scoped_nsobject<CRBTestAlloc> original_result(
+      [CRBTestAlloc allocWithZone:kZone]);
+  ASSERT_TRUE([original_result isMemberOfClass:[CRBTestAlloc class]]);
+  NSString* const kResult = @"test-result";
+  AddAllocWithZoneMethod([CRBTestAlloc class], ^id(Class klass, NSZone* zone) {
+    EXPECT_EQ([CRBTestAlloc class], klass);
+    EXPECT_EQ(kZone, zone);
+    return kResult;
+  });
+  EXPECT_NSEQ(kResult, [CRBTestAlloc allocWithZone:kZone]);
+}
+
+}  // namespace
+}  // namespace web
diff --git a/ios/web/ios_web.gyp b/ios/web/ios_web.gyp
index 35ca42b..dc71de0 100644
--- a/ios/web/ios_web.gyp
+++ b/ios/web/ios_web.gyp
@@ -27,6 +27,8 @@
         '../../url/url.gyp:url_lib',
       ],
       'sources': [
+        'alloc_with_zone_interceptor.h',
+        'alloc_with_zone_interceptor.mm',
         'browser_state.cc',
         'load_committed_details.cc',
         'navigation/navigation_item_impl.h',
diff --git a/ios/web/ios_web_unittests.gyp b/ios/web/ios_web_unittests.gyp
index 8919547..7000f88 100644
--- a/ios/web/ios_web_unittests.gyp
+++ b/ios/web/ios_web_unittests.gyp
@@ -21,6 +21,7 @@
         'ios_web.gyp:test_support_ios_web',
       ],
       'sources': [
+        'alloc_with_zone_interceptor_unittest.mm',
         'browser_state_unittest.cc',
         'crw_network_activity_indicator_manager_unittest.mm',
         'history_state_util_unittest.mm',