--- /dev/null
+#include <sys/mman.h>
+#include <stdint.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <unistd.h>
+#include <string.h>
+#include <sys/time.h>
+#include <sys/resource.h>
+#include <syscall.h>
+#include <errno.h>
+#include <stdbool.h>
+
+#ifndef MLOCK_ONFAULT
+#define MLOCK_ONFAULT 1
+#endif
+
+#ifndef MCL_ONFAULT
+#define MCL_ONFAULT (MCL_FUTURE << 1)
+#endif
+
+static int mlock2_(void *start, size_t len, int flags)
+{
+#ifdef __NR_mlock2
+       return syscall(__NR_mlock2, start, len, flags);
+#else
+       errno = ENOSYS;
+       return -1;
+#endif
+}
+
+struct vm_boundaries {
+       unsigned long start;
+       unsigned long end;
+};
+
+static int get_vm_area(unsigned long addr, struct vm_boundaries *area)
+{
+       FILE *file;
+       int ret = 1;
+       char line[1024] = {0};
+       char *end_addr;
+       char *stop;
+       unsigned long start;
+       unsigned long end;
+
+       if (!area)
+               return ret;
+
+       file = fopen("/proc/self/maps", "r");
+       if (!file) {
+               perror("fopen");
+               return ret;
+       }
+
+       memset(area, 0, sizeof(struct vm_boundaries));
+
+       while(fgets(line, 1024, file)) {
+               end_addr = strchr(line, '-');
+               if (!end_addr) {
+                       printf("cannot parse /proc/self/maps\n");
+                       goto out;
+               }
+               *end_addr = '\0';
+               end_addr++;
+               stop = strchr(end_addr, ' ');
+               if (!stop) {
+                       printf("cannot parse /proc/self/maps\n");
+                       goto out;
+               }
+               stop = '\0';
+
+               sscanf(line, "%lx", &start);
+               sscanf(end_addr, "%lx", &end);
+
+               if (start <= addr && end > addr) {
+                       area->start = start;
+                       area->end = end;
+                       ret = 0;
+                       goto out;
+               }
+       }
+out:
+       fclose(file);
+       return ret;
+}
+
+static uint64_t get_pageflags(unsigned long addr)
+{
+       FILE *file;
+       uint64_t pfn;
+       unsigned long offset;
+
+       file = fopen("/proc/self/pagemap", "r");
+       if (!file) {
+               perror("fopen pagemap");
+               _exit(1);
+       }
+
+       offset = addr / getpagesize() * sizeof(pfn);
+
+       if (fseek(file, offset, SEEK_SET)) {
+               perror("fseek pagemap");
+               _exit(1);
+       }
+
+       if (fread(&pfn, sizeof(pfn), 1, file) != 1) {
+               perror("fread pagemap");
+               _exit(1);
+       }
+
+       fclose(file);
+       return pfn;
+}
+
+static uint64_t get_kpageflags(unsigned long pfn)
+{
+       uint64_t flags;
+       FILE *file;
+
+       file = fopen("/proc/kpageflags", "r");
+       if (!file) {
+               perror("fopen kpageflags");
+               _exit(1);
+       }
+
+       if (fseek(file, pfn * sizeof(flags), SEEK_SET)) {
+               perror("fseek kpageflags");
+               _exit(1);
+       }
+
+       if (fread(&flags, sizeof(flags), 1, file) != 1) {
+               perror("fread kpageflags");
+               _exit(1);
+       }
+
+       fclose(file);
+       return flags;
+}
+
+static FILE *seek_to_smaps_entry(unsigned long addr)
+{
+       FILE *file;
+       char *line = NULL;
+       size_t size = 0;
+       unsigned long start, end;
+       char perms[5];
+       unsigned long offset;
+       char dev[32];
+       unsigned long inode;
+       char path[BUFSIZ];
+
+       file = fopen("/proc/self/smaps", "r");
+       if (!file) {
+               perror("fopen smaps");
+               _exit(1);
+       }
+
+       while (getline(&line, &size, file) > 0) {
+               if (sscanf(line, "%lx-%lx %s %lx %s %lu %s\n",
+                          &start, &end, perms, &offset, dev, &inode, path) < 6)
+                       goto next;
+
+               if (start <= addr && addr < end)
+                       goto out;
+
+next:
+               free(line);
+               line = NULL;
+               size = 0;
+       }
+
+       fclose(file);
+       file = NULL;
+
+out:
+       free(line);
+       return file;
+}
+
+#define VMFLAGS "VmFlags:"
+
+static bool is_vmflag_set(unsigned long addr, const char *vmflag)
+{
+       char *line = NULL;
+       char *flags;
+       size_t size = 0;
+       bool ret = false;
+       FILE *smaps;
+
+       smaps = seek_to_smaps_entry(addr);
+       if (!smaps) {
+               printf("Unable to parse /proc/self/smaps\n");
+               goto out;
+       }
+
+       while (getline(&line, &size, smaps) > 0) {
+               if (!strstr(line, VMFLAGS)) {
+                       free(line);
+                       line = NULL;
+                       size = 0;
+                       continue;
+               }
+
+               flags = line + strlen(VMFLAGS);
+               ret = (strstr(flags, vmflag) != NULL);
+               goto out;
+       }
+
+out:
+       free(line);
+       fclose(smaps);
+       return ret;
+}
+
+#define SIZE "Size:"
+#define RSS  "Rss:"
+#define LOCKED "lo"
+
+static bool is_vma_lock_on_fault(unsigned long addr)
+{
+       bool ret = false;
+       bool locked;
+       FILE *smaps = NULL;
+       unsigned long vma_size, vma_rss;
+       char *line = NULL;
+       char *value;
+       size_t size = 0;
+
+       locked = is_vmflag_set(addr, LOCKED);
+       if (!locked)
+               goto out;
+
+       smaps = seek_to_smaps_entry(addr);
+       if (!smaps) {
+               printf("Unable to parse /proc/self/smaps\n");
+               goto out;
+       }
+
+       while (getline(&line, &size, smaps) > 0) {
+               if (!strstr(line, SIZE)) {
+                       free(line);
+                       line = NULL;
+                       size = 0;
+                       continue;
+               }
+
+               value = line + strlen(SIZE);
+               if (sscanf(value, "%lu kB", &vma_size) < 1) {
+                       printf("Unable to parse smaps entry for Size\n");
+                       goto out;
+               }
+               break;
+       }
+
+       while (getline(&line, &size, smaps) > 0) {
+               if (!strstr(line, RSS)) {
+                       free(line);
+                       line = NULL;
+                       size = 0;
+                       continue;
+               }
+
+               value = line + strlen(RSS);
+               if (sscanf(value, "%lu kB", &vma_rss) < 1) {
+                       printf("Unable to parse smaps entry for Rss\n");
+                       goto out;
+               }
+               break;
+       }
+
+       ret = locked && (vma_rss < vma_size);
+out:
+       free(line);
+       if (smaps)
+               fclose(smaps);
+       return ret;
+}
+
+#define PRESENT_BIT     0x8000000000000000
+#define PFN_MASK        0x007FFFFFFFFFFFFF
+#define UNEVICTABLE_BIT (1UL << 18)
+
+static int lock_check(char *map)
+{
+       unsigned long page_size = getpagesize();
+       uint64_t page1_flags, page2_flags;
+
+       page1_flags = get_pageflags((unsigned long)map);
+       page2_flags = get_pageflags((unsigned long)map + page_size);
+
+       /* Both pages should be present */
+       if (((page1_flags & PRESENT_BIT) == 0) ||
+           ((page2_flags & PRESENT_BIT) == 0)) {
+               printf("Failed to make both pages present\n");
+               return 1;
+       }
+
+       page1_flags = get_kpageflags(page1_flags & PFN_MASK);
+       page2_flags = get_kpageflags(page2_flags & PFN_MASK);
+
+       /* Both pages should be unevictable */
+       if (((page1_flags & UNEVICTABLE_BIT) == 0) ||
+           ((page2_flags & UNEVICTABLE_BIT) == 0)) {
+               printf("Failed to make both pages unevictable\n");
+               return 1;
+       }
+
+       if (!is_vmflag_set((unsigned long)map, LOCKED)) {
+               printf("VMA flag %s is missing on page 1\n", LOCKED);
+               return 1;
+       }
+
+       if (!is_vmflag_set((unsigned long)map + page_size, LOCKED)) {
+               printf("VMA flag %s is missing on page 2\n", LOCKED);
+               return 1;
+       }
+
+       return 0;
+}
+
+static int unlock_lock_check(char *map)
+{
+       unsigned long page_size = getpagesize();
+       uint64_t page1_flags, page2_flags;
+
+       page1_flags = get_pageflags((unsigned long)map);
+       page2_flags = get_pageflags((unsigned long)map + page_size);
+       page1_flags = get_kpageflags(page1_flags & PFN_MASK);
+       page2_flags = get_kpageflags(page2_flags & PFN_MASK);
+
+       if ((page1_flags & UNEVICTABLE_BIT) || (page2_flags & UNEVICTABLE_BIT)) {
+               printf("A page is still marked unevictable after unlock\n");
+               return 1;
+       }
+
+       if (is_vmflag_set((unsigned long)map, LOCKED)) {
+               printf("VMA flag %s is present on page 1 after unlock\n", LOCKED);
+               return 1;
+       }
+
+       if (is_vmflag_set((unsigned long)map + page_size, LOCKED)) {
+               printf("VMA flag %s is present on page 2 after unlock\n", LOCKED);
+               return 1;
+       }
+
+       return 0;
+}
+
+static int test_mlock_lock()
+{
+       char *map;
+       int ret = 1;
+       unsigned long page_size = getpagesize();
+
+       map = mmap(NULL, 2 * page_size, PROT_READ | PROT_WRITE,
+                  MAP_ANONYMOUS | MAP_PRIVATE, 0, 0);
+       if (map == MAP_FAILED) {
+               perror("test_mlock_locked mmap");
+               goto out;
+       }
+
+       if (mlock2_(map, 2 * page_size, 0)) {
+               if (errno == ENOSYS) {
+                       printf("Cannot call new mlock family, skipping test\n");
+                       _exit(0);
+               }
+               perror("mlock2(0)");
+               goto unmap;
+       }
+
+       if (lock_check(map))
+               goto unmap;
+
+       /* Now unlock and recheck attributes */
+       if (munlock(map, 2 * page_size)) {
+               perror("munlock()");
+               goto unmap;
+       }
+
+       ret = unlock_lock_check(map);
+
+unmap:
+       munmap(map, 2 * page_size);
+out:
+       return ret;
+}
+
+static int onfault_check(char *map)
+{
+       unsigned long page_size = getpagesize();
+       uint64_t page1_flags, page2_flags;
+
+       page1_flags = get_pageflags((unsigned long)map);
+       page2_flags = get_pageflags((unsigned long)map + page_size);
+
+       /* Neither page should be present */
+       if ((page1_flags & PRESENT_BIT) || (page2_flags & PRESENT_BIT)) {
+               printf("Pages were made present by MLOCK_ONFAULT\n");
+               return 1;
+       }
+
+       *map = 'a';
+       page1_flags = get_pageflags((unsigned long)map);
+       page2_flags = get_pageflags((unsigned long)map + page_size);
+
+       /* Only page 1 should be present */
+       if ((page1_flags & PRESENT_BIT) == 0) {
+               printf("Page 1 is not present after fault\n");
+               return 1;
+       } else if (page2_flags & PRESENT_BIT) {
+               printf("Page 2 was made present\n");
+               return 1;
+       }
+
+       page1_flags = get_kpageflags(page1_flags & PFN_MASK);
+
+       /* Page 1 should be unevictable */
+       if ((page1_flags & UNEVICTABLE_BIT) == 0) {
+               printf("Failed to make faulted page unevictable\n");
+               return 1;
+       }
+
+       if (!is_vma_lock_on_fault((unsigned long)map)) {
+               printf("VMA is not marked for lock on fault\n");
+               return 1;
+       }
+
+       if (!is_vma_lock_on_fault((unsigned long)map + page_size)) {
+               printf("VMA is not marked for lock on fault\n");
+               return 1;
+       }
+
+       return 0;
+}
+
+static int unlock_onfault_check(char *map)
+{
+       unsigned long page_size = getpagesize();
+       uint64_t page1_flags;
+
+       page1_flags = get_pageflags((unsigned long)map);
+       page1_flags = get_kpageflags(page1_flags & PFN_MASK);
+
+       if (page1_flags & UNEVICTABLE_BIT) {
+               printf("Page 1 is still marked unevictable after unlock\n");
+               return 1;
+       }
+
+       if (is_vma_lock_on_fault((unsigned long)map) ||
+           is_vma_lock_on_fault((unsigned long)map + page_size)) {
+               printf("VMA is still lock on fault after unlock\n");
+               return 1;
+       }
+
+       return 0;
+}
+
+static int test_mlock_onfault()
+{
+       char *map;
+       int ret = 1;
+       unsigned long page_size = getpagesize();
+
+       map = mmap(NULL, 2 * page_size, PROT_READ | PROT_WRITE,
+                  MAP_ANONYMOUS | MAP_PRIVATE, 0, 0);
+       if (map == MAP_FAILED) {
+               perror("test_mlock_locked mmap");
+               goto out;
+       }
+
+       if (mlock2_(map, 2 * page_size, MLOCK_ONFAULT)) {
+               if (errno == ENOSYS) {
+                       printf("Cannot call new mlock family, skipping test\n");
+                       _exit(0);
+               }
+               perror("mlock2(MLOCK_ONFAULT)");
+               goto unmap;
+       }
+
+       if (onfault_check(map))
+               goto unmap;
+
+       /* Now unlock and recheck attributes */
+       if (munlock(map, 2 * page_size)) {
+               if (errno == ENOSYS) {
+                       printf("Cannot call new mlock family, skipping test\n");
+                       _exit(0);
+               }
+               perror("munlock()");
+               goto unmap;
+       }
+
+       ret = unlock_onfault_check(map);
+unmap:
+       munmap(map, 2 * page_size);
+out:
+       return ret;
+}
+
+static int test_lock_onfault_of_present()
+{
+       char *map;
+       int ret = 1;
+       unsigned long page_size = getpagesize();
+       uint64_t page1_flags, page2_flags;
+
+       map = mmap(NULL, 2 * page_size, PROT_READ | PROT_WRITE,
+                  MAP_ANONYMOUS | MAP_PRIVATE, 0, 0);
+       if (map == MAP_FAILED) {
+               perror("test_mlock_locked mmap");
+               goto out;
+       }
+
+       *map = 'a';
+
+       if (mlock2_(map, 2 * page_size, MLOCK_ONFAULT)) {
+               if (errno == ENOSYS) {
+                       printf("Cannot call new mlock family, skipping test\n");
+                       _exit(0);
+               }
+               perror("mlock2(MLOCK_ONFAULT)");
+               goto unmap;
+       }
+
+       page1_flags = get_pageflags((unsigned long)map);
+       page2_flags = get_pageflags((unsigned long)map + page_size);
+       page1_flags = get_kpageflags(page1_flags & PFN_MASK);
+       page2_flags = get_kpageflags(page2_flags & PFN_MASK);
+
+       /* Page 1 should be unevictable */
+       if ((page1_flags & UNEVICTABLE_BIT) == 0) {
+               printf("Failed to make present page unevictable\n");
+               goto unmap;
+       }
+
+       if (!is_vma_lock_on_fault((unsigned long)map) ||
+           !is_vma_lock_on_fault((unsigned long)map + page_size)) {
+               printf("VMA with present pages is not marked lock on fault\n");
+               goto unmap;
+       }
+       ret = 0;
+unmap:
+       munmap(map, 2 * page_size);
+out:
+       return ret;
+}
+
+static int test_munlockall()
+{
+       char *map;
+       int ret = 1;
+       unsigned long page_size = getpagesize();
+
+       map = mmap(NULL, 2 * page_size, PROT_READ | PROT_WRITE,
+                  MAP_ANONYMOUS | MAP_PRIVATE, 0, 0);
+
+       if (map == MAP_FAILED) {
+               perror("test_munlockall mmap");
+               goto out;
+       }
+
+       if (mlockall(MCL_CURRENT)) {
+               perror("mlockall(MCL_CURRENT)");
+               goto out;
+       }
+
+       if (lock_check(map))
+               goto unmap;
+
+       if (munlockall()) {
+               perror("munlockall()");
+               goto unmap;
+       }
+
+       if (unlock_lock_check(map))
+               goto unmap;
+
+       munmap(map, 2 * page_size);
+
+       map = mmap(NULL, 2 * page_size, PROT_READ | PROT_WRITE,
+                  MAP_ANONYMOUS | MAP_PRIVATE, 0, 0);
+
+       if (map == MAP_FAILED) {
+               perror("test_munlockall second mmap");
+               goto out;
+       }
+
+       if (mlockall(MCL_CURRENT | MCL_ONFAULT)) {
+               perror("mlockall(MCL_CURRENT | MCL_ONFAULT)");
+               goto unmap;
+       }
+
+       if (onfault_check(map))
+               goto unmap;
+
+       if (munlockall()) {
+               perror("munlockall()");
+               goto unmap;
+       }
+
+       if (unlock_onfault_check(map))
+               goto unmap;
+
+       if (mlockall(MCL_CURRENT | MCL_FUTURE)) {
+               perror("mlockall(MCL_CURRENT | MCL_FUTURE)");
+               goto out;
+       }
+
+       if (lock_check(map))
+               goto unmap;
+
+       if (munlockall()) {
+               perror("munlockall()");
+               goto unmap;
+       }
+
+       ret = unlock_lock_check(map);
+
+unmap:
+       munmap(map, 2 * page_size);
+out:
+       munlockall();
+       return ret;
+}
+
+static int test_vma_management(bool call_mlock)
+{
+       int ret = 1;
+       void *map;
+       unsigned long page_size = getpagesize();
+       struct vm_boundaries page1;
+       struct vm_boundaries page2;
+       struct vm_boundaries page3;
+
+       map = mmap(NULL, 3 * page_size, PROT_READ | PROT_WRITE,
+                  MAP_ANONYMOUS | MAP_PRIVATE, 0, 0);
+       if (map == MAP_FAILED) {
+               perror("mmap()");
+               return ret;
+       }
+
+       if (call_mlock && mlock2_(map, 3 * page_size, MLOCK_ONFAULT)) {
+               if (errno == ENOSYS) {
+                       printf("Cannot call new mlock family, skipping test\n");
+                       _exit(0);
+               }
+               perror("mlock(ONFAULT)\n");
+               goto out;
+       }
+
+       if (get_vm_area((unsigned long)map, &page1) ||
+           get_vm_area((unsigned long)map + page_size, &page2) ||
+           get_vm_area((unsigned long)map + page_size * 2, &page3)) {
+               printf("couldn't find mapping in /proc/self/maps\n");
+               goto out;
+       }
+
+       /*
+        * Before we unlock a portion, we need to that all three pages are in
+        * the same VMA.  If they are not we abort this test (Note that this is
+        * not a failure)
+        */
+       if (page1.start != page2.start || page2.start != page3.start) {
+               printf("VMAs are not merged to start, aborting test\n");
+               ret = 0;
+               goto out;
+       }
+
+       if (munlock(map + page_size, page_size)) {
+               perror("munlock()");
+               goto out;
+       }
+
+       if (get_vm_area((unsigned long)map, &page1) ||
+           get_vm_area((unsigned long)map + page_size, &page2) ||
+           get_vm_area((unsigned long)map + page_size * 2, &page3)) {
+               printf("couldn't find mapping in /proc/self/maps\n");
+               goto out;
+       }
+
+       /* All three VMAs should be different */
+       if (page1.start == page2.start || page2.start == page3.start) {
+               printf("failed to split VMA for munlock\n");
+               goto out;
+       }
+
+       /* Now unlock the first and third page and check the VMAs again */
+       if (munlock(map, page_size * 3)) {
+               perror("munlock()");
+               goto out;
+       }
+
+       if (get_vm_area((unsigned long)map, &page1) ||
+           get_vm_area((unsigned long)map + page_size, &page2) ||
+           get_vm_area((unsigned long)map + page_size * 2, &page3)) {
+               printf("couldn't find mapping in /proc/self/maps\n");
+               goto out;
+       }
+
+       /* Now all three VMAs should be the same */
+       if (page1.start != page2.start || page2.start != page3.start) {
+               printf("failed to merge VMAs after munlock\n");
+               goto out;
+       }
+
+       ret = 0;
+out:
+       munmap(map, 3 * page_size);
+       return ret;
+}
+
+static int test_mlockall(int (test_function)(bool call_mlock))
+{
+       int ret = 1;
+
+       if (mlockall(MCL_CURRENT | MCL_ONFAULT | MCL_FUTURE)) {
+               perror("mlockall");
+               return ret;
+       }
+
+       ret = test_function(false);
+       munlockall();
+       return ret;
+}
+
+int main(int argc, char **argv)
+{
+       int ret = 0;
+       ret += test_mlock_lock();
+       ret += test_mlock_onfault();
+       ret += test_munlockall();
+       ret += test_lock_onfault_of_present();
+       ret += test_vma_management(true);
+       ret += test_mlockall(test_vma_management);
+       return ret;
+}