1#define _GNU_SOURCE
2#include <sys/mman.h>
3#include <stdint.h>
4#include <stdio.h>
5#include <stdlib.h>
6#include <unistd.h>
7#include <string.h>
8#include <sys/time.h>
9#include <sys/resource.h>
10#include <syscall.h>
11#include <errno.h>
12#include <stdbool.h>
13
14#ifndef MLOCK_ONFAULT
15#define MLOCK_ONFAULT 1
16#endif
17
18#ifndef MCL_ONFAULT
19#define MCL_ONFAULT (MCL_FUTURE << 1)
20#endif
21
22static int mlock2_(void *start, size_t len, int flags)
23{
24#ifdef __NR_mlock2
25	return syscall(__NR_mlock2, start, len, flags);
26#else
27	errno = ENOSYS;
28	return -1;
29#endif
30}
31
32struct vm_boundaries {
33	unsigned long start;
34	unsigned long end;
35};
36
37static int get_vm_area(unsigned long addr, struct vm_boundaries *area)
38{
39	FILE *file;
40	int ret = 1;
41	char line[1024] = {0};
42	char *end_addr;
43	char *stop;
44	unsigned long start;
45	unsigned long end;
46
47	if (!area)
48		return ret;
49
50	file = fopen("/proc/self/maps", "r");
51	if (!file) {
52		perror("fopen");
53		return ret;
54	}
55
56	memset(area, 0, sizeof(struct vm_boundaries));
57
58	while(fgets(line, 1024, file)) {
59		end_addr = strchr(line, '-');
60		if (!end_addr) {
61			printf("cannot parse /proc/self/maps\n");
62			goto out;
63		}
64		*end_addr = '\0';
65		end_addr++;
66		stop = strchr(end_addr, ' ');
67		if (!stop) {
68			printf("cannot parse /proc/self/maps\n");
69			goto out;
70		}
71		stop = '\0';
72
73		sscanf(line, "%lx", &start);
74		sscanf(end_addr, "%lx", &end);
75
76		if (start <= addr && end > addr) {
77			area->start = start;
78			area->end = end;
79			ret = 0;
80			goto out;
81		}
82	}
83out:
84	fclose(file);
85	return ret;
86}
87
88static uint64_t get_pageflags(unsigned long addr)
89{
90	FILE *file;
91	uint64_t pfn;
92	unsigned long offset;
93
94	file = fopen("/proc/self/pagemap", "r");
95	if (!file) {
96		perror("fopen pagemap");
97		_exit(1);
98	}
99
100	offset = addr / getpagesize() * sizeof(pfn);
101
102	if (fseek(file, offset, SEEK_SET)) {
103		perror("fseek pagemap");
104		_exit(1);
105	}
106
107	if (fread(&pfn, sizeof(pfn), 1, file) != 1) {
108		perror("fread pagemap");
109		_exit(1);
110	}
111
112	fclose(file);
113	return pfn;
114}
115
116static uint64_t get_kpageflags(unsigned long pfn)
117{
118	uint64_t flags;
119	FILE *file;
120
121	file = fopen("/proc/kpageflags", "r");
122	if (!file) {
123		perror("fopen kpageflags");
124		_exit(1);
125	}
126
127	if (fseek(file, pfn * sizeof(flags), SEEK_SET)) {
128		perror("fseek kpageflags");
129		_exit(1);
130	}
131
132	if (fread(&flags, sizeof(flags), 1, file) != 1) {
133		perror("fread kpageflags");
134		_exit(1);
135	}
136
137	fclose(file);
138	return flags;
139}
140
141static FILE *seek_to_smaps_entry(unsigned long addr)
142{
143	FILE *file;
144	char *line = NULL;
145	size_t size = 0;
146	unsigned long start, end;
147	char perms[5];
148	unsigned long offset;
149	char dev[32];
150	unsigned long inode;
151	char path[BUFSIZ];
152
153	file = fopen("/proc/self/smaps", "r");
154	if (!file) {
155		perror("fopen smaps");
156		_exit(1);
157	}
158
159	while (getline(&line, &size, file) > 0) {
160		if (sscanf(line, "%lx-%lx %s %lx %s %lu %s\n",
161			   &start, &end, perms, &offset, dev, &inode, path) < 6)
162			goto next;
163
164		if (start <= addr && addr < end)
165			goto out;
166
167next:
168		free(line);
169		line = NULL;
170		size = 0;
171	}
172
173	fclose(file);
174	file = NULL;
175
176out:
177	free(line);
178	return file;
179}
180
181#define VMFLAGS "VmFlags:"
182
183static bool is_vmflag_set(unsigned long addr, const char *vmflag)
184{
185	char *line = NULL;
186	char *flags;
187	size_t size = 0;
188	bool ret = false;
189	FILE *smaps;
190
191	smaps = seek_to_smaps_entry(addr);
192	if (!smaps) {
193		printf("Unable to parse /proc/self/smaps\n");
194		goto out;
195	}
196
197	while (getline(&line, &size, smaps) > 0) {
198		if (!strstr(line, VMFLAGS)) {
199			free(line);
200			line = NULL;
201			size = 0;
202			continue;
203		}
204
205		flags = line + strlen(VMFLAGS);
206		ret = (strstr(flags, vmflag) != NULL);
207		goto out;
208	}
209
210out:
211	free(line);
212	fclose(smaps);
213	return ret;
214}
215
216#define SIZE "Size:"
217#define RSS  "Rss:"
218#define LOCKED "lo"
219
220static bool is_vma_lock_on_fault(unsigned long addr)
221{
222	bool ret = false;
223	bool locked;
224	FILE *smaps = NULL;
225	unsigned long vma_size, vma_rss;
226	char *line = NULL;
227	char *value;
228	size_t size = 0;
229
230	locked = is_vmflag_set(addr, LOCKED);
231	if (!locked)
232		goto out;
233
234	smaps = seek_to_smaps_entry(addr);
235	if (!smaps) {
236		printf("Unable to parse /proc/self/smaps\n");
237		goto out;
238	}
239
240	while (getline(&line, &size, smaps) > 0) {
241		if (!strstr(line, SIZE)) {
242			free(line);
243			line = NULL;
244			size = 0;
245			continue;
246		}
247
248		value = line + strlen(SIZE);
249		if (sscanf(value, "%lu kB", &vma_size) < 1) {
250			printf("Unable to parse smaps entry for Size\n");
251			goto out;
252		}
253		break;
254	}
255
256	while (getline(&line, &size, smaps) > 0) {
257		if (!strstr(line, RSS)) {
258			free(line);
259			line = NULL;
260			size = 0;
261			continue;
262		}
263
264		value = line + strlen(RSS);
265		if (sscanf(value, "%lu kB", &vma_rss) < 1) {
266			printf("Unable to parse smaps entry for Rss\n");
267			goto out;
268		}
269		break;
270	}
271
272	ret = locked && (vma_rss < vma_size);
273out:
274	free(line);
275	if (smaps)
276		fclose(smaps);
277	return ret;
278}
279
280#define PRESENT_BIT     0x8000000000000000ULL
281#define PFN_MASK        0x007FFFFFFFFFFFFFULL
282#define UNEVICTABLE_BIT (1UL << 18)
283
284static int lock_check(char *map)
285{
286	unsigned long page_size = getpagesize();
287	uint64_t page1_flags, page2_flags;
288
289	page1_flags = get_pageflags((unsigned long)map);
290	page2_flags = get_pageflags((unsigned long)map + page_size);
291
292	/* Both pages should be present */
293	if (((page1_flags & PRESENT_BIT) == 0) ||
294	    ((page2_flags & PRESENT_BIT) == 0)) {
295		printf("Failed to make both pages present\n");
296		return 1;
297	}
298
299	page1_flags = get_kpageflags(page1_flags & PFN_MASK);
300	page2_flags = get_kpageflags(page2_flags & PFN_MASK);
301
302	/* Both pages should be unevictable */
303	if (((page1_flags & UNEVICTABLE_BIT) == 0) ||
304	    ((page2_flags & UNEVICTABLE_BIT) == 0)) {
305		printf("Failed to make both pages unevictable\n");
306		return 1;
307	}
308
309	if (!is_vmflag_set((unsigned long)map, LOCKED)) {
310		printf("VMA flag %s is missing on page 1\n", LOCKED);
311		return 1;
312	}
313
314	if (!is_vmflag_set((unsigned long)map + page_size, LOCKED)) {
315		printf("VMA flag %s is missing on page 2\n", LOCKED);
316		return 1;
317	}
318
319	return 0;
320}
321
322static int unlock_lock_check(char *map)
323{
324	unsigned long page_size = getpagesize();
325	uint64_t page1_flags, page2_flags;
326
327	page1_flags = get_pageflags((unsigned long)map);
328	page2_flags = get_pageflags((unsigned long)map + page_size);
329	page1_flags = get_kpageflags(page1_flags & PFN_MASK);
330	page2_flags = get_kpageflags(page2_flags & PFN_MASK);
331
332	if ((page1_flags & UNEVICTABLE_BIT) || (page2_flags & UNEVICTABLE_BIT)) {
333		printf("A page is still marked unevictable after unlock\n");
334		return 1;
335	}
336
337	if (is_vmflag_set((unsigned long)map, LOCKED)) {
338		printf("VMA flag %s is present on page 1 after unlock\n", LOCKED);
339		return 1;
340	}
341
342	if (is_vmflag_set((unsigned long)map + page_size, LOCKED)) {
343		printf("VMA flag %s is present on page 2 after unlock\n", LOCKED);
344		return 1;
345	}
346
347	return 0;
348}
349
350static int test_mlock_lock()
351{
352	char *map;
353	int ret = 1;
354	unsigned long page_size = getpagesize();
355
356	map = mmap(NULL, 2 * page_size, PROT_READ | PROT_WRITE,
357		   MAP_ANONYMOUS | MAP_PRIVATE, 0, 0);
358	if (map == MAP_FAILED) {
359		perror("test_mlock_locked mmap");
360		goto out;
361	}
362
363	if (mlock2_(map, 2 * page_size, 0)) {
364		if (errno == ENOSYS) {
365			printf("Cannot call new mlock family, skipping test\n");
366			_exit(0);
367		}
368		perror("mlock2(0)");
369		goto unmap;
370	}
371
372	if (lock_check(map))
373		goto unmap;
374
375	/* Now unlock and recheck attributes */
376	if (munlock(map, 2 * page_size)) {
377		perror("munlock()");
378		goto unmap;
379	}
380
381	ret = unlock_lock_check(map);
382
383unmap:
384	munmap(map, 2 * page_size);
385out:
386	return ret;
387}
388
389static int onfault_check(char *map)
390{
391	unsigned long page_size = getpagesize();
392	uint64_t page1_flags, page2_flags;
393
394	page1_flags = get_pageflags((unsigned long)map);
395	page2_flags = get_pageflags((unsigned long)map + page_size);
396
397	/* Neither page should be present */
398	if ((page1_flags & PRESENT_BIT) || (page2_flags & PRESENT_BIT)) {
399		printf("Pages were made present by MLOCK_ONFAULT\n");
400		return 1;
401	}
402
403	*map = 'a';
404	page1_flags = get_pageflags((unsigned long)map);
405	page2_flags = get_pageflags((unsigned long)map + page_size);
406
407	/* Only page 1 should be present */
408	if ((page1_flags & PRESENT_BIT) == 0) {
409		printf("Page 1 is not present after fault\n");
410		return 1;
411	} else if (page2_flags & PRESENT_BIT) {
412		printf("Page 2 was made present\n");
413		return 1;
414	}
415
416	page1_flags = get_kpageflags(page1_flags & PFN_MASK);
417
418	/* Page 1 should be unevictable */
419	if ((page1_flags & UNEVICTABLE_BIT) == 0) {
420		printf("Failed to make faulted page unevictable\n");
421		return 1;
422	}
423
424	if (!is_vma_lock_on_fault((unsigned long)map)) {
425		printf("VMA is not marked for lock on fault\n");
426		return 1;
427	}
428
429	if (!is_vma_lock_on_fault((unsigned long)map + page_size)) {
430		printf("VMA is not marked for lock on fault\n");
431		return 1;
432	}
433
434	return 0;
435}
436
437static int unlock_onfault_check(char *map)
438{
439	unsigned long page_size = getpagesize();
440	uint64_t page1_flags;
441
442	page1_flags = get_pageflags((unsigned long)map);
443	page1_flags = get_kpageflags(page1_flags & PFN_MASK);
444
445	if (page1_flags & UNEVICTABLE_BIT) {
446		printf("Page 1 is still marked unevictable after unlock\n");
447		return 1;
448	}
449
450	if (is_vma_lock_on_fault((unsigned long)map) ||
451	    is_vma_lock_on_fault((unsigned long)map + page_size)) {
452		printf("VMA is still lock on fault after unlock\n");
453		return 1;
454	}
455
456	return 0;
457}
458
459static int test_mlock_onfault()
460{
461	char *map;
462	int ret = 1;
463	unsigned long page_size = getpagesize();
464
465	map = mmap(NULL, 2 * page_size, PROT_READ | PROT_WRITE,
466		   MAP_ANONYMOUS | MAP_PRIVATE, 0, 0);
467	if (map == MAP_FAILED) {
468		perror("test_mlock_locked mmap");
469		goto out;
470	}
471
472	if (mlock2_(map, 2 * page_size, MLOCK_ONFAULT)) {
473		if (errno == ENOSYS) {
474			printf("Cannot call new mlock family, skipping test\n");
475			_exit(0);
476		}
477		perror("mlock2(MLOCK_ONFAULT)");
478		goto unmap;
479	}
480
481	if (onfault_check(map))
482		goto unmap;
483
484	/* Now unlock and recheck attributes */
485	if (munlock(map, 2 * page_size)) {
486		if (errno == ENOSYS) {
487			printf("Cannot call new mlock family, skipping test\n");
488			_exit(0);
489		}
490		perror("munlock()");
491		goto unmap;
492	}
493
494	ret = unlock_onfault_check(map);
495unmap:
496	munmap(map, 2 * page_size);
497out:
498	return ret;
499}
500
501static int test_lock_onfault_of_present()
502{
503	char *map;
504	int ret = 1;
505	unsigned long page_size = getpagesize();
506	uint64_t page1_flags, page2_flags;
507
508	map = mmap(NULL, 2 * page_size, PROT_READ | PROT_WRITE,
509		   MAP_ANONYMOUS | MAP_PRIVATE, 0, 0);
510	if (map == MAP_FAILED) {
511		perror("test_mlock_locked mmap");
512		goto out;
513	}
514
515	*map = 'a';
516
517	if (mlock2_(map, 2 * page_size, MLOCK_ONFAULT)) {
518		if (errno == ENOSYS) {
519			printf("Cannot call new mlock family, skipping test\n");
520			_exit(0);
521		}
522		perror("mlock2(MLOCK_ONFAULT)");
523		goto unmap;
524	}
525
526	page1_flags = get_pageflags((unsigned long)map);
527	page2_flags = get_pageflags((unsigned long)map + page_size);
528	page1_flags = get_kpageflags(page1_flags & PFN_MASK);
529	page2_flags = get_kpageflags(page2_flags & PFN_MASK);
530
531	/* Page 1 should be unevictable */
532	if ((page1_flags & UNEVICTABLE_BIT) == 0) {
533		printf("Failed to make present page unevictable\n");
534		goto unmap;
535	}
536
537	if (!is_vma_lock_on_fault((unsigned long)map) ||
538	    !is_vma_lock_on_fault((unsigned long)map + page_size)) {
539		printf("VMA with present pages is not marked lock on fault\n");
540		goto unmap;
541	}
542	ret = 0;
543unmap:
544	munmap(map, 2 * page_size);
545out:
546	return ret;
547}
548
549static int test_munlockall()
550{
551	char *map;
552	int ret = 1;
553	unsigned long page_size = getpagesize();
554
555	map = mmap(NULL, 2 * page_size, PROT_READ | PROT_WRITE,
556		   MAP_ANONYMOUS | MAP_PRIVATE, 0, 0);
557
558	if (map == MAP_FAILED) {
559		perror("test_munlockall mmap");
560		goto out;
561	}
562
563	if (mlockall(MCL_CURRENT)) {
564		perror("mlockall(MCL_CURRENT)");
565		goto out;
566	}
567
568	if (lock_check(map))
569		goto unmap;
570
571	if (munlockall()) {
572		perror("munlockall()");
573		goto unmap;
574	}
575
576	if (unlock_lock_check(map))
577		goto unmap;
578
579	munmap(map, 2 * page_size);
580
581	map = mmap(NULL, 2 * page_size, PROT_READ | PROT_WRITE,
582		   MAP_ANONYMOUS | MAP_PRIVATE, 0, 0);
583
584	if (map == MAP_FAILED) {
585		perror("test_munlockall second mmap");
586		goto out;
587	}
588
589	if (mlockall(MCL_CURRENT | MCL_ONFAULT)) {
590		perror("mlockall(MCL_CURRENT | MCL_ONFAULT)");
591		goto unmap;
592	}
593
594	if (onfault_check(map))
595		goto unmap;
596
597	if (munlockall()) {
598		perror("munlockall()");
599		goto unmap;
600	}
601
602	if (unlock_onfault_check(map))
603		goto unmap;
604
605	if (mlockall(MCL_CURRENT | MCL_FUTURE)) {
606		perror("mlockall(MCL_CURRENT | MCL_FUTURE)");
607		goto out;
608	}
609
610	if (lock_check(map))
611		goto unmap;
612
613	if (munlockall()) {
614		perror("munlockall()");
615		goto unmap;
616	}
617
618	ret = unlock_lock_check(map);
619
620unmap:
621	munmap(map, 2 * page_size);
622out:
623	munlockall();
624	return ret;
625}
626
627static int test_vma_management(bool call_mlock)
628{
629	int ret = 1;
630	void *map;
631	unsigned long page_size = getpagesize();
632	struct vm_boundaries page1;
633	struct vm_boundaries page2;
634	struct vm_boundaries page3;
635
636	map = mmap(NULL, 3 * page_size, PROT_READ | PROT_WRITE,
637		   MAP_ANONYMOUS | MAP_PRIVATE, 0, 0);
638	if (map == MAP_FAILED) {
639		perror("mmap()");
640		return ret;
641	}
642
643	if (call_mlock && mlock2_(map, 3 * page_size, MLOCK_ONFAULT)) {
644		if (errno == ENOSYS) {
645			printf("Cannot call new mlock family, skipping test\n");
646			_exit(0);
647		}
648		perror("mlock(ONFAULT)\n");
649		goto out;
650	}
651
652	if (get_vm_area((unsigned long)map, &page1) ||
653	    get_vm_area((unsigned long)map + page_size, &page2) ||
654	    get_vm_area((unsigned long)map + page_size * 2, &page3)) {
655		printf("couldn't find mapping in /proc/self/maps\n");
656		goto out;
657	}
658
659	/*
660	 * Before we unlock a portion, we need to that all three pages are in
661	 * the same VMA.  If they are not we abort this test (Note that this is
662	 * not a failure)
663	 */
664	if (page1.start != page2.start || page2.start != page3.start) {
665		printf("VMAs are not merged to start, aborting test\n");
666		ret = 0;
667		goto out;
668	}
669
670	if (munlock(map + page_size, page_size)) {
671		perror("munlock()");
672		goto out;
673	}
674
675	if (get_vm_area((unsigned long)map, &page1) ||
676	    get_vm_area((unsigned long)map + page_size, &page2) ||
677	    get_vm_area((unsigned long)map + page_size * 2, &page3)) {
678		printf("couldn't find mapping in /proc/self/maps\n");
679		goto out;
680	}
681
682	/* All three VMAs should be different */
683	if (page1.start == page2.start || page2.start == page3.start) {
684		printf("failed to split VMA for munlock\n");
685		goto out;
686	}
687
688	/* Now unlock the first and third page and check the VMAs again */
689	if (munlock(map, page_size * 3)) {
690		perror("munlock()");
691		goto out;
692	}
693
694	if (get_vm_area((unsigned long)map, &page1) ||
695	    get_vm_area((unsigned long)map + page_size, &page2) ||
696	    get_vm_area((unsigned long)map + page_size * 2, &page3)) {
697		printf("couldn't find mapping in /proc/self/maps\n");
698		goto out;
699	}
700
701	/* Now all three VMAs should be the same */
702	if (page1.start != page2.start || page2.start != page3.start) {
703		printf("failed to merge VMAs after munlock\n");
704		goto out;
705	}
706
707	ret = 0;
708out:
709	munmap(map, 3 * page_size);
710	return ret;
711}
712
713static int test_mlockall(int (test_function)(bool call_mlock))
714{
715	int ret = 1;
716
717	if (mlockall(MCL_CURRENT | MCL_ONFAULT | MCL_FUTURE)) {
718		perror("mlockall");
719		return ret;
720	}
721
722	ret = test_function(false);
723	munlockall();
724	return ret;
725}
726
727int main(int argc, char **argv)
728{
729	int ret = 0;
730	ret += test_mlock_lock();
731	ret += test_mlock_onfault();
732	ret += test_munlockall();
733	ret += test_lock_onfault_of_present();
734	ret += test_vma_management(true);
735	ret += test_mlockall(test_vma_management);
736	return ret;
737}
738