I now have two solutions to my own question. They're compiled and ran in the exact same way as my original question.

They are based on yyyy's answer and suggestions in the comments, so massive thanks to them!

As noted in the comments of the below solution, I think this GCC page is saying that how I'm setting the rsp and rbp registers here is UB? If so, I'd love to hear of alternative ways to set them:

Another restriction is that the clobber list should not contain the stack pointer register. This is because the compiler requires the value of the stack pointer to be the same after an asm statement as it was on entry to the statement. However, previous versions of GCC did not enforce this rule and allowed the stack pointer to appear in the list, with unclear semantics. This behavior is deprecated and listing the stack pointer may become an error in future versions of GCC.

Solution 1

Here are its steps:

1. mmap() to create a new stack (that has a guard page provided by MAP_GROWSDOWN)
2. __asm__ volatile() to save, set, and restore the rsp and rbp registers to this block of mmap()ed memory

#include <assert.h>
#include <dlfcn.h>
#include <jni.h>
#include <setjmp.h>
#include <signal.h>
#include <stdlib.h>
#include <sys/mman.h>
#include <unistd.h>

jmp_buf jmp_buffer;

static char *base;
static char *top;

static void segv_handler(int sig) {
    (void)sig;
    siglongjmp(jmp_buffer, 1);
}

static void mod_fn() {
    if (sigsetjmp(jmp_buffer, 1)) {
        fprintf(stderr, "Jumped %p %p %ldn", base, top, (base - top) / 1024);
        return;
    }

    char c;
    top = &c;
    while (1) {
        top--;
        *top = 1;
    }
}

JNIEXPORT void JNICALL Java_Main_foo(JNIEnv *env, jobject obj) {
    (void)env;
    (void)obj;

    char b;
    base = &b;

    struct sigaction sigsegv_sa = {
        .sa_handler = segv_handler,
        .sa_flags = SA_ONSTACK, // SA_ONSTACK gives SIGSEGV its own stack
    };

    // Set up an emergency stack for SIGSEGV
    // See https://stackoverflow.com/a/7342398/13279557
    static char emergency_stack[SIGSTKSZ];
    stack_t ss = {
        .ss_size = SIGSTKSZ,
        .ss_sp = emergency_stack,
    };

    if (sigaltstack(&ss, NULL) == -1) {
        perror("sigaltstack");
        exit(EXIT_FAILURE);
    }

    if (sigfillset(&sigsegv_sa.sa_mask) == -1) {
        perror("sigfillset");
        exit(EXIT_FAILURE);
    }

    if (sigaction(SIGSEGV, &sigsegv_sa, NULL) == -1) {
        perror("sigaction");
        exit(EXIT_FAILURE);
    }

    void *dll = dlopen("./mage.so", RTLD_NOW);
    if (!dll) {
        fprintf(stderr, "dlopen(): %sn", dlerror());
    }

    size_t page_count = 8192;
    size_t page_size = sysconf(_SC_PAGE_SIZE);
    size_t length = page_count * page_size;

    void *map = mmap(NULL, length, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS | MAP_GROWSDOWN, -1, 0);
    if (map == MAP_FAILED) {
        perror("mmap");
        exit(EXIT_FAILURE);
    }

    // Asserting 16-byte alignment here is not necessary,
    // since mmap() guarantees it with the args we pass it
    assert(((size_t)map & 0xf) == 0);

    void *stack = (char *)map + length;

    // Save rbp and rsp
    // Marking these static is necessary for restoring
    static int64_t rsp;
    static int64_t rbp;
    __asm__ volatile("mov %%rsp, %0nt" : "=r" (rsp));
    __asm__ volatile("mov %%rbp, %0nt" : "=r" (rbp));

    // Set rbp and rsp to the very start of the mmap-ed memory
    //
    // TODO: I think setting rsp and rbp here is UB?:
    // "Another restriction is that the clobber list should not contain
    // the stack pointer register. This is because the compiler requires
    // the value of the stack pointer to be the same after an asm statement
    // as it was on entry to the statement. However, previous versions
    // of GCC did not enforce this rule and allowed the stack pointer
    // to appear in the list, with unclear semantics. This behavior
    // is deprecated and listing the stack pointer may become an error
    // in future versions of GCC."
    // From https://gcc.gnu.org/onlinedocs/gcc/Extended-Asm.html
    __asm__ volatile("mov %0, %%rspnt" : : "r" (stack));
    __asm__ volatile("mov %0, %%rbpnt" : : "r" (stack));

    mod_fn();

    // Restore rbp and rsp
    __asm__ volatile("mov %0, %%rspnt" : : "r" (rsp));
    __asm__ volatile("mov %0, %%rbpnt" : : "r" (rbp));

    if (munmap(map, length) == -1) {
        perror("munmap");
        exit(EXIT_FAILURE);
    }

    printf("Success!n");
}

