Test Intel's Protection Key

pkey.c

/**
 * This is a program to test Intel's Memory Protection Key Unit.
 *
 * Check CPU    support: cat /proc/cpuinfo | grep pku
 * Check Kernel support: man pkeys
 *
 * Amazon Support:
 *   Image: Ubuntu Server 18.04 LTS (HVM), SSD Volume Type
 *   Instance: c5.large
 *
 * Reference:
 *   manpage for pkeys, pkey_alloc, pkey_free, pkey_mprotect.
 */

#define _GNU_SOURCE
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <unistd.h>
#include <sys/mman.h>

void my_read(char* src) {
    printf("[READ BEGIN]\n");
    printf("[READ      ] %s\n", src);
    printf("[READ END  ]\n");
}

void my_write(char* src, char* dst) {
    printf("[WRITE BEGIN]\n");
    strcpy(src, dst);
    printf("[WRITE END  ]\n");
}

void test(char* buf) {
    printf("---------- Read & Write w/o pkey.\n");
    my_read(buf);
    my_write(buf, "------");
    my_read(buf);
}

void test_pkey(char* buf, int pkey) {
    printf("---------- Read & Write w pkey.\n");
    pkey_mprotect(buf, getpagesize(), PROT_READ | PROT_WRITE, pkey);
    my_read(buf);
    my_write(buf, "------");
    my_read(buf);
}

int choice() {
    int choice;
    printf("Available pkey types:\n");
    printf("[1] PKEY_DISABLE_ACCESS\n");
    printf("[2] PKEY_DISABLE_WRITE\n");
    printf("Your choice [1/2]: ");
    scanf("%d", &choice);
    return choice == 1 ? PKEY_DISABLE_ACCESS : PKEY_DISABLE_WRITE;
}

int main() {
    // Allocate a new protection key descriptor
    int pkey_rights = choice();
    int pkey = pkey_alloc(0, pkey_rights);

    // Allocate a page of memory
    char *buf = mmap(NULL, getpagesize(), PROT_READ | PROT_WRITE, 
                     MAP_SHARED | MAP_ANONYMOUS, -1, 0);
    strcpy(buf, "++++++");

    // Run test
    test(buf);
    test_pkey(buf, pkey);

    // Free memory
    free(buf);

    // Free pkey descriptor
    pkey_free(pkey);
}


/* Sample Output */

/*
ubuntu@ip-172-31-11-13:~$ gcc pkey.c && ./a.out
Available pkey types:
[1] PKEY_DISABLE_ACCESS
[2] PKEY_DISABLE_WRITE
Your choice [1/2]: 1
---------- Read & Write w/o pkey.
[READ BEGIN]
[READ      ] ++++++
[READ END  ]
[WRITE BEGIN]
[WRITE END  ]
[READ BEGIN]
[READ      ] ------
[READ END  ]
---------- Read & Write w pkey.
[READ BEGIN]
Segmentation fault (core dumped)

ubuntu@ip-172-31-11-13:~$ gcc pkey.c && ./a.out
Available pkey types:
[1] PKEY_DISABLE_ACCESS
[2] PKEY_DISABLE_WRITE
Your choice [1/2]: 2
---------- Read & Write w/o pkey.
[READ BEGIN]
[READ      ] ++++++
[READ END  ]
[WRITE BEGIN]
[WRITE END  ]
[READ BEGIN]
[READ      ] ------
[READ END  ]
---------- Read & Write w pkey.
[READ BEGIN]
[READ      ] ------
[READ END  ]
[WRITE BEGIN]
Segmentation fault (core dumped)
*/