1// Copyright 2017 The Chromium Authors
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5#include "partition_alloc/shim/allocator_interception_apple.h"
6
7#include "partition_alloc/partition_alloc_buildflags.h"
8
9#if BUILDFLAG(USE_ALLOCATOR_SHIM)
10#include <mach/mach.h>
11
12#include "partition_alloc/shim/allocator_shim.h"
13#include "partition_alloc/shim/malloc_zone_functions_apple.h"
14#include "testing/gtest/include/gtest/gtest.h"
15
16namespace allocator_shim {
17
18namespace {
19void ResetMallocZone(ChromeMallocZone* zone) {
20  MallocZoneFunctions& functions = GetFunctionsForZone(zone);
21  ReplaceZoneFunctions(zone, &functions);
22}
23
24void ResetAllMallocZones() {
25  ChromeMallocZone* default_malloc_zone =
26      reinterpret_cast<ChromeMallocZone*>(malloc_default_zone());
27  ResetMallocZone(default_malloc_zone);
28
29  vm_address_t* zones;
30  unsigned int count;
31  kern_return_t kr = malloc_get_all_zones(mach_task_self(), /*reader=*/nullptr,
32                                          &zones, &count);
33  if (kr != KERN_SUCCESS) {
34    return;
35  }
36  for (unsigned int i = 0; i < count; ++i) {
37    ChromeMallocZone* zone = reinterpret_cast<ChromeMallocZone*>(zones[i]);
38    ResetMallocZone(zone);
39  }
40}
41}  // namespace
42
43class AllocatorInterceptionTest : public testing::Test {
44 protected:
45  void TearDown() override {
46    ResetAllMallocZones();
47    ClearAllMallocZonesForTesting();
48  }
49};
50
51#if !defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
52TEST_F(AllocatorInterceptionTest, ShimNewMallocZones) {
53  InitializeAllocatorShim();
54  ChromeMallocZone* default_malloc_zone =
55      reinterpret_cast<ChromeMallocZone*>(malloc_default_zone());
56
57  malloc_zone_t new_zone;
58  // Uninitialized: this is not a problem, because `ShimNewMalloc()` will just
59  // copy a member from this struct, not use it, which is why its pointer in
60  // `new_zone` must be valid, but its content can be meaningless.
61  malloc_introspection_t introspection;
62  memset(&new_zone, 1, sizeof(malloc_zone_t));
63  new_zone.introspect = &introspection;
64  malloc_zone_register(&new_zone);
65  EXPECT_NE(new_zone.malloc, default_malloc_zone->malloc);
66  ShimNewMallocZones();
67  EXPECT_EQ(new_zone.malloc, default_malloc_zone->malloc);
68
69  malloc_zone_unregister(&new_zone);
70}
71#endif
72
73}  // namespace allocator_shim
74
75#endif  // BUILDFLAG(USE_ALLOCATOR_SHIM)
76