1 //===-- enable_disable.cpp --------------------------------------*- C++ -*-===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
9 #include "gwp_asan/tests/harness.h"
10 
11 constexpr size_t Size = 100;
12 
13 TEST_F(DefaultGuardedPoolAllocatorDeathTest, Fork) {
14   void *P;
15   pid_t Pid = fork();
16   EXPECT_GE(Pid, 0);
17   if (Pid == 0) {
18     P = GPA.allocate(Size);
19     EXPECT_NE(P, nullptr);
20     memset(P, 0x42, Size);
21     GPA.deallocate(P);
22     _exit(0);
23   }
24   waitpid(Pid, nullptr, 0);
25   P = GPA.allocate(Size);
26   EXPECT_NE(P, nullptr);
27   memset(P, 0x42, Size);
28   GPA.deallocate(P);
29 
30   // fork should stall if the allocator has been disabled.
31   EXPECT_DEATH(
32       {
33         GPA.disable();
34         alarm(1);
35         Pid = fork();
36         EXPECT_GE(Pid, 0);
37       },
38       "");
39 }
40 
41 namespace {
42 pthread_mutex_t Mutex;
43 pthread_cond_t Conditional = PTHREAD_COND_INITIALIZER;
44 bool ThreadReady = false;
45 
46 void *enableMalloc(void *arg) {
47   auto &GPA = *reinterpret_cast<gwp_asan::GuardedPoolAllocator *>(arg);
48 
49   // Signal the main thread we are ready.
50   pthread_mutex_lock(&Mutex);
51   ThreadReady = true;
52   pthread_cond_signal(&Conditional);
53   pthread_mutex_unlock(&Mutex);
54 
55   // Wait for the malloc_disable & fork, then enable the allocator again.
56   sleep(1);
57   GPA.enable();
58 
59   return nullptr;
60 }
61 
62 TEST_F(DefaultGuardedPoolAllocator, DisableForkEnable) {
63   pthread_t ThreadId;
64   EXPECT_EQ(pthread_create(&ThreadId, nullptr, &enableMalloc, &GPA), 0);
65 
66   // Do not lock the allocator right away, the other thread may need it to start
67   // up.
68   pthread_mutex_lock(&Mutex);
69   while (!ThreadReady)
70     pthread_cond_wait(&Conditional, &Mutex);
71   pthread_mutex_unlock(&Mutex);
72 
73   // Disable the allocator and fork. fork should succeed after malloc_enable.
74   GPA.disable();
75   pid_t Pid = fork();
76   EXPECT_GE(Pid, 0);
77   if (Pid == 0) {
78     void *P = GPA.allocate(Size);
79     EXPECT_NE(P, nullptr);
80     GPA.deallocate(P);
81     _exit(0);
82   }
83   waitpid(Pid, nullptr, 0);
84   EXPECT_EQ(pthread_join(ThreadId, 0), 0);
85 }
86 } // namespace
87