Solution 2

Here are its steps:

1. mmap() to create a new stack (that has a guard page provided by MAP_GROWSDOWN)
2. __asm__ volatile() to set the rsp and rbp registers to this new block of mmap()ed memory
3. setjmp() + longjmp() to restore the rsp and rbp registers, and whatever needs to be preserved across fn calls

#include <assert.h>
#include <dlfcn.h>
#include <jni.h>
#include <setjmp.h>
#include <signal.h>
#include <stdlib.h>
#include <sys/mman.h>
#include <unistd.h>

jmp_buf jmp_buffer;
jmp_buf mmap_jmp_buffer;

static char *base;
static char *top;

static void segv_handler(int sig) {
    (void)sig;
    siglongjmp(jmp_buffer, 1);
}

static void mod_fn() {
    if (sigsetjmp(jmp_buffer, 1)) {
        fprintf(stderr, "Jumped %p %p %ldn", base, top, (base - top) / 1024);
        return;
    }

    char c;
    top = &c;
    while (1) {
        top--;
        *top = 1;
    }
}

JNIEXPORT void JNICALL Java_Main_foo(JNIEnv *env, jobject obj) {
    (void)env;
    (void)obj;

    char b;
    base = &b;

    struct sigaction sigsegv_sa = {
        .sa_handler = segv_handler,
        .sa_flags = SA_ONSTACK, // SA_ONSTACK gives SIGSEGV its own stack
    };

    // Set up an emergency stack for SIGSEGV
    // See https://stackoverflow.com/a/7342398/13279557
    static char emergency_stack[SIGSTKSZ];
    stack_t ss = {
        .ss_size = SIGSTKSZ,
        .ss_sp = emergency_stack,
    };

    if (sigaltstack(&ss, NULL) == -1) {
        perror("sigaltstack");
        exit(EXIT_FAILURE);
    }

    if (sigfillset(&sigsegv_sa.sa_mask) == -1) {
        perror("sigfillset");
        exit(EXIT_FAILURE);
    }

    if (sigaction(SIGSEGV, &sigsegv_sa, NULL) == -1) {
        perror("sigaction");
        exit(EXIT_FAILURE);
    }

    void *dll = dlopen("./mage.so", RTLD_NOW);
    if (!dll) {
        fprintf(stderr, "dlopen(): %sn", dlerror());
    }

    size_t page_count = 8192;
    size_t page_size = sysconf(_SC_PAGE_SIZE);
    size_t length = page_count * page_size;

    void *map = mmap(NULL, length, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS | MAP_GROWSDOWN, -1, 0);
    if (map == MAP_FAILED) {
        perror("mmap");
        exit(EXIT_FAILURE);
    }

    // Asserting 16-byte alignment here is not necessary,
    // since mmap() guarantees it with the args we pass it
    assert(((size_t)map & 0xf) == 0);

    void *stack = (char *)map + length;

    if (setjmp(mmap_jmp_buffer) == 0) {
        // Set rbp and rsp to the very start of the mmap-ed memory
        //
        // TODO: I think setting rsp and rbp here is UB?:
        // "Another restriction is that the clobber list should not contain
        // the stack pointer register. This is because the compiler requires
        // the value of the stack pointer to be the same after an asm statement
        // as it was on entry to the statement. However, previous versions
        // of GCC did not enforce this rule and allowed the stack pointer
        // to appear in the list, with unclear semantics. This behavior
        // is deprecated and listing the stack pointer may become an error
        // in future versions of GCC."
        // From https://gcc.gnu.org/onlinedocs/gcc/Extended-Asm.html
        __asm__ volatile("mov %0, %%rspnt" : : "r" (stack));
        __asm__ volatile("mov %0, %%rbpnt" : : "r" (stack));

        mod_fn();

        // Restore rbp and rsp, and whatever needs to be preserved
        // across fn calls: https://stackoverflow.com/a/25266891/13279557
        longjmp(mmap_jmp_buffer, 1);
    }

    if (munmap(map, length) == -1) {
        perror("munmap");
        exit(EXIT_FAILURE);
    }

    printf("Success!n");
}