1 // SPDX-License-Identifier: GPL-2.0
2 /*
3  * Tests Memory Protection Keys (see Documentation/core-api/protection-keys.rst)
4  *
5  * There are examples in here of:
6  *  * how to set protection keys on memory
7  *  * how to set/clear bits in pkey registers (the rights register)
8  *  * how to handle SEGV_PKUERR signals and extract pkey-relevant
9  *    information from the siginfo
10  *
11  * Things to add:
12  *	make sure KSM and KSM COW breaking works
13  *	prefault pages in at malloc, or not
14  *	protect MPX bounds tables with protection keys?
15  *	make sure VMA splitting/merging is working correctly
16  *	OOMs can destroy mm->mmap (see exit_mmap()), so make sure it is immune to pkeys
17  *	look for pkey "leaks" where it is still set on a VMA but "freed" back to the kernel
18  *	do a plain mprotect() to a mprotect_pkey() area and make sure the pkey sticks
19  *
20  * Compile like this:
21  *	gcc -mxsave      -o protection_keys    -O2 -g -std=gnu99 -pthread -Wall protection_keys.c -lrt -ldl -lm
22  *	gcc -mxsave -m32 -o protection_keys_32 -O2 -g -std=gnu99 -pthread -Wall protection_keys.c -lrt -ldl -lm
23  */
24 #define _GNU_SOURCE
25 #define __SANE_USERSPACE_TYPES__
26 #include <errno.h>
27 #include <linux/elf.h>
28 #include <linux/futex.h>
29 #include <time.h>
30 #include <sys/time.h>
31 #include <sys/syscall.h>
32 #include <string.h>
33 #include <stdio.h>
34 #include <stdint.h>
35 #include <stdbool.h>
36 #include <signal.h>
37 #include <assert.h>
38 #include <stdlib.h>
39 #include <ucontext.h>
40 #include <sys/mman.h>
41 #include <sys/types.h>
42 #include <sys/wait.h>
43 #include <sys/stat.h>
44 #include <fcntl.h>
45 #include <unistd.h>
46 #include <sys/ptrace.h>
47 #include <setjmp.h>
48 
49 #include "pkey-helpers.h"
50 
51 int iteration_nr = 1;
52 int test_nr;
53 
54 u64 shadow_pkey_reg;
55 int dprint_in_signal;
56 char dprint_in_signal_buffer[DPRINT_IN_SIGNAL_BUF_SIZE];
57 
58 void cat_into_file(char *str, char *file)
59 {
60 	int fd = open(file, O_RDWR);
61 	int ret;
62 
63 	dprintf2("%s(): writing '%s' to '%s'\n", __func__, str, file);
64 	/*
65 	 * these need to be raw because they are called under
66 	 * pkey_assert()
67 	 */
68 	if (fd < 0) {
69 		fprintf(stderr, "error opening '%s'\n", str);
70 		perror("error: ");
71 		exit(__LINE__);
72 	}
73 
74 	ret = write(fd, str, strlen(str));
75 	if (ret != strlen(str)) {
76 		perror("write to file failed");
77 		fprintf(stderr, "filename: '%s' str: '%s'\n", file, str);
78 		exit(__LINE__);
79 	}
80 	close(fd);
81 }
82 
83 #if CONTROL_TRACING > 0
84 static int warned_tracing;
85 int tracing_root_ok(void)
86 {
87 	if (geteuid() != 0) {
88 		if (!warned_tracing)
89 			fprintf(stderr, "WARNING: not run as root, "
90 					"can not do tracing control\n");
91 		warned_tracing = 1;
92 		return 0;
93 	}
94 	return 1;
95 }
96 #endif
97 
98 void tracing_on(void)
99 {
100 #if CONTROL_TRACING > 0
101 #define TRACEDIR "/sys/kernel/tracing"
102 	char pidstr[32];
103 
104 	if (!tracing_root_ok())
105 		return;
106 
107 	sprintf(pidstr, "%d", getpid());
108 	cat_into_file("0", TRACEDIR "/tracing_on");
109 	cat_into_file("\n", TRACEDIR "/trace");
110 	if (1) {
111 		cat_into_file("function_graph", TRACEDIR "/current_tracer");
112 		cat_into_file("1", TRACEDIR "/options/funcgraph-proc");
113 	} else {
114 		cat_into_file("nop", TRACEDIR "/current_tracer");
115 	}
116 	cat_into_file(pidstr, TRACEDIR "/set_ftrace_pid");
117 	cat_into_file("1", TRACEDIR "/tracing_on");
118 	dprintf1("enabled tracing\n");
119 #endif
120 }
121 
122 void tracing_off(void)
123 {
124 #if CONTROL_TRACING > 0
125 	if (!tracing_root_ok())
126 		return;
127 	cat_into_file("0", "/sys/kernel/tracing/tracing_on");
128 #endif
129 }
130 
131 void abort_hooks(void)
132 {
133 	fprintf(stderr, "running %s()...\n", __func__);
134 	tracing_off();
135 #ifdef SLEEP_ON_ABORT
136 	sleep(SLEEP_ON_ABORT);
137 #endif
138 }
139 
140 /*
141  * This attempts to have roughly a page of instructions followed by a few
142  * instructions that do a write, and another page of instructions.  That
143  * way, we are pretty sure that the write is in the second page of
144  * instructions and has at least a page of padding behind it.
145  *
146  * *That* lets us be sure to madvise() away the write instruction, which
147  * will then fault, which makes sure that the fault code handles
148  * execute-only memory properly.
149  */
150 #ifdef __powerpc64__
151 /* This way, both 4K and 64K alignment are maintained */
152 __attribute__((__aligned__(65536)))
153 #else
154 __attribute__((__aligned__(PAGE_SIZE)))
155 #endif
156 void lots_o_noops_around_write(int *write_to_me)
157 {
158 	dprintf3("running %s()\n", __func__);
159 	__page_o_noops();
160 	/* Assume this happens in the second page of instructions: */
161 	*write_to_me = __LINE__;
162 	/* pad out by another page: */
163 	__page_o_noops();
164 	dprintf3("%s() done\n", __func__);
165 }
166 
167 void dump_mem(void *dumpme, int len_bytes)
168 {
169 	char *c = (void *)dumpme;
170 	int i;
171 
172 	for (i = 0; i < len_bytes; i += sizeof(u64)) {
173 		u64 *ptr = (u64 *)(c + i);
174 		dprintf1("dump[%03d][@%p]: %016llx\n", i, ptr, *ptr);
175 	}
176 }
177 
178 static u32 hw_pkey_get(int pkey, unsigned long flags)
179 {
180 	u64 pkey_reg = __read_pkey_reg();
181 
182 	dprintf1("%s(pkey=%d, flags=%lx) = %x / %d\n",
183 			__func__, pkey, flags, 0, 0);
184 	dprintf2("%s() raw pkey_reg: %016llx\n", __func__, pkey_reg);
185 
186 	return (u32) get_pkey_bits(pkey_reg, pkey);
187 }
188 
189 static int hw_pkey_set(int pkey, unsigned long rights, unsigned long flags)
190 {
191 	u32 mask = (PKEY_DISABLE_ACCESS|PKEY_DISABLE_WRITE);
192 	u64 old_pkey_reg = __read_pkey_reg();
193 	u64 new_pkey_reg;
194 
195 	/* make sure that 'rights' only contains the bits we expect: */
196 	assert(!(rights & ~mask));
197 
198 	/* modify bits accordingly in old pkey_reg and assign it */
199 	new_pkey_reg = set_pkey_bits(old_pkey_reg, pkey, rights);
200 
201 	__write_pkey_reg(new_pkey_reg);
202 
203 	dprintf3("%s(pkey=%d, rights=%lx, flags=%lx) = %x"
204 		" pkey_reg now: %016llx old_pkey_reg: %016llx\n",
205 		__func__, pkey, rights, flags, 0, __read_pkey_reg(),
206 		old_pkey_reg);
207 	return 0;
208 }
209 
210 void pkey_disable_set(int pkey, int flags)
211 {
212 	unsigned long syscall_flags = 0;
213 	int ret;
214 	int pkey_rights;
215 	u64 orig_pkey_reg = read_pkey_reg();
216 
217 	dprintf1("START->%s(%d, 0x%x)\n", __func__,
218 		pkey, flags);
219 	pkey_assert(flags & (PKEY_DISABLE_ACCESS | PKEY_DISABLE_WRITE));
220 
221 	pkey_rights = hw_pkey_get(pkey, syscall_flags);
222 
223 	dprintf1("%s(%d) hw_pkey_get(%d): %x\n", __func__,
224 			pkey, pkey, pkey_rights);
225 
226 	pkey_assert(pkey_rights >= 0);
227 
228 	pkey_rights |= flags;
229 
230 	ret = hw_pkey_set(pkey, pkey_rights, syscall_flags);
231 	assert(!ret);
232 	/* pkey_reg and flags have the same format */
233 	shadow_pkey_reg = set_pkey_bits(shadow_pkey_reg, pkey, pkey_rights);
234 	dprintf1("%s(%d) shadow: 0x%016llx\n",
235 		__func__, pkey, shadow_pkey_reg);
236 
237 	pkey_assert(ret >= 0);
238 
239 	pkey_rights = hw_pkey_get(pkey, syscall_flags);
240 	dprintf1("%s(%d) hw_pkey_get(%d): %x\n", __func__,
241 			pkey, pkey, pkey_rights);
242 
243 	dprintf1("%s(%d) pkey_reg: 0x%016llx\n",
244 		__func__, pkey, read_pkey_reg());
245 	if (flags)
246 		pkey_assert(read_pkey_reg() >= orig_pkey_reg);
247 	dprintf1("END<---%s(%d, 0x%x)\n", __func__,
248 		pkey, flags);
249 }
250 
251 void pkey_disable_clear(int pkey, int flags)
252 {
253 	unsigned long syscall_flags = 0;
254 	int ret;
255 	int pkey_rights = hw_pkey_get(pkey, syscall_flags);
256 	u64 orig_pkey_reg = read_pkey_reg();
257 
258 	pkey_assert(flags & (PKEY_DISABLE_ACCESS | PKEY_DISABLE_WRITE));
259 
260 	dprintf1("%s(%d) hw_pkey_get(%d): %x\n", __func__,
261 			pkey, pkey, pkey_rights);
262 	pkey_assert(pkey_rights >= 0);
263 
264 	pkey_rights &= ~flags;
265 
266 	ret = hw_pkey_set(pkey, pkey_rights, 0);
267 	shadow_pkey_reg = set_pkey_bits(shadow_pkey_reg, pkey, pkey_rights);
268 	pkey_assert(ret >= 0);
269 
270 	pkey_rights = hw_pkey_get(pkey, syscall_flags);
271 	dprintf1("%s(%d) hw_pkey_get(%d): %x\n", __func__,
272 			pkey, pkey, pkey_rights);
273 
274 	dprintf1("%s(%d) pkey_reg: 0x%016llx\n", __func__,
275 			pkey, read_pkey_reg());
276 	if (flags)
277 		assert(read_pkey_reg() <= orig_pkey_reg);
278 }
279 
280 void pkey_write_allow(int pkey)
281 {
282 	pkey_disable_clear(pkey, PKEY_DISABLE_WRITE);
283 }
284 void pkey_write_deny(int pkey)
285 {
286 	pkey_disable_set(pkey, PKEY_DISABLE_WRITE);
287 }
288 void pkey_access_allow(int pkey)
289 {
290 	pkey_disable_clear(pkey, PKEY_DISABLE_ACCESS);
291 }
292 void pkey_access_deny(int pkey)
293 {
294 	pkey_disable_set(pkey, PKEY_DISABLE_ACCESS);
295 }
296 
297 static char *si_code_str(int si_code)
298 {
299 	if (si_code == SEGV_MAPERR)
300 		return "SEGV_MAPERR";
301 	if (si_code == SEGV_ACCERR)
302 		return "SEGV_ACCERR";
303 	if (si_code == SEGV_BNDERR)
304 		return "SEGV_BNDERR";
305 	if (si_code == SEGV_PKUERR)
306 		return "SEGV_PKUERR";
307 	return "UNKNOWN";
308 }
309 
310 int pkey_faults;
311 int last_si_pkey = -1;
312 void signal_handler(int signum, siginfo_t *si, void *vucontext)
313 {
314 	ucontext_t *uctxt = vucontext;
315 	int trapno;
316 	unsigned long ip;
317 	char *fpregs;
318 #if defined(__i386__) || defined(__x86_64__) /* arch */
319 	u32 *pkey_reg_ptr;
320 	int pkey_reg_offset;
321 #endif /* arch */
322 	u64 siginfo_pkey;
323 	u32 *si_pkey_ptr;
324 
325 	dprint_in_signal = 1;
326 	dprintf1(">>>>===============SIGSEGV============================\n");
327 	dprintf1("%s()::%d, pkey_reg: 0x%016llx shadow: %016llx\n",
328 			__func__, __LINE__,
329 			__read_pkey_reg(), shadow_pkey_reg);
330 
331 	trapno = uctxt->uc_mcontext.gregs[REG_TRAPNO];
332 	ip = uctxt->uc_mcontext.gregs[REG_IP_IDX];
333 	fpregs = (char *) uctxt->uc_mcontext.fpregs;
334 
335 	dprintf2("%s() trapno: %d ip: 0x%016lx info->si_code: %s/%d\n",
336 			__func__, trapno, ip, si_code_str(si->si_code),
337 			si->si_code);
338 
339 #if defined(__i386__) || defined(__x86_64__) /* arch */
340 #ifdef __i386__
341 	/*
342 	 * 32-bit has some extra padding so that userspace can tell whether
343 	 * the XSTATE header is present in addition to the "legacy" FPU
344 	 * state.  We just assume that it is here.
345 	 */
346 	fpregs += 0x70;
347 #endif /* i386 */
348 	pkey_reg_offset = pkey_reg_xstate_offset();
349 	pkey_reg_ptr = (void *)(&fpregs[pkey_reg_offset]);
350 
351 	/*
352 	 * If we got a PKEY fault, we *HAVE* to have at least one bit set in
353 	 * here.
354 	 */
355 	dprintf1("pkey_reg_xstate_offset: %d\n", pkey_reg_xstate_offset());
356 	if (DEBUG_LEVEL > 4)
357 		dump_mem(pkey_reg_ptr - 128, 256);
358 	pkey_assert(*pkey_reg_ptr);
359 #endif /* arch */
360 
361 	dprintf1("siginfo: %p\n", si);
362 	dprintf1(" fpregs: %p\n", fpregs);
363 
364 	if ((si->si_code == SEGV_MAPERR) ||
365 	    (si->si_code == SEGV_ACCERR) ||
366 	    (si->si_code == SEGV_BNDERR)) {
367 		printf("non-PK si_code, exiting...\n");
368 		exit(4);
369 	}
370 
371 	si_pkey_ptr = siginfo_get_pkey_ptr(si);
372 	dprintf1("si_pkey_ptr: %p\n", si_pkey_ptr);
373 	dump_mem((u8 *)si_pkey_ptr - 8, 24);
374 	siginfo_pkey = *si_pkey_ptr;
375 	pkey_assert(siginfo_pkey < NR_PKEYS);
376 	last_si_pkey = siginfo_pkey;
377 
378 	/*
379 	 * need __read_pkey_reg() version so we do not do shadow_pkey_reg
380 	 * checking
381 	 */
382 	dprintf1("signal pkey_reg from  pkey_reg: %016llx\n",
383 			__read_pkey_reg());
384 	dprintf1("pkey from siginfo: %016llx\n", siginfo_pkey);
385 #if defined(__i386__) || defined(__x86_64__) /* arch */
386 	dprintf1("signal pkey_reg from xsave: %08x\n", *pkey_reg_ptr);
387 	*(u64 *)pkey_reg_ptr = 0x00000000;
388 	dprintf1("WARNING: set PKEY_REG=0 to allow faulting instruction to continue\n");
389 #elif defined(__powerpc64__) /* arch */
390 	/* restore access and let the faulting instruction continue */
391 	pkey_access_allow(siginfo_pkey);
392 #endif /* arch */
393 	pkey_faults++;
394 	dprintf1("<<<<==================================================\n");
395 	dprint_in_signal = 0;
396 }
397 
398 int wait_all_children(void)
399 {
400 	int status;
401 	return waitpid(-1, &status, 0);
402 }
403 
404 void sig_chld(int x)
405 {
406 	dprint_in_signal = 1;
407 	dprintf2("[%d] SIGCHLD: %d\n", getpid(), x);
408 	dprint_in_signal = 0;
409 }
410 
411 void setup_sigsegv_handler(void)
412 {
413 	int r, rs;
414 	struct sigaction newact;
415 	struct sigaction oldact;
416 
417 	/* #PF is mapped to sigsegv */
418 	int signum  = SIGSEGV;
419 
420 	newact.sa_handler = 0;
421 	newact.sa_sigaction = signal_handler;
422 
423 	/*sigset_t - signals to block while in the handler */
424 	/* get the old signal mask. */
425 	rs = sigprocmask(SIG_SETMASK, 0, &newact.sa_mask);
426 	pkey_assert(rs == 0);
427 
428 	/* call sa_sigaction, not sa_handler*/
429 	newact.sa_flags = SA_SIGINFO;
430 
431 	newact.sa_restorer = 0;  /* void(*)(), obsolete */
432 	r = sigaction(signum, &newact, &oldact);
433 	r = sigaction(SIGALRM, &newact, &oldact);
434 	pkey_assert(r == 0);
435 }
436 
437 void setup_handlers(void)
438 {
439 	signal(SIGCHLD, &sig_chld);
440 	setup_sigsegv_handler();
441 }
442 
443 pid_t fork_lazy_child(void)
444 {
445 	pid_t forkret;
446 
447 	forkret = fork();
448 	pkey_assert(forkret >= 0);
449 	dprintf3("[%d] fork() ret: %d\n", getpid(), forkret);
450 
451 	if (!forkret) {
452 		/* in the child */
453 		while (1) {
454 			dprintf1("child sleeping...\n");
455 			sleep(30);
456 		}
457 	}
458 	return forkret;
459 }
460 
461 int sys_mprotect_pkey(void *ptr, size_t size, unsigned long orig_prot,
462 		unsigned long pkey)
463 {
464 	int sret;
465 
466 	dprintf2("%s(0x%p, %zx, prot=%lx, pkey=%lx)\n", __func__,
467 			ptr, size, orig_prot, pkey);
468 
469 	errno = 0;
470 	sret = syscall(__NR_pkey_mprotect, ptr, size, orig_prot, pkey);
471 	if (errno) {
472 		dprintf2("SYS_mprotect_key sret: %d\n", sret);
473 		dprintf2("SYS_mprotect_key prot: 0x%lx\n", orig_prot);
474 		dprintf2("SYS_mprotect_key failed, errno: %d\n", errno);
475 		if (DEBUG_LEVEL >= 2)
476 			perror("SYS_mprotect_pkey");
477 	}
478 	return sret;
479 }
480 
481 int sys_pkey_alloc(unsigned long flags, unsigned long init_val)
482 {
483 	int ret = syscall(SYS_pkey_alloc, flags, init_val);
484 	dprintf1("%s(flags=%lx, init_val=%lx) syscall ret: %d errno: %d\n",
485 			__func__, flags, init_val, ret, errno);
486 	return ret;
487 }
488 
489 int alloc_pkey(void)
490 {
491 	int ret;
492 	unsigned long init_val = 0x0;
493 
494 	dprintf1("%s()::%d, pkey_reg: 0x%016llx shadow: %016llx\n",
495 			__func__, __LINE__, __read_pkey_reg(), shadow_pkey_reg);
496 	ret = sys_pkey_alloc(0, init_val);
497 	/*
498 	 * pkey_alloc() sets PKEY register, so we need to reflect it in
499 	 * shadow_pkey_reg:
500 	 */
501 	dprintf4("%s()::%d, ret: %d pkey_reg: 0x%016llx"
502 			" shadow: 0x%016llx\n",
503 			__func__, __LINE__, ret, __read_pkey_reg(),
504 			shadow_pkey_reg);
505 	if (ret > 0) {
506 		/* clear both the bits: */
507 		shadow_pkey_reg = set_pkey_bits(shadow_pkey_reg, ret,
508 						~PKEY_MASK);
509 		dprintf4("%s()::%d, ret: %d pkey_reg: 0x%016llx"
510 				" shadow: 0x%016llx\n",
511 				__func__,
512 				__LINE__, ret, __read_pkey_reg(),
513 				shadow_pkey_reg);
514 		/*
515 		 * move the new state in from init_val
516 		 * (remember, we cheated and init_val == pkey_reg format)
517 		 */
518 		shadow_pkey_reg = set_pkey_bits(shadow_pkey_reg, ret,
519 						init_val);
520 	}
521 	dprintf4("%s()::%d, ret: %d pkey_reg: 0x%016llx"
522 			" shadow: 0x%016llx\n",
523 			__func__, __LINE__, ret, __read_pkey_reg(),
524 			shadow_pkey_reg);
525 	dprintf1("%s()::%d errno: %d\n", __func__, __LINE__, errno);
526 	/* for shadow checking: */
527 	read_pkey_reg();
528 	dprintf4("%s()::%d, ret: %d pkey_reg: 0x%016llx"
529 		 " shadow: 0x%016llx\n",
530 		__func__, __LINE__, ret, __read_pkey_reg(),
531 		shadow_pkey_reg);
532 	return ret;
533 }
534 
535 int sys_pkey_free(unsigned long pkey)
536 {
537 	int ret = syscall(SYS_pkey_free, pkey);
538 	dprintf1("%s(pkey=%ld) syscall ret: %d\n", __func__, pkey, ret);
539 	return ret;
540 }
541 
542 /*
543  * I had a bug where pkey bits could be set by mprotect() but
544  * not cleared.  This ensures we get lots of random bit sets
545  * and clears on the vma and pte pkey bits.
546  */
547 int alloc_random_pkey(void)
548 {
549 	int max_nr_pkey_allocs;
550 	int ret;
551 	int i;
552 	int alloced_pkeys[NR_PKEYS];
553 	int nr_alloced = 0;
554 	int random_index;
555 	memset(alloced_pkeys, 0, sizeof(alloced_pkeys));
556 
557 	/* allocate every possible key and make a note of which ones we got */
558 	max_nr_pkey_allocs = NR_PKEYS;
559 	for (i = 0; i < max_nr_pkey_allocs; i++) {
560 		int new_pkey = alloc_pkey();
561 		if (new_pkey < 0)
562 			break;
563 		alloced_pkeys[nr_alloced++] = new_pkey;
564 	}
565 
566 	pkey_assert(nr_alloced > 0);
567 	/* select a random one out of the allocated ones */
568 	random_index = rand() % nr_alloced;
569 	ret = alloced_pkeys[random_index];
570 	/* now zero it out so we don't free it next */
571 	alloced_pkeys[random_index] = 0;
572 
573 	/* go through the allocated ones that we did not want and free them */
574 	for (i = 0; i < nr_alloced; i++) {
575 		int free_ret;
576 		if (!alloced_pkeys[i])
577 			continue;
578 		free_ret = sys_pkey_free(alloced_pkeys[i]);
579 		pkey_assert(!free_ret);
580 	}
581 	dprintf1("%s()::%d, ret: %d pkey_reg: 0x%016llx"
582 			 " shadow: 0x%016llx\n", __func__,
583 			__LINE__, ret, __read_pkey_reg(), shadow_pkey_reg);
584 	return ret;
585 }
586 
587 int mprotect_pkey(void *ptr, size_t size, unsigned long orig_prot,
588 		unsigned long pkey)
589 {
590 	int nr_iterations = random() % 100;
591 	int ret;
592 
593 	while (0) {
594 		int rpkey = alloc_random_pkey();
595 		ret = sys_mprotect_pkey(ptr, size, orig_prot, pkey);
596 		dprintf1("sys_mprotect_pkey(%p, %zx, prot=0x%lx, pkey=%ld) ret: %d\n",
597 				ptr, size, orig_prot, pkey, ret);
598 		if (nr_iterations-- < 0)
599 			break;
600 
601 		dprintf1("%s()::%d, ret: %d pkey_reg: 0x%016llx"
602 			" shadow: 0x%016llx\n",
603 			__func__, __LINE__, ret, __read_pkey_reg(),
604 			shadow_pkey_reg);
605 		sys_pkey_free(rpkey);
606 		dprintf1("%s()::%d, ret: %d pkey_reg: 0x%016llx"
607 			" shadow: 0x%016llx\n",
608 			__func__, __LINE__, ret, __read_pkey_reg(),
609 			shadow_pkey_reg);
610 	}
611 	pkey_assert(pkey < NR_PKEYS);
612 
613 	ret = sys_mprotect_pkey(ptr, size, orig_prot, pkey);
614 	dprintf1("mprotect_pkey(%p, %zx, prot=0x%lx, pkey=%ld) ret: %d\n",
615 			ptr, size, orig_prot, pkey, ret);
616 	pkey_assert(!ret);
617 	dprintf1("%s()::%d, ret: %d pkey_reg: 0x%016llx"
618 			" shadow: 0x%016llx\n", __func__,
619 			__LINE__, ret, __read_pkey_reg(), shadow_pkey_reg);
620 	return ret;
621 }
622 
623 struct pkey_malloc_record {
624 	void *ptr;
625 	long size;
626 	int prot;
627 };
628 struct pkey_malloc_record *pkey_malloc_records;
629 struct pkey_malloc_record *pkey_last_malloc_record;
630 long nr_pkey_malloc_records;
631 void record_pkey_malloc(void *ptr, long size, int prot)
632 {
633 	long i;
634 	struct pkey_malloc_record *rec = NULL;
635 
636 	for (i = 0; i < nr_pkey_malloc_records; i++) {
637 		rec = &pkey_malloc_records[i];
638 		/* find a free record */
639 		if (rec)
640 			break;
641 	}
642 	if (!rec) {
643 		/* every record is full */
644 		size_t old_nr_records = nr_pkey_malloc_records;
645 		size_t new_nr_records = (nr_pkey_malloc_records * 2 + 1);
646 		size_t new_size = new_nr_records * sizeof(struct pkey_malloc_record);
647 		dprintf2("new_nr_records: %zd\n", new_nr_records);
648 		dprintf2("new_size: %zd\n", new_size);
649 		pkey_malloc_records = realloc(pkey_malloc_records, new_size);
650 		pkey_assert(pkey_malloc_records != NULL);
651 		rec = &pkey_malloc_records[nr_pkey_malloc_records];
652 		/*
653 		 * realloc() does not initialize memory, so zero it from
654 		 * the first new record all the way to the end.
655 		 */
656 		for (i = 0; i < new_nr_records - old_nr_records; i++)
657 			memset(rec + i, 0, sizeof(*rec));
658 	}
659 	dprintf3("filling malloc record[%d/%p]: {%p, %ld}\n",
660 		(int)(rec - pkey_malloc_records), rec, ptr, size);
661 	rec->ptr = ptr;
662 	rec->size = size;
663 	rec->prot = prot;
664 	pkey_last_malloc_record = rec;
665 	nr_pkey_malloc_records++;
666 }
667 
668 void free_pkey_malloc(void *ptr)
669 {
670 	long i;
671 	int ret;
672 	dprintf3("%s(%p)\n", __func__, ptr);
673 	for (i = 0; i < nr_pkey_malloc_records; i++) {
674 		struct pkey_malloc_record *rec = &pkey_malloc_records[i];
675 		dprintf4("looking for ptr %p at record[%ld/%p]: {%p, %ld}\n",
676 				ptr, i, rec, rec->ptr, rec->size);
677 		if ((ptr <  rec->ptr) ||
678 		    (ptr >= rec->ptr + rec->size))
679 			continue;
680 
681 		dprintf3("found ptr %p at record[%ld/%p]: {%p, %ld}\n",
682 				ptr, i, rec, rec->ptr, rec->size);
683 		nr_pkey_malloc_records--;
684 		ret = munmap(rec->ptr, rec->size);
685 		dprintf3("munmap ret: %d\n", ret);
686 		pkey_assert(!ret);
687 		dprintf3("clearing rec->ptr, rec: %p\n", rec);
688 		rec->ptr = NULL;
689 		dprintf3("done clearing rec->ptr, rec: %p\n", rec);
690 		return;
691 	}
692 	pkey_assert(false);
693 }
694 
695 
696 void *malloc_pkey_with_mprotect(long size, int prot, u16 pkey)
697 {
698 	void *ptr;
699 	int ret;
700 
701 	read_pkey_reg();
702 	dprintf1("doing %s(size=%ld, prot=0x%x, pkey=%d)\n", __func__,
703 			size, prot, pkey);
704 	pkey_assert(pkey < NR_PKEYS);
705 	ptr = mmap(NULL, size, prot, MAP_ANONYMOUS|MAP_PRIVATE, -1, 0);
706 	pkey_assert(ptr != (void *)-1);
707 	ret = mprotect_pkey((void *)ptr, PAGE_SIZE, prot, pkey);
708 	pkey_assert(!ret);
709 	record_pkey_malloc(ptr, size, prot);
710 	read_pkey_reg();
711 
712 	dprintf1("%s() for pkey %d @ %p\n", __func__, pkey, ptr);
713 	return ptr;
714 }
715 
716 void *malloc_pkey_anon_huge(long size, int prot, u16 pkey)
717 {
718 	int ret;
719 	void *ptr;
720 
721 	dprintf1("doing %s(size=%ld, prot=0x%x, pkey=%d)\n", __func__,
722 			size, prot, pkey);
723 	/*
724 	 * Guarantee we can fit at least one huge page in the resulting
725 	 * allocation by allocating space for 2:
726 	 */
727 	size = ALIGN_UP(size, HPAGE_SIZE * 2);
728 	ptr = mmap(NULL, size, PROT_NONE, MAP_ANONYMOUS|MAP_PRIVATE, -1, 0);
729 	pkey_assert(ptr != (void *)-1);
730 	record_pkey_malloc(ptr, size, prot);
731 	mprotect_pkey(ptr, size, prot, pkey);
732 
733 	dprintf1("unaligned ptr: %p\n", ptr);
734 	ptr = ALIGN_PTR_UP(ptr, HPAGE_SIZE);
735 	dprintf1("  aligned ptr: %p\n", ptr);
736 	ret = madvise(ptr, HPAGE_SIZE, MADV_HUGEPAGE);
737 	dprintf1("MADV_HUGEPAGE ret: %d\n", ret);
738 	ret = madvise(ptr, HPAGE_SIZE, MADV_WILLNEED);
739 	dprintf1("MADV_WILLNEED ret: %d\n", ret);
740 	memset(ptr, 0, HPAGE_SIZE);
741 
742 	dprintf1("mmap()'d thp for pkey %d @ %p\n", pkey, ptr);
743 	return ptr;
744 }
745 
746 int hugetlb_setup_ok;
747 #define SYSFS_FMT_NR_HUGE_PAGES "/sys/kernel/mm/hugepages/hugepages-%ldkB/nr_hugepages"
748 #define GET_NR_HUGE_PAGES 10
749 void setup_hugetlbfs(void)
750 {
751 	int err;
752 	int fd;
753 	char buf[256];
754 	long hpagesz_kb;
755 	long hpagesz_mb;
756 
757 	if (geteuid() != 0) {
758 		fprintf(stderr, "WARNING: not run as root, can not do hugetlb test\n");
759 		return;
760 	}
761 
762 	cat_into_file(__stringify(GET_NR_HUGE_PAGES), "/proc/sys/vm/nr_hugepages");
763 
764 	/*
765 	 * Now go make sure that we got the pages and that they
766 	 * are PMD-level pages. Someone might have made PUD-level
767 	 * pages the default.
768 	 */
769 	hpagesz_kb = HPAGE_SIZE / 1024;
770 	hpagesz_mb = hpagesz_kb / 1024;
771 	sprintf(buf, SYSFS_FMT_NR_HUGE_PAGES, hpagesz_kb);
772 	fd = open(buf, O_RDONLY);
773 	if (fd < 0) {
774 		fprintf(stderr, "opening sysfs %ldM hugetlb config: %s\n",
775 			hpagesz_mb, strerror(errno));
776 		return;
777 	}
778 
779 	/* -1 to guarantee leaving the trailing \0 */
780 	err = read(fd, buf, sizeof(buf)-1);
781 	close(fd);
782 	if (err <= 0) {
783 		fprintf(stderr, "reading sysfs %ldM hugetlb config: %s\n",
784 			hpagesz_mb, strerror(errno));
785 		return;
786 	}
787 
788 	if (atoi(buf) != GET_NR_HUGE_PAGES) {
789 		fprintf(stderr, "could not confirm %ldM pages, got: '%s' expected %d\n",
790 			hpagesz_mb, buf, GET_NR_HUGE_PAGES);
791 		return;
792 	}
793 
794 	hugetlb_setup_ok = 1;
795 }
796 
797 void *malloc_pkey_hugetlb(long size, int prot, u16 pkey)
798 {
799 	void *ptr;
800 	int flags = MAP_ANONYMOUS|MAP_PRIVATE|MAP_HUGETLB;
801 
802 	if (!hugetlb_setup_ok)
803 		return PTR_ERR_ENOTSUP;
804 
805 	dprintf1("doing %s(%ld, %x, %x)\n", __func__, size, prot, pkey);
806 	size = ALIGN_UP(size, HPAGE_SIZE * 2);
807 	pkey_assert(pkey < NR_PKEYS);
808 	ptr = mmap(NULL, size, PROT_NONE, flags, -1, 0);
809 	pkey_assert(ptr != (void *)-1);
810 	mprotect_pkey(ptr, size, prot, pkey);
811 
812 	record_pkey_malloc(ptr, size, prot);
813 
814 	dprintf1("mmap()'d hugetlbfs for pkey %d @ %p\n", pkey, ptr);
815 	return ptr;
816 }
817 
818 void *malloc_pkey_mmap_dax(long size, int prot, u16 pkey)
819 {
820 	void *ptr;
821 	int fd;
822 
823 	dprintf1("doing %s(size=%ld, prot=0x%x, pkey=%d)\n", __func__,
824 			size, prot, pkey);
825 	pkey_assert(pkey < NR_PKEYS);
826 	fd = open("/dax/foo", O_RDWR);
827 	pkey_assert(fd >= 0);
828 
829 	ptr = mmap(0, size, prot, MAP_SHARED, fd, 0);
830 	pkey_assert(ptr != (void *)-1);
831 
832 	mprotect_pkey(ptr, size, prot, pkey);
833 
834 	record_pkey_malloc(ptr, size, prot);
835 
836 	dprintf1("mmap()'d for pkey %d @ %p\n", pkey, ptr);
837 	close(fd);
838 	return ptr;
839 }
840 
841 void *(*pkey_malloc[])(long size, int prot, u16 pkey) = {
842 
843 	malloc_pkey_with_mprotect,
844 	malloc_pkey_with_mprotect_subpage,
845 	malloc_pkey_anon_huge,
846 	malloc_pkey_hugetlb
847 /* can not do direct with the pkey_mprotect() API:
848 	malloc_pkey_mmap_direct,
849 	malloc_pkey_mmap_dax,
850 */
851 };
852 
853 void *malloc_pkey(long size, int prot, u16 pkey)
854 {
855 	void *ret;
856 	static int malloc_type;
857 	int nr_malloc_types = ARRAY_SIZE(pkey_malloc);
858 
859 	pkey_assert(pkey < NR_PKEYS);
860 
861 	while (1) {
862 		pkey_assert(malloc_type < nr_malloc_types);
863 
864 		ret = pkey_malloc[malloc_type](size, prot, pkey);
865 		pkey_assert(ret != (void *)-1);
866 
867 		malloc_type++;
868 		if (malloc_type >= nr_malloc_types)
869 			malloc_type = (random()%nr_malloc_types);
870 
871 		/* try again if the malloc_type we tried is unsupported */
872 		if (ret == PTR_ERR_ENOTSUP)
873 			continue;
874 
875 		break;
876 	}
877 
878 	dprintf3("%s(%ld, prot=%x, pkey=%x) returning: %p\n", __func__,
879 			size, prot, pkey, ret);
880 	return ret;
881 }
882 
883 int last_pkey_faults;
884 #define UNKNOWN_PKEY -2
885 void expected_pkey_fault(int pkey)
886 {
887 	dprintf2("%s(): last_pkey_faults: %d pkey_faults: %d\n",
888 			__func__, last_pkey_faults, pkey_faults);
889 	dprintf2("%s(%d): last_si_pkey: %d\n", __func__, pkey, last_si_pkey);
890 	pkey_assert(last_pkey_faults + 1 == pkey_faults);
891 
892        /*
893 	* For exec-only memory, we do not know the pkey in
894 	* advance, so skip this check.
895 	*/
896 	if (pkey != UNKNOWN_PKEY)
897 		pkey_assert(last_si_pkey == pkey);
898 
899 #if defined(__i386__) || defined(__x86_64__) /* arch */
900 	/*
901 	 * The signal handler shold have cleared out PKEY register to let the
902 	 * test program continue.  We now have to restore it.
903 	 */
904 	if (__read_pkey_reg() != 0)
905 #else /* arch */
906 	if (__read_pkey_reg() != shadow_pkey_reg)
907 #endif /* arch */
908 		pkey_assert(0);
909 
910 	__write_pkey_reg(shadow_pkey_reg);
911 	dprintf1("%s() set pkey_reg=%016llx to restore state after signal "
912 		       "nuked it\n", __func__, shadow_pkey_reg);
913 	last_pkey_faults = pkey_faults;
914 	last_si_pkey = -1;
915 }
916 
917 #define do_not_expect_pkey_fault(msg)	do {			\
918 	if (last_pkey_faults != pkey_faults)			\
919 		dprintf0("unexpected PKey fault: %s\n", msg);	\
920 	pkey_assert(last_pkey_faults == pkey_faults);		\
921 } while (0)
922 
923 int test_fds[10] = { -1 };
924 int nr_test_fds;
925 void __save_test_fd(int fd)
926 {
927 	pkey_assert(fd >= 0);
928 	pkey_assert(nr_test_fds < ARRAY_SIZE(test_fds));
929 	test_fds[nr_test_fds] = fd;
930 	nr_test_fds++;
931 }
932 
933 int get_test_read_fd(void)
934 {
935 	int test_fd = open("/etc/passwd", O_RDONLY);
936 	__save_test_fd(test_fd);
937 	return test_fd;
938 }
939 
940 void close_test_fds(void)
941 {
942 	int i;
943 
944 	for (i = 0; i < nr_test_fds; i++) {
945 		if (test_fds[i] < 0)
946 			continue;
947 		close(test_fds[i]);
948 		test_fds[i] = -1;
949 	}
950 	nr_test_fds = 0;
951 }
952 
953 #define barrier() __asm__ __volatile__("": : :"memory")
954 __attribute__((noinline)) int read_ptr(int *ptr)
955 {
956 	/*
957 	 * Keep GCC from optimizing this away somehow
958 	 */
959 	barrier();
960 	return *ptr;
961 }
962 
963 void test_pkey_alloc_free_attach_pkey0(int *ptr, u16 pkey)
964 {
965 	int i, err;
966 	int max_nr_pkey_allocs;
967 	int alloced_pkeys[NR_PKEYS];
968 	int nr_alloced = 0;
969 	long size;
970 
971 	pkey_assert(pkey_last_malloc_record);
972 	size = pkey_last_malloc_record->size;
973 	/*
974 	 * This is a bit of a hack.  But mprotect() requires
975 	 * huge-page-aligned sizes when operating on hugetlbfs.
976 	 * So, make sure that we use something that's a multiple
977 	 * of a huge page when we can.
978 	 */
979 	if (size >= HPAGE_SIZE)
980 		size = HPAGE_SIZE;
981 
982 	/* allocate every possible key and make sure key-0 never got allocated */
983 	max_nr_pkey_allocs = NR_PKEYS;
984 	for (i = 0; i < max_nr_pkey_allocs; i++) {
985 		int new_pkey = alloc_pkey();
986 		pkey_assert(new_pkey != 0);
987 
988 		if (new_pkey < 0)
989 			break;
990 		alloced_pkeys[nr_alloced++] = new_pkey;
991 	}
992 	/* free all the allocated keys */
993 	for (i = 0; i < nr_alloced; i++) {
994 		int free_ret;
995 
996 		if (!alloced_pkeys[i])
997 			continue;
998 		free_ret = sys_pkey_free(alloced_pkeys[i]);
999 		pkey_assert(!free_ret);
1000 	}
1001 
1002 	/* attach key-0 in various modes */
1003 	err = sys_mprotect_pkey(ptr, size, PROT_READ, 0);
1004 	pkey_assert(!err);
1005 	err = sys_mprotect_pkey(ptr, size, PROT_WRITE, 0);
1006 	pkey_assert(!err);
1007 	err = sys_mprotect_pkey(ptr, size, PROT_EXEC, 0);
1008 	pkey_assert(!err);
1009 	err = sys_mprotect_pkey(ptr, size, PROT_READ|PROT_WRITE, 0);
1010 	pkey_assert(!err);
1011 	err = sys_mprotect_pkey(ptr, size, PROT_READ|PROT_WRITE|PROT_EXEC, 0);
1012 	pkey_assert(!err);
1013 }
1014 
1015 void test_read_of_write_disabled_region(int *ptr, u16 pkey)
1016 {
1017 	int ptr_contents;
1018 
1019 	dprintf1("disabling write access to PKEY[1], doing read\n");
1020 	pkey_write_deny(pkey);
1021 	ptr_contents = read_ptr(ptr);
1022 	dprintf1("*ptr: %d\n", ptr_contents);
1023 	dprintf1("\n");
1024 }
1025 void test_read_of_access_disabled_region(int *ptr, u16 pkey)
1026 {
1027 	int ptr_contents;
1028 
1029 	dprintf1("disabling access to PKEY[%02d], doing read @ %p\n", pkey, ptr);
1030 	read_pkey_reg();
1031 	pkey_access_deny(pkey);
1032 	ptr_contents = read_ptr(ptr);
1033 	dprintf1("*ptr: %d\n", ptr_contents);
1034 	expected_pkey_fault(pkey);
1035 }
1036 
1037 void test_read_of_access_disabled_region_with_page_already_mapped(int *ptr,
1038 		u16 pkey)
1039 {
1040 	int ptr_contents;
1041 
1042 	dprintf1("disabling access to PKEY[%02d], doing read @ %p\n",
1043 				pkey, ptr);
1044 	ptr_contents = read_ptr(ptr);
1045 	dprintf1("reading ptr before disabling the read : %d\n",
1046 			ptr_contents);
1047 	read_pkey_reg();
1048 	pkey_access_deny(pkey);
1049 	ptr_contents = read_ptr(ptr);
1050 	dprintf1("*ptr: %d\n", ptr_contents);
1051 	expected_pkey_fault(pkey);
1052 }
1053 
1054 void test_write_of_write_disabled_region_with_page_already_mapped(int *ptr,
1055 		u16 pkey)
1056 {
1057 	*ptr = __LINE__;
1058 	dprintf1("disabling write access; after accessing the page, "
1059 		"to PKEY[%02d], doing write\n", pkey);
1060 	pkey_write_deny(pkey);
1061 	*ptr = __LINE__;
1062 	expected_pkey_fault(pkey);
1063 }
1064 
1065 void test_write_of_write_disabled_region(int *ptr, u16 pkey)
1066 {
1067 	dprintf1("disabling write access to PKEY[%02d], doing write\n", pkey);
1068 	pkey_write_deny(pkey);
1069 	*ptr = __LINE__;
1070 	expected_pkey_fault(pkey);
1071 }
1072 void test_write_of_access_disabled_region(int *ptr, u16 pkey)
1073 {
1074 	dprintf1("disabling access to PKEY[%02d], doing write\n", pkey);
1075 	pkey_access_deny(pkey);
1076 	*ptr = __LINE__;
1077 	expected_pkey_fault(pkey);
1078 }
1079 
1080 void test_write_of_access_disabled_region_with_page_already_mapped(int *ptr,
1081 			u16 pkey)
1082 {
1083 	*ptr = __LINE__;
1084 	dprintf1("disabling access; after accessing the page, "
1085 		" to PKEY[%02d], doing write\n", pkey);
1086 	pkey_access_deny(pkey);
1087 	*ptr = __LINE__;
1088 	expected_pkey_fault(pkey);
1089 }
1090 
1091 void test_kernel_write_of_access_disabled_region(int *ptr, u16 pkey)
1092 {
1093 	int ret;
1094 	int test_fd = get_test_read_fd();
1095 
1096 	dprintf1("disabling access to PKEY[%02d], "
1097 		 "having kernel read() to buffer\n", pkey);
1098 	pkey_access_deny(pkey);
1099 	ret = read(test_fd, ptr, 1);
1100 	dprintf1("read ret: %d\n", ret);
1101 	pkey_assert(ret);
1102 }
1103 void test_kernel_write_of_write_disabled_region(int *ptr, u16 pkey)
1104 {
1105 	int ret;
1106 	int test_fd = get_test_read_fd();
1107 
1108 	pkey_write_deny(pkey);
1109 	ret = read(test_fd, ptr, 100);
1110 	dprintf1("read ret: %d\n", ret);
1111 	if (ret < 0 && (DEBUG_LEVEL > 0))
1112 		perror("verbose read result (OK for this to be bad)");
1113 	pkey_assert(ret);
1114 }
1115 
1116 void test_kernel_gup_of_access_disabled_region(int *ptr, u16 pkey)
1117 {
1118 	int pipe_ret, vmsplice_ret;
1119 	struct iovec iov;
1120 	int pipe_fds[2];
1121 
1122 	pipe_ret = pipe(pipe_fds);
1123 
1124 	pkey_assert(pipe_ret == 0);
1125 	dprintf1("disabling access to PKEY[%02d], "
1126 		 "having kernel vmsplice from buffer\n", pkey);
1127 	pkey_access_deny(pkey);
1128 	iov.iov_base = ptr;
1129 	iov.iov_len = PAGE_SIZE;
1130 	vmsplice_ret = vmsplice(pipe_fds[1], &iov, 1, SPLICE_F_GIFT);
1131 	dprintf1("vmsplice() ret: %d\n", vmsplice_ret);
1132 	pkey_assert(vmsplice_ret == -1);
1133 
1134 	close(pipe_fds[0]);
1135 	close(pipe_fds[1]);
1136 }
1137 
1138 void test_kernel_gup_write_to_write_disabled_region(int *ptr, u16 pkey)
1139 {
1140 	int ignored = 0xdada;
1141 	int futex_ret;
1142 	int some_int = __LINE__;
1143 
1144 	dprintf1("disabling write to PKEY[%02d], "
1145 		 "doing futex gunk in buffer\n", pkey);
1146 	*ptr = some_int;
1147 	pkey_write_deny(pkey);
1148 	futex_ret = syscall(SYS_futex, ptr, FUTEX_WAIT, some_int-1, NULL,
1149 			&ignored, ignored);
1150 	if (DEBUG_LEVEL > 0)
1151 		perror("futex");
1152 	dprintf1("futex() ret: %d\n", futex_ret);
1153 }
1154 
1155 /* Assumes that all pkeys other than 'pkey' are unallocated */
1156 void test_pkey_syscalls_on_non_allocated_pkey(int *ptr, u16 pkey)
1157 {
1158 	int err;
1159 	int i;
1160 
1161 	/* Note: 0 is the default pkey, so don't mess with it */
1162 	for (i = 1; i < NR_PKEYS; i++) {
1163 		if (pkey == i)
1164 			continue;
1165 
1166 		dprintf1("trying get/set/free to non-allocated pkey: %2d\n", i);
1167 		err = sys_pkey_free(i);
1168 		pkey_assert(err);
1169 
1170 		err = sys_pkey_free(i);
1171 		pkey_assert(err);
1172 
1173 		err = sys_mprotect_pkey(ptr, PAGE_SIZE, PROT_READ, i);
1174 		pkey_assert(err);
1175 	}
1176 }
1177 
1178 /* Assumes that all pkeys other than 'pkey' are unallocated */
1179 void test_pkey_syscalls_bad_args(int *ptr, u16 pkey)
1180 {
1181 	int err;
1182 	int bad_pkey = NR_PKEYS+99;
1183 
1184 	/* pass a known-invalid pkey in: */
1185 	err = sys_mprotect_pkey(ptr, PAGE_SIZE, PROT_READ, bad_pkey);
1186 	pkey_assert(err);
1187 }
1188 
1189 void become_child(void)
1190 {
1191 	pid_t forkret;
1192 
1193 	forkret = fork();
1194 	pkey_assert(forkret >= 0);
1195 	dprintf3("[%d] fork() ret: %d\n", getpid(), forkret);
1196 
1197 	if (!forkret) {
1198 		/* in the child */
1199 		return;
1200 	}
1201 	exit(0);
1202 }
1203 
1204 /* Assumes that all pkeys other than 'pkey' are unallocated */
1205 void test_pkey_alloc_exhaust(int *ptr, u16 pkey)
1206 {
1207 	int err;
1208 	int allocated_pkeys[NR_PKEYS] = {0};
1209 	int nr_allocated_pkeys = 0;
1210 	int i;
1211 
1212 	for (i = 0; i < NR_PKEYS*3; i++) {
1213 		int new_pkey;
1214 		dprintf1("%s() alloc loop: %d\n", __func__, i);
1215 		new_pkey = alloc_pkey();
1216 		dprintf4("%s()::%d, err: %d pkey_reg: 0x%016llx"
1217 				" shadow: 0x%016llx\n",
1218 				__func__, __LINE__, err, __read_pkey_reg(),
1219 				shadow_pkey_reg);
1220 		read_pkey_reg(); /* for shadow checking */
1221 		dprintf2("%s() errno: %d ENOSPC: %d\n", __func__, errno, ENOSPC);
1222 		if ((new_pkey == -1) && (errno == ENOSPC)) {
1223 			dprintf2("%s() failed to allocate pkey after %d tries\n",
1224 				__func__, nr_allocated_pkeys);
1225 		} else {
1226 			/*
1227 			 * Ensure the number of successes never
1228 			 * exceeds the number of keys supported
1229 			 * in the hardware.
1230 			 */
1231 			pkey_assert(nr_allocated_pkeys < NR_PKEYS);
1232 			allocated_pkeys[nr_allocated_pkeys++] = new_pkey;
1233 		}
1234 
1235 		/*
1236 		 * Make sure that allocation state is properly
1237 		 * preserved across fork().
1238 		 */
1239 		if (i == NR_PKEYS*2)
1240 			become_child();
1241 	}
1242 
1243 	dprintf3("%s()::%d\n", __func__, __LINE__);
1244 
1245 	/*
1246 	 * On x86:
1247 	 * There are 16 pkeys supported in hardware.  Three are
1248 	 * allocated by the time we get here:
1249 	 *   1. The default key (0)
1250 	 *   2. One possibly consumed by an execute-only mapping.
1251 	 *   3. One allocated by the test code and passed in via
1252 	 *      'pkey' to this function.
1253 	 * Ensure that we can allocate at least another 13 (16-3).
1254 	 *
1255 	 * On powerpc:
1256 	 * There are either 5, 28, 29 or 32 pkeys supported in
1257 	 * hardware depending on the page size (4K or 64K) and
1258 	 * platform (powernv or powervm). Four are allocated by
1259 	 * the time we get here. These include pkey-0, pkey-1,
1260 	 * exec-only pkey and the one allocated by the test code.
1261 	 * Ensure that we can allocate the remaining.
1262 	 */
1263 	pkey_assert(i >= (NR_PKEYS - get_arch_reserved_keys() - 1));
1264 
1265 	for (i = 0; i < nr_allocated_pkeys; i++) {
1266 		err = sys_pkey_free(allocated_pkeys[i]);
1267 		pkey_assert(!err);
1268 		read_pkey_reg(); /* for shadow checking */
1269 	}
1270 }
1271 
1272 void arch_force_pkey_reg_init(void)
1273 {
1274 #if defined(__i386__) || defined(__x86_64__) /* arch */
1275 	u64 *buf;
1276 
1277 	/*
1278 	 * All keys should be allocated and set to allow reads and
1279 	 * writes, so the register should be all 0.  If not, just
1280 	 * skip the test.
1281 	 */
1282 	if (read_pkey_reg())
1283 		return;
1284 
1285 	/*
1286 	 * Just allocate an absurd about of memory rather than
1287 	 * doing the XSAVE size enumeration dance.
1288 	 */
1289 	buf = mmap(NULL, 1*MB, PROT_READ|PROT_WRITE, MAP_ANONYMOUS|MAP_PRIVATE, -1, 0);
1290 
1291 	/* These __builtins require compiling with -mxsave */
1292 
1293 	/* XSAVE to build a valid buffer: */
1294 	__builtin_ia32_xsave(buf, XSTATE_PKEY);
1295 	/* Clear XSTATE_BV[PKRU]: */
1296 	buf[XSTATE_BV_OFFSET/sizeof(u64)] &= ~XSTATE_PKEY;
1297 	/* XRSTOR will likely get PKRU back to the init state: */
1298 	__builtin_ia32_xrstor(buf, XSTATE_PKEY);
1299 
1300 	munmap(buf, 1*MB);
1301 #endif
1302 }
1303 
1304 
1305 /*
1306  * This is mostly useless on ppc for now.  But it will not
1307  * hurt anything and should give some better coverage as
1308  * a long-running test that continually checks the pkey
1309  * register.
1310  */
1311 void test_pkey_init_state(int *ptr, u16 pkey)
1312 {
1313 	int err;
1314 	int allocated_pkeys[NR_PKEYS] = {0};
1315 	int nr_allocated_pkeys = 0;
1316 	int i;
1317 
1318 	for (i = 0; i < NR_PKEYS; i++) {
1319 		int new_pkey = alloc_pkey();
1320 
1321 		if (new_pkey < 0)
1322 			continue;
1323 		allocated_pkeys[nr_allocated_pkeys++] = new_pkey;
1324 	}
1325 
1326 	dprintf3("%s()::%d\n", __func__, __LINE__);
1327 
1328 	arch_force_pkey_reg_init();
1329 
1330 	/*
1331 	 * Loop for a bit, hoping to get exercise the kernel
1332 	 * context switch code.
1333 	 */
1334 	for (i = 0; i < 1000000; i++)
1335 		read_pkey_reg();
1336 
1337 	for (i = 0; i < nr_allocated_pkeys; i++) {
1338 		err = sys_pkey_free(allocated_pkeys[i]);
1339 		pkey_assert(!err);
1340 		read_pkey_reg(); /* for shadow checking */
1341 	}
1342 }
1343 
1344 /*
1345  * pkey 0 is special.  It is allocated by default, so you do not
1346  * have to call pkey_alloc() to use it first.  Make sure that it
1347  * is usable.
1348  */
1349 void test_mprotect_with_pkey_0(int *ptr, u16 pkey)
1350 {
1351 	long size;
1352 	int prot;
1353 
1354 	assert(pkey_last_malloc_record);
1355 	size = pkey_last_malloc_record->size;
1356 	/*
1357 	 * This is a bit of a hack.  But mprotect() requires
1358 	 * huge-page-aligned sizes when operating on hugetlbfs.
1359 	 * So, make sure that we use something that's a multiple
1360 	 * of a huge page when we can.
1361 	 */
1362 	if (size >= HPAGE_SIZE)
1363 		size = HPAGE_SIZE;
1364 	prot = pkey_last_malloc_record->prot;
1365 
1366 	/* Use pkey 0 */
1367 	mprotect_pkey(ptr, size, prot, 0);
1368 
1369 	/* Make sure that we can set it back to the original pkey. */
1370 	mprotect_pkey(ptr, size, prot, pkey);
1371 }
1372 
1373 void test_ptrace_of_child(int *ptr, u16 pkey)
1374 {
1375 	__attribute__((__unused__)) int peek_result;
1376 	pid_t child_pid;
1377 	void *ignored = 0;
1378 	long ret;
1379 	int status;
1380 	/*
1381 	 * This is the "control" for our little expermient.  Make sure
1382 	 * we can always access it when ptracing.
1383 	 */
1384 	int *plain_ptr_unaligned = malloc(HPAGE_SIZE);
1385 	int *plain_ptr = ALIGN_PTR_UP(plain_ptr_unaligned, PAGE_SIZE);
1386 
1387 	/*
1388 	 * Fork a child which is an exact copy of this process, of course.
1389 	 * That means we can do all of our tests via ptrace() and then plain
1390 	 * memory access and ensure they work differently.
1391 	 */
1392 	child_pid = fork_lazy_child();
1393 	dprintf1("[%d] child pid: %d\n", getpid(), child_pid);
1394 
1395 	ret = ptrace(PTRACE_ATTACH, child_pid, ignored, ignored);
1396 	if (ret)
1397 		perror("attach");
1398 	dprintf1("[%d] attach ret: %ld %d\n", getpid(), ret, __LINE__);
1399 	pkey_assert(ret != -1);
1400 	ret = waitpid(child_pid, &status, WUNTRACED);
1401 	if ((ret != child_pid) || !(WIFSTOPPED(status))) {
1402 		fprintf(stderr, "weird waitpid result %ld stat %x\n",
1403 				ret, status);
1404 		pkey_assert(0);
1405 	}
1406 	dprintf2("waitpid ret: %ld\n", ret);
1407 	dprintf2("waitpid status: %d\n", status);
1408 
1409 	pkey_access_deny(pkey);
1410 	pkey_write_deny(pkey);
1411 
1412 	/* Write access, untested for now:
1413 	ret = ptrace(PTRACE_POKEDATA, child_pid, peek_at, data);
1414 	pkey_assert(ret != -1);
1415 	dprintf1("poke at %p: %ld\n", peek_at, ret);
1416 	*/
1417 
1418 	/*
1419 	 * Try to access the pkey-protected "ptr" via ptrace:
1420 	 */
1421 	ret = ptrace(PTRACE_PEEKDATA, child_pid, ptr, ignored);
1422 	/* expect it to work, without an error: */
1423 	pkey_assert(ret != -1);
1424 	/* Now access from the current task, and expect an exception: */
1425 	peek_result = read_ptr(ptr);
1426 	expected_pkey_fault(pkey);
1427 
1428 	/*
1429 	 * Try to access the NON-pkey-protected "plain_ptr" via ptrace:
1430 	 */
1431 	ret = ptrace(PTRACE_PEEKDATA, child_pid, plain_ptr, ignored);
1432 	/* expect it to work, without an error: */
1433 	pkey_assert(ret != -1);
1434 	/* Now access from the current task, and expect NO exception: */
1435 	peek_result = read_ptr(plain_ptr);
1436 	do_not_expect_pkey_fault("read plain pointer after ptrace");
1437 
1438 	ret = ptrace(PTRACE_DETACH, child_pid, ignored, 0);
1439 	pkey_assert(ret != -1);
1440 
1441 	ret = kill(child_pid, SIGKILL);
1442 	pkey_assert(ret != -1);
1443 
1444 	wait(&status);
1445 
1446 	free(plain_ptr_unaligned);
1447 }
1448 
1449 void *get_pointer_to_instructions(void)
1450 {
1451 	void *p1;
1452 
1453 	p1 = ALIGN_PTR_UP(&lots_o_noops_around_write, PAGE_SIZE);
1454 	dprintf3("&lots_o_noops: %p\n", &lots_o_noops_around_write);
1455 	/* lots_o_noops_around_write should be page-aligned already */
1456 	assert(p1 == &lots_o_noops_around_write);
1457 
1458 	/* Point 'p1' at the *second* page of the function: */
1459 	p1 += PAGE_SIZE;
1460 
1461 	/*
1462 	 * Try to ensure we fault this in on next touch to ensure
1463 	 * we get an instruction fault as opposed to a data one
1464 	 */
1465 	madvise(p1, PAGE_SIZE, MADV_DONTNEED);
1466 
1467 	return p1;
1468 }
1469 
1470 void test_executing_on_unreadable_memory(int *ptr, u16 pkey)
1471 {
1472 	void *p1;
1473 	int scratch;
1474 	int ptr_contents;
1475 	int ret;
1476 
1477 	p1 = get_pointer_to_instructions();
1478 	lots_o_noops_around_write(&scratch);
1479 	ptr_contents = read_ptr(p1);
1480 	dprintf2("ptr (%p) contents@%d: %x\n", p1, __LINE__, ptr_contents);
1481 
1482 	ret = mprotect_pkey(p1, PAGE_SIZE, PROT_EXEC, (u64)pkey);
1483 	pkey_assert(!ret);
1484 	pkey_access_deny(pkey);
1485 
1486 	dprintf2("pkey_reg: %016llx\n", read_pkey_reg());
1487 
1488 	/*
1489 	 * Make sure this is an *instruction* fault
1490 	 */
1491 	madvise(p1, PAGE_SIZE, MADV_DONTNEED);
1492 	lots_o_noops_around_write(&scratch);
1493 	do_not_expect_pkey_fault("executing on PROT_EXEC memory");
1494 	expect_fault_on_read_execonly_key(p1, pkey);
1495 }
1496 
1497 void test_implicit_mprotect_exec_only_memory(int *ptr, u16 pkey)
1498 {
1499 	void *p1;
1500 	int scratch;
1501 	int ptr_contents;
1502 	int ret;
1503 
1504 	dprintf1("%s() start\n", __func__);
1505 
1506 	p1 = get_pointer_to_instructions();
1507 	lots_o_noops_around_write(&scratch);
1508 	ptr_contents = read_ptr(p1);
1509 	dprintf2("ptr (%p) contents@%d: %x\n", p1, __LINE__, ptr_contents);
1510 
1511 	/* Use a *normal* mprotect(), not mprotect_pkey(): */
1512 	ret = mprotect(p1, PAGE_SIZE, PROT_EXEC);
1513 	pkey_assert(!ret);
1514 
1515 	/*
1516 	 * Reset the shadow, assuming that the above mprotect()
1517 	 * correctly changed PKRU, but to an unknown value since
1518 	 * the actual allocated pkey is unknown.
1519 	 */
1520 	shadow_pkey_reg = __read_pkey_reg();
1521 
1522 	dprintf2("pkey_reg: %016llx\n", read_pkey_reg());
1523 
1524 	/* Make sure this is an *instruction* fault */
1525 	madvise(p1, PAGE_SIZE, MADV_DONTNEED);
1526 	lots_o_noops_around_write(&scratch);
1527 	do_not_expect_pkey_fault("executing on PROT_EXEC memory");
1528 	expect_fault_on_read_execonly_key(p1, UNKNOWN_PKEY);
1529 
1530 	/*
1531 	 * Put the memory back to non-PROT_EXEC.  Should clear the
1532 	 * exec-only pkey off the VMA and allow it to be readable
1533 	 * again.  Go to PROT_NONE first to check for a kernel bug
1534 	 * that did not clear the pkey when doing PROT_NONE.
1535 	 */
1536 	ret = mprotect(p1, PAGE_SIZE, PROT_NONE);
1537 	pkey_assert(!ret);
1538 
1539 	ret = mprotect(p1, PAGE_SIZE, PROT_READ|PROT_EXEC);
1540 	pkey_assert(!ret);
1541 	ptr_contents = read_ptr(p1);
1542 	do_not_expect_pkey_fault("plain read on recently PROT_EXEC area");
1543 }
1544 
1545 #if defined(__i386__) || defined(__x86_64__)
1546 void test_ptrace_modifies_pkru(int *ptr, u16 pkey)
1547 {
1548 	u32 new_pkru;
1549 	pid_t child;
1550 	int status, ret;
1551 	int pkey_offset = pkey_reg_xstate_offset();
1552 	size_t xsave_size = cpu_max_xsave_size();
1553 	void *xsave;
1554 	u32 *pkey_register;
1555 	u64 *xstate_bv;
1556 	struct iovec iov;
1557 
1558 	new_pkru = ~read_pkey_reg();
1559 	/* Don't make PROT_EXEC mappings inaccessible */
1560 	new_pkru &= ~3;
1561 
1562 	child = fork();
1563 	pkey_assert(child >= 0);
1564 	dprintf3("[%d] fork() ret: %d\n", getpid(), child);
1565 	if (!child) {
1566 		ptrace(PTRACE_TRACEME, 0, 0, 0);
1567 		/* Stop and allow the tracer to modify PKRU directly */
1568 		raise(SIGSTOP);
1569 
1570 		/*
1571 		 * need __read_pkey_reg() version so we do not do shadow_pkey_reg
1572 		 * checking
1573 		 */
1574 		if (__read_pkey_reg() != new_pkru)
1575 			exit(1);
1576 
1577 		/* Stop and allow the tracer to clear XSTATE_BV for PKRU */
1578 		raise(SIGSTOP);
1579 
1580 		if (__read_pkey_reg() != 0)
1581 			exit(1);
1582 
1583 		/* Stop and allow the tracer to examine PKRU */
1584 		raise(SIGSTOP);
1585 
1586 		exit(0);
1587 	}
1588 
1589 	pkey_assert(child == waitpid(child, &status, 0));
1590 	dprintf3("[%d] waitpid(%d) status: %x\n", getpid(), child, status);
1591 	pkey_assert(WIFSTOPPED(status) && WSTOPSIG(status) == SIGSTOP);
1592 
1593 	xsave = (void *)malloc(xsave_size);
1594 	pkey_assert(xsave > 0);
1595 
1596 	/* Modify the PKRU register directly */
1597 	iov.iov_base = xsave;
1598 	iov.iov_len = xsave_size;
1599 	ret = ptrace(PTRACE_GETREGSET, child, (void *)NT_X86_XSTATE, &iov);
1600 	pkey_assert(ret == 0);
1601 
1602 	pkey_register = (u32 *)(xsave + pkey_offset);
1603 	pkey_assert(*pkey_register == read_pkey_reg());
1604 
1605 	*pkey_register = new_pkru;
1606 
1607 	ret = ptrace(PTRACE_SETREGSET, child, (void *)NT_X86_XSTATE, &iov);
1608 	pkey_assert(ret == 0);
1609 
1610 	/* Test that the modification is visible in ptrace before any execution */
1611 	memset(xsave, 0xCC, xsave_size);
1612 	ret = ptrace(PTRACE_GETREGSET, child, (void *)NT_X86_XSTATE, &iov);
1613 	pkey_assert(ret == 0);
1614 	pkey_assert(*pkey_register == new_pkru);
1615 
1616 	/* Execute the tracee */
1617 	ret = ptrace(PTRACE_CONT, child, 0, 0);
1618 	pkey_assert(ret == 0);
1619 
1620 	/* Test that the tracee saw the PKRU value change */
1621 	pkey_assert(child == waitpid(child, &status, 0));
1622 	dprintf3("[%d] waitpid(%d) status: %x\n", getpid(), child, status);
1623 	pkey_assert(WIFSTOPPED(status) && WSTOPSIG(status) == SIGSTOP);
1624 
1625 	/* Test that the modification is visible in ptrace after execution */
1626 	memset(xsave, 0xCC, xsave_size);
1627 	ret = ptrace(PTRACE_GETREGSET, child, (void *)NT_X86_XSTATE, &iov);
1628 	pkey_assert(ret == 0);
1629 	pkey_assert(*pkey_register == new_pkru);
1630 
1631 	/* Clear the PKRU bit from XSTATE_BV */
1632 	xstate_bv = (u64 *)(xsave + 512);
1633 	*xstate_bv &= ~(1 << 9);
1634 
1635 	ret = ptrace(PTRACE_SETREGSET, child, (void *)NT_X86_XSTATE, &iov);
1636 	pkey_assert(ret == 0);
1637 
1638 	/* Test that the modification is visible in ptrace before any execution */
1639 	memset(xsave, 0xCC, xsave_size);
1640 	ret = ptrace(PTRACE_GETREGSET, child, (void *)NT_X86_XSTATE, &iov);
1641 	pkey_assert(ret == 0);
1642 	pkey_assert(*pkey_register == 0);
1643 
1644 	ret = ptrace(PTRACE_CONT, child, 0, 0);
1645 	pkey_assert(ret == 0);
1646 
1647 	/* Test that the tracee saw the PKRU value go to 0 */
1648 	pkey_assert(child == waitpid(child, &status, 0));
1649 	dprintf3("[%d] waitpid(%d) status: %x\n", getpid(), child, status);
1650 	pkey_assert(WIFSTOPPED(status) && WSTOPSIG(status) == SIGSTOP);
1651 
1652 	/* Test that the modification is visible in ptrace after execution */
1653 	memset(xsave, 0xCC, xsave_size);
1654 	ret = ptrace(PTRACE_GETREGSET, child, (void *)NT_X86_XSTATE, &iov);
1655 	pkey_assert(ret == 0);
1656 	pkey_assert(*pkey_register == 0);
1657 
1658 	ret = ptrace(PTRACE_CONT, child, 0, 0);
1659 	pkey_assert(ret == 0);
1660 	pkey_assert(child == waitpid(child, &status, 0));
1661 	dprintf3("[%d] waitpid(%d) status: %x\n", getpid(), child, status);
1662 	pkey_assert(WIFEXITED(status));
1663 	pkey_assert(WEXITSTATUS(status) == 0);
1664 	free(xsave);
1665 }
1666 #endif
1667 
1668 void test_mprotect_pkey_on_unsupported_cpu(int *ptr, u16 pkey)
1669 {
1670 	int size = PAGE_SIZE;
1671 	int sret;
1672 
1673 	if (cpu_has_pkeys()) {
1674 		dprintf1("SKIP: %s: no CPU support\n", __func__);
1675 		return;
1676 	}
1677 
1678 	sret = syscall(__NR_pkey_mprotect, ptr, size, PROT_READ, pkey);
1679 	pkey_assert(sret < 0);
1680 }
1681 
1682 void (*pkey_tests[])(int *ptr, u16 pkey) = {
1683 	test_read_of_write_disabled_region,
1684 	test_read_of_access_disabled_region,
1685 	test_read_of_access_disabled_region_with_page_already_mapped,
1686 	test_write_of_write_disabled_region,
1687 	test_write_of_write_disabled_region_with_page_already_mapped,
1688 	test_write_of_access_disabled_region,
1689 	test_write_of_access_disabled_region_with_page_already_mapped,
1690 	test_kernel_write_of_access_disabled_region,
1691 	test_kernel_write_of_write_disabled_region,
1692 	test_kernel_gup_of_access_disabled_region,
1693 	test_kernel_gup_write_to_write_disabled_region,
1694 	test_executing_on_unreadable_memory,
1695 	test_implicit_mprotect_exec_only_memory,
1696 	test_mprotect_with_pkey_0,
1697 	test_ptrace_of_child,
1698 	test_pkey_init_state,
1699 	test_pkey_syscalls_on_non_allocated_pkey,
1700 	test_pkey_syscalls_bad_args,
1701 	test_pkey_alloc_exhaust,
1702 	test_pkey_alloc_free_attach_pkey0,
1703 #if defined(__i386__) || defined(__x86_64__)
1704 	test_ptrace_modifies_pkru,
1705 #endif
1706 };
1707 
1708 void run_tests_once(void)
1709 {
1710 	int *ptr;
1711 	int prot = PROT_READ|PROT_WRITE;
1712 
1713 	for (test_nr = 0; test_nr < ARRAY_SIZE(pkey_tests); test_nr++) {
1714 		int pkey;
1715 		int orig_pkey_faults = pkey_faults;
1716 
1717 		dprintf1("======================\n");
1718 		dprintf1("test %d preparing...\n", test_nr);
1719 
1720 		tracing_on();
1721 		pkey = alloc_random_pkey();
1722 		dprintf1("test %d starting with pkey: %d\n", test_nr, pkey);
1723 		ptr = malloc_pkey(PAGE_SIZE, prot, pkey);
1724 		dprintf1("test %d starting...\n", test_nr);
1725 		pkey_tests[test_nr](ptr, pkey);
1726 		dprintf1("freeing test memory: %p\n", ptr);
1727 		free_pkey_malloc(ptr);
1728 		sys_pkey_free(pkey);
1729 
1730 		dprintf1("pkey_faults: %d\n", pkey_faults);
1731 		dprintf1("orig_pkey_faults: %d\n", orig_pkey_faults);
1732 
1733 		tracing_off();
1734 		close_test_fds();
1735 
1736 		printf("test %2d PASSED (iteration %d)\n", test_nr, iteration_nr);
1737 		dprintf1("======================\n\n");
1738 	}
1739 	iteration_nr++;
1740 }
1741 
1742 void pkey_setup_shadow(void)
1743 {
1744 	shadow_pkey_reg = __read_pkey_reg();
1745 }
1746 
1747 int main(void)
1748 {
1749 	int nr_iterations = 22;
1750 	int pkeys_supported = is_pkeys_supported();
1751 
1752 	srand((unsigned int)time(NULL));
1753 
1754 	setup_handlers();
1755 
1756 	printf("has pkeys: %d\n", pkeys_supported);
1757 
1758 	if (!pkeys_supported) {
1759 		int size = PAGE_SIZE;
1760 		int *ptr;
1761 
1762 		printf("running PKEY tests for unsupported CPU/OS\n");
1763 
1764 		ptr  = mmap(NULL, size, PROT_NONE, MAP_ANONYMOUS|MAP_PRIVATE, -1, 0);
1765 		assert(ptr != (void *)-1);
1766 		test_mprotect_pkey_on_unsupported_cpu(ptr, 1);
1767 		exit(0);
1768 	}
1769 
1770 	pkey_setup_shadow();
1771 	printf("startup pkey_reg: %016llx\n", read_pkey_reg());
1772 	setup_hugetlbfs();
1773 
1774 	while (nr_iterations-- > 0)
1775 		run_tests_once();
1776 
1777 	printf("done (all tests OK)\n");
1778 	return 0;
1779 }
1780