1 // Copyright 2022 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/address_pool_manager.h"
6 #include "partition_alloc/partition_alloc_buildflags.h"
7 #include "partition_alloc/partition_alloc_constants.h"
8 #include "partition_alloc/partition_root.h"
9 #include "partition_alloc/thread_isolation/thread_isolation.h"
10 
11 #if BUILDFLAG(ENABLE_PKEYS)
12 
13 #include <link.h>
14 #include <sys/mman.h>
15 #include <sys/syscall.h>
16 
17 #include "partition_alloc/address_space_stats.h"
18 #include "partition_alloc/page_allocator.h"
19 #include "partition_alloc/page_allocator_constants.h"
20 #include "partition_alloc/partition_alloc.h"
21 #include "partition_alloc/partition_alloc_base/no_destructor.h"
22 #include "partition_alloc/partition_alloc_forward.h"
23 #include "partition_alloc/thread_isolation/pkey.h"
24 #include "testing/gtest/include/gtest/gtest.h"
25 
26 #define ISOLATED_FUNCTION extern "C" __attribute__((used))
27 constexpr size_t kIsolatedThreadStackSize = 64 * 1024;
28 constexpr int kNumPkey = 16;
29 constexpr size_t kTestReturnValue = 0x8765432187654321llu;
30 constexpr uint32_t kPKRUAllowAccessNoWrite = 0b10101010101010101010101010101000;
31 
32 namespace partition_alloc::internal {
33 
34 struct PA_THREAD_ISOLATED_ALIGN IsolatedGlobals {
35   int pkey = kInvalidPkey;
36   void* stack;
37   partition_alloc::internal::base::NoDestructor<
38       partition_alloc::PartitionAllocator>
39       allocator{};
40 } isolated_globals;
41 
ProtFromSegmentFlags(ElfW (Word)flags)42 int ProtFromSegmentFlags(ElfW(Word) flags) {
43   int prot = 0;
44   if (flags & PF_R) {
45     prot |= PROT_READ;
46   }
47   if (flags & PF_W) {
48     prot |= PROT_WRITE;
49   }
50   if (flags & PF_X) {
51     prot |= PROT_EXEC;
52   }
53   return prot;
54 }
55 
ProtectROSegments(struct dl_phdr_info * info,size_t info_size,void * data)56 int ProtectROSegments(struct dl_phdr_info* info, size_t info_size, void* data) {
57   if (!strcmp(info->dlpi_name, "linux-vdso.so.1")) {
58     return 0;
59   }
60   for (int i = 0; i < info->dlpi_phnum; i++) {
61     const ElfW(Phdr)* phdr = &info->dlpi_phdr[i];
62     if (phdr->p_type != PT_LOAD && phdr->p_type != PT_GNU_RELRO) {
63       continue;
64     }
65     if (phdr->p_flags & PF_W) {
66       continue;
67     }
68     uintptr_t start = info->dlpi_addr + phdr->p_vaddr;
69     uintptr_t end = start + phdr->p_memsz;
70     uintptr_t start_page = RoundDownToSystemPage(start);
71     uintptr_t end_page = RoundUpToSystemPage(end);
72     uintptr_t size = end_page - start_page;
73     PA_PCHECK(PkeyMprotect(reinterpret_cast<void*>(start_page), size,
74                            ProtFromSegmentFlags(phdr->p_flags),
75                            isolated_globals.pkey) == 0);
76   }
77   return 0;
78 }
79 
80 class PkeyTest : public testing::Test {
81  protected:
PkeyProtectMemory()82   static void PkeyProtectMemory() {
83     PA_PCHECK(dl_iterate_phdr(ProtectROSegments, nullptr) == 0);
84 
85     PA_PCHECK(PkeyMprotect(&isolated_globals, sizeof(isolated_globals),
86                            PROT_READ | PROT_WRITE, isolated_globals.pkey) == 0);
87 
88     PA_PCHECK(PkeyMprotect(isolated_globals.stack, kIsolatedThreadStackSize,
89                            PROT_READ | PROT_WRITE, isolated_globals.pkey) == 0);
90   }
91 
InitializeIsolatedThread()92   static void InitializeIsolatedThread() {
93     isolated_globals.stack =
94         mmap(nullptr, kIsolatedThreadStackSize, PROT_READ | PROT_WRITE,
95              MAP_ANONYMOUS | MAP_PRIVATE | MAP_STACK, -1, 0);
96     PA_PCHECK(isolated_globals.stack != MAP_FAILED);
97 
98     PkeyProtectMemory();
99   }
100 
SetUp()101   void SetUp() override {
102     // SetUp only once, but we can't do it in SetUpTestSuite since that runs
103     // before other PartitionAlloc initialization happened.
104     if (isolated_globals.pkey != kInvalidPkey) {
105       return;
106     }
107 
108     int pkey = PkeyAlloc(0);
109     if (pkey == -1) {
110       return;
111     }
112     isolated_globals.pkey = pkey;
113 
114     isolated_globals.allocator->init([]() {
115       partition_alloc::PartitionOptions opts;
116       opts.thread_isolation = ThreadIsolationOption(isolated_globals.pkey);
117       return opts;
118     }());
119 
120     InitializeIsolatedThread();
121 
122     Wrpkru(kPKRUAllowAccessNoWrite);
123   }
124 
TearDownTestSuite()125   static void TearDownTestSuite() {
126     if (isolated_globals.pkey == kInvalidPkey) {
127       return;
128     }
129     PA_PCHECK(PkeyMprotect(&isolated_globals, sizeof(isolated_globals),
130                            PROT_READ | PROT_WRITE, kDefaultPkey) == 0);
131     isolated_globals.pkey = kDefaultPkey;
132     InitializeIsolatedThread();
133     PkeyFree(isolated_globals.pkey);
134   }
135 };
136 
137 // This code will run with access limited to pkey 1, no default pkey access.
138 // Note that we're stricter than required for debugging purposes.
139 // In the final use, we'll likely allow at least read access to the default
140 // pkey.
IsolatedAllocFree(void * arg)141 ISOLATED_FUNCTION uint64_t IsolatedAllocFree(void* arg) {
142   char* buf = (char*)isolated_globals.allocator->root()
143                   ->Alloc<partition_alloc::AllocFlags::kNoHooks>(1024);
144   if (!buf) {
145     return 0xffffffffffffffffllu;
146   }
147   isolated_globals.allocator->root()->Free<FreeFlags::kNoHooks>(buf);
148 
149   return kTestReturnValue;
150 }
151 
152 // This test is a bit compliated. We want to ensure that the code
153 // allocating/freeing from the pkey pool doesn't *unexpectedly* access memory
154 // tagged with the default pkey (pkey 0). This could be a security issue since
155 // in our CFI threat model that memory might be attacker controlled.
156 // To test for this, we run alloc/free without access to the default pkey. In
157 // order to do this, we need to tag all global read-only memory with our pkey as
158 // well as switch to a pkey-tagged stack.
TEST_F(PkeyTest,AllocWithoutDefaultPkey)159 TEST_F(PkeyTest, AllocWithoutDefaultPkey) {
160   if (isolated_globals.pkey == kInvalidPkey) {
161     return;
162   }
163 
164   uint64_t ret;
165   uint32_t pkru_value = 0;
166   for (int pkey = 0; pkey < kNumPkey; pkey++) {
167     if (pkey != isolated_globals.pkey) {
168       pkru_value |= (PKEY_DISABLE_ACCESS | PKEY_DISABLE_WRITE) << (2 * pkey);
169     }
170   }
171 
172   // Switch to the safe stack with inline assembly.
173   //
174   // The simple solution would be to use one asm statement as a prologue to
175   // switch to the protected stack and a second one to switch it back. However,
176   // that doesn't work since inline assembly doesn't support a clobbered stack
177   // register. So instead, we switch the stack, perform a function call
178   // to the
179   // actual code and switch back afterwards.
180   //
181   // The inline asm docs mention that special care must be taken
182   // when calling a function in inline assembly. I.e. we will
183   // need to make sure that we follow the ABI of the platform.
184   // In this example, we use the System-V ABI.
185   //
186   // == Caller-saved registers ==
187   // We had two ideas for handling caller-saved registers. Option 1 was chosen,
188   // but I'll describe both to show why option 2 didn't work out:
189   // * Option 1) mark all caller-saved registers as clobbered. This should be
190   //             in line with how the compiler would create the function call.
191   //             Problem: future additions to caller-saved registers can break
192   //             this.
193   // * Option 2) use attribute no_caller_saved_registers. This prohibits use of
194   //             sse/mmx/x87. We can disable sse/mmx with a "target" attribute,
195   //             but I couldn't find a way to disable x87.
196   //             The docs tell you to use -mgeneral-regs-only. Maybe we
197   //             could move the isolated code to a separate file and then
198   //             use that flag for compiling that file only.
199   //             !!! This doesn't work: the inner function can call out to code
200   //             that uses caller-saved registers and won't save
201   //             them itself.
202   //
203   // == stack alignment ==
204   // The ABI requires us to have a 16 byte aligned rsp on function
205   // entry. We push one qword onto the stack so we need to subtract
206   // an additional 8 bytes from the stack pointer.
207   //
208   // == additional clobbering ==
209   // As described above, we need to clobber everything besides
210   // callee-saved registers. The ABI requires all x87 registers to
211   // be set to empty on fn entry / return,
212   // so we should tell the compiler that this is the case. As I understand the
213   // docs, this is done by marking them as clobbered. Worst case, we'll notice
214   // any issues quickly and can fix them if it turned out to be false>
215   //
216   // == direction flag ==
217   // Theoretically, the DF flag could be set to 1 at asm entry. If this
218   // leads to problems, we might have to zero it before the fn call and
219   // restore it afterwards. I would'ave assumed that marking flags as
220   // clobbered would require the compiler to reset the DF before the next fn
221   // call, but that doesn't seem to be the case.
222   asm volatile(
223       // Set pkru to only allow access to pkey 1 memory.
224       ".byte 0x0f,0x01,0xef\n"  // wrpkru
225 
226       // Move to the isolated stack and store the old value
227       "xchg %4, %%rsp\n"
228       "push %4\n"
229       "call IsolatedAllocFree\n"
230       // We need rax below, so move the return value to the stack
231       "push %%rax\n"
232 
233       // Set pkru to only allow access to pkey 0 memory.
234       "mov $0b10101010101010101010101010101000, %%rax\n"
235       "xor %%rcx, %%rcx\n"
236       "xor %%rdx, %%rdx\n"
237       ".byte 0x0f,0x01,0xef\n"  // wrpkru
238 
239       // Pop the return value
240       "pop %0\n"
241       // Restore the original stack
242       "pop %%rsp\n"
243 
244       : "=r"(ret)
245       : "a"(pkru_value), "c"(0), "d"(0),
246         "r"(reinterpret_cast<uintptr_t>(isolated_globals.stack) +
247             kIsolatedThreadStackSize - 8)
248       : "memory", "cc", "r8", "r9", "r10", "r11", "xmm0", "xmm1", "xmm2",
249         "xmm3", "xmm4", "xmm5", "xmm6", "xmm7", "xmm8", "xmm9", "xmm10",
250         "xmm11", "xmm12", "xmm13", "xmm14", "xmm15", "flags", "fpsr", "st",
251         "st(1)", "st(2)", "st(3)", "st(4)", "st(5)", "st(6)", "st(7)");
252 
253   ASSERT_EQ(ret, kTestReturnValue);
254 }
255 
256 class MockAddressSpaceStatsDumper : public AddressSpaceStatsDumper {
257  public:
258   MockAddressSpaceStatsDumper() = default;
DumpStats(const AddressSpaceStats * address_space_stats)259   void DumpStats(const AddressSpaceStats* address_space_stats) override {}
260 };
261 
TEST_F(PkeyTest,DumpPkeyPoolStats)262 TEST_F(PkeyTest, DumpPkeyPoolStats) {
263   if (isolated_globals.pkey == kInvalidPkey) {
264     return;
265   }
266 
267   MockAddressSpaceStatsDumper mock_stats_dumper;
268   partition_alloc::internal::AddressPoolManager::GetInstance().DumpStats(
269       &mock_stats_dumper);
270 }
271 
272 }  // namespace partition_alloc::internal
273 
274 #endif  // BUILDFLAG(ENABLE_THREAD_ISOLATION)
275