map.c

/*
 * Given a test pattern (number in base10), for each digit D starting from the right,
 * spawn D threads to burn cpu in separate dynamically allocated executable pages,
 * sleep for user-configurable delay, cancel all threads and revert executable pages
 * to PROT_NONE.
 *
 * Example for pattern 432:
 *   spawn 2 threads, sleep, cancel/mprotect, sleep, spawn 3 threads, sleep, cancel/mprotect..
 *
 * The addresses used for the executable mappings are of the form:
 *   0xITDDXXXXXX with:
 *   IT the current test iteration (index of current digit in the pattern)
 *   DD the current digit (number of threads spawned in this iteration) repeated
 *
 *
 * NOTE: This should be considered throw-away code that makes a number of assumptions.
 *       If it breaks, you get to keep the pieces!
 */

#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <unistd.h>

#include <time.h>
#include <errno.h>
#include <string.h>
#include <pthread.h>
#include <sys/mman.h>

#define RESERVE_START  0x10000000000ULL
#define RESERVE_LENGTH 0x2000000000ULL // 128GB

// Can be overridden through cmdline
#define TEST_PATTERN 432
#define DELAY        8
#define PAGE_GAP     1

// Set by main(), size of burn_cpu PIC chunk.
static size_t CODE_LEN;

extern void * __start__burn;
extern void * __stop__burn;

// PIC chunk that burns CPU and can be asynchronously canceled.
static __attribute__((section ("_burn"))) void*
burn_cpu(void *fptr)
{
    int (*set_cancel_type)(int, int*) = fptr;
    int old_type;

    // We're just burning CPU, so asynchronous cancellation is safe
    set_cancel_type(PTHREAD_CANCEL_ASYNCHRONOUS, &old_type);

    for (;;) {}

    // Never reached
    return NULL;
}

// Reserves address space and returns a pointer to the reserved region.
static void*
reserve(void)
{
    for (uintptr_t s=RESERVE_START;; s+=RESERVE_START) {
        // Reserve but do not commit RESERVE_LENGTH bytes of address space.
        // MAP_NORESERVE ensures that the kernel will not check the allocation.
        void *ret = mmap((void*)s, RESERVE_LENGTH, PROT_NONE,
                         MAP_PRIVATE|MAP_ANONYMOUS|MAP_NORESERVE|MAP_FIXED_NOREPLACE,
                         -1, 0);
        if (ret == MAP_FAILED) {
            if (errno == EEXIST) {
                continue;
            }
            perror("reserve/mmap");
            exit(EXIT_FAILURE);
        }
        return ret;
    }
}

// Spawns a new thread, executing burn_cpu() in a page starting at addr.
static void
spawn(pthread_t *th, uintptr_t addr)
{
    void *dst = mmap((void*)addr, CODE_LEN,
                     PROT_EXEC|PROT_READ|PROT_WRITE,
                     MAP_PRIVATE|MAP_ANONYMOUS|MAP_FIXED,
                     -1, 0);
    if (dst == MAP_FAILED) {
        perror("spawn/mmap");
        exit(EXIT_FAILURE);
    }

    // Copy PIC chunk to addr
    memcpy(dst, burn_cpu, CODE_LEN);

    int ret = pthread_create(th, NULL, dst, pthread_setcanceltype);
    if (ret != 0) {
        fprintf(stderr, "pthread_create: %d\n", ret);
        exit(EXIT_FAILURE);
    }
}

static void
usage(const char *name)
{
    fprintf(stderr, "Usage: %s [pattern: %d] [delay: %d] [page_gap: %d]\n",
            name, TEST_PATTERN, DELAY, PAGE_GAP);
    fprintf(stderr, "Options:\n");
    fprintf(stderr, "  pattern   For every digit X of base10 pattern, spawn X threads in X different mappings.\n");
    fprintf(stderr, "  delay     Delay, in seconds, in-between operations.\n");
    fprintf(stderr, "  page_gap  Number of PROT_NONE pages in-between mappings.\n");
    fprintf(stderr, "\n");
}

static void
delay(size_t sec)
{
    if (sec > 0) {
        printf("Sleeping for %zus\n", sec);
        sleep(sec);
    }
}

int
main(int argc, char **argv)
{
    size_t test_pattern = TEST_PATTERN;
    size_t delay_sec    = DELAY;
    size_t page_gap     = PAGE_GAP;

    if (argc > 4 || (argc == 2 && !strcmp(argv[1], "-h"))) {
        usage(argv[0]);
        exit(EXIT_FAILURE);
    }

    if (argc > 1) {
        if (sscanf(argv[1], "%zu", &test_pattern) != 1) {
            fprintf(stderr, "sscanf: error parsing test pattern\n");
            exit(EXIT_FAILURE);
        }
    }

    if (argc > 2) {
        if (sscanf(argv[2], "%zu", &delay_sec) != 1) {
            fprintf(stderr, "sscanf: error parsing delay\n");
            exit(EXIT_FAILURE);
        }
    }

    if (argc > 3) {
        if (sscanf(argv[3], "%zu", &page_gap) != 1) {
            fprintf(stderr, "sscanf: error parsing page gap\n");
            exit(EXIT_FAILURE);
        }
    }

    size_t code_start = (size_t)&__start__burn;
    size_t code_stop  = (size_t)&__stop__burn;
    void *region = reserve();
    long page_sz = sysconf(_SC_PAGESIZE);
    CODE_LEN = code_stop - code_start;

    printf("Reserved 0x%llX bytes at %p\n", RESERVE_LENGTH, region);
    printf("PID: %d\n", getpid());
    printf("Code len: %zu\n", CODE_LEN);
    printf("Page size: %lu\n", page_sz);
    printf("=== ARGS:\n");
    printf("Test pattern: %zu\n", test_pattern);
    printf("Delay: %zu\n", delay_sec);
    printf("Page gap: %zu\n", page_gap);
    printf("\n");

    page_gap += 1;

    pthread_t tids[10];
    for (size_t iter=0; test_pattern > 0; iter++, test_pattern /= 10) {
        size_t max_threads = test_pattern % 10;
        uintptr_t base = (uintptr_t)region + (iter<<32);

        for (size_t i=0; i<max_threads; i++) {
            uintptr_t addr = base + (max_threads<<28) + (max_threads << 24) + i*page_sz*page_gap;
            spawn(&tids[i], addr);
            printf("+ 0x%zX\n", (size_t)addr);
        }

        delay(delay_sec);

        for (size_t i=0; i<max_threads; i++) {
            uintptr_t addr = base + (max_threads<<28) + (max_threads << 24) + i*page_sz*page_gap;
            void *tret;

            pthread_cancel(tids[i]);
            pthread_join(tids[i], &tret);

            if (tret != PTHREAD_CANCELED) {
                fprintf(stderr, "pthread_join: %p at 0x%zX\n", tret, (size_t)addr);
                exit(EXIT_FAILURE);
            }
            mprotect((void*)addr, CODE_LEN, PROT_NONE);
            printf("- 0x%zX\n", (size_t)addr);
        }

        delay(delay_sec);
    }

    exit(EXIT_SUCCESS);
}

/* Local Variables: */
/* c-basic-offset: 4 */
/* End: */