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