jsonjit.c

/* Maps a page of the request size near the text page that contains
 * a symbol.  Useful for hand-rolled JIT compilers to avoid indirect
 * calls to symbols in a program on x86-64.  (Part of the code here is
 * a JIT compiler to encode JSON based on a struct descriptor, but it's
 * not finished yet.)
 *
 * Written by L. Pereira <l@tia.mat.br>
 *
 * Thanks to Ole André V. Ravnås for the idea to use the NOREPLACE flag
 * in mmap(), and to Paul Khuong for helping me make it more robust in
 * systems where that flag might not exist.
 * 
 */
#include <assert.h>
#include <errno.h>
#include <stddef.h>
#include <stdint.h>
#include <stdio.h>
#include <sys/mman.h>
#include <unistd.h>
#include <string.h>
#include <stdbool.h>
#include <stdarg.h>
#include <stdlib.h>
#include <limits.h>

#include "lwan-array.h"
#include "json.h"

static void *find_free_page_near_symbol(void *symbol, size_t want_size)
{
    /* This is quite inneficient, being O(n²) and all, but should be fine for the
     * toy I'm working on.
     */
    FILE *maps;
    char buffer[512];
    char *ptr;
    uintptr_t max_want_addr = 0;
    uintptr_t symbol_addr = (uintptr_t)symbol;
    long page_size = sysconf(_SC_PAGE_SIZE);

    maps = fopen("/proc/self/maps", "re");
    if (!maps)
        return NULL;

    want_size = (want_size + (page_size - 1)) & ~(page_size - 1);

    while ((ptr = fgets(buffer, 512, maps))) {
        uintptr_t start, end;

        if (sscanf(ptr, "%lx-%lx", &start, &end) != 2)
            continue;

        if (symbol_addr >= start && symbol_addr <= end) {
           max_want_addr = start - want_size;
           rewind(maps);
           break;
        }
    }

    if (max_want_addr) {
        while ((ptr = fgets(buffer, 512, maps))) {
            uintptr_t start, end;

            if (sscanf(ptr, "%lx-%lx", &start, &end) != 2)
                continue;

            if (start >= max_want_addr && end <= max_want_addr + want_size) {
                max_want_addr -= page_size;
                rewind(maps);
            }
        }
    }

    fclose(maps);

    return (void *)max_want_addr;
}

static void *try_map_near_symbol(void *symbol, size_t want_size, int prot)
{
    int flags = MAP_ANONYMOUS | MAP_FIXED_NOREPLACE | MAP_PRIVATE;
    void *addr = symbol;

    for (int try = 0; try < 32; try ++) {
        addr = find_free_page_near_symbol(addr, want_size);
        if (!addr)
            break;

        void *ptr = mmap(addr, want_size, prot, flags, -1, 0);
        if (ptr == addr)
            return ptr;

        if (ptr != MAP_FAILED) {
            munmap(ptr, want_size);
        } else {
            if (errno == EEXIST && (flags & MAP_FIXED_NOREPLACE)) {
                /* If we get EEXIST here, then we know that this kernel
                 * has NOREPLACE support -- we just got unlucky with a
                 * race condition and lost.  Try finding another address
                 * again. */
                continue;
            }

            /* Any other error condition means we might want to try
             * just passing a hint to mmap() and seeing if the address
             * it returns us is the same that we want. */
        }

        /* This kernel probably has no NOREPLACE support, so disable it
         * for the next tries. */
        flags &= ~MAP_FIXED_NOREPLACE;
    }

    return NULL;
}

struct jit_call_patch {
    void *target;
    size_t patch;
};

DEFINE_ARRAY_TYPE(jit_code, uint8_t)
DEFINE_ARRAY_TYPE(jit_call, struct jit_call_patch)

struct jit {
    struct jit_code code;
    struct jit_call calls;

    int (*append_bytes)(const void *buf, size_t len, void *user_data);

    size_t aux_func_encode_str_loc;

    size_t size_for_unmapping;
    void *mapped;
};

static void jit_emit(struct jit *j, int n, ...)
{
    va_list ap;

    va_start(ap, n);
    while (n--) {
        uint8_t *b = jit_code_append(&j->code);
        *b = (uint8_t)va_arg(ap, int);
    }
    va_end(ap);
}

#define JIT32_ARGS(v)                                                          \
    (uint8_t)((uint32_t)(v) & 0xff), \
    (uint8_t)(((uint32_t)(v) >> 8) & 0xff), \
    (uint8_t)(((uint32_t)(v) >> 16) & 0xff), \
    (uint8_t)(((uint32_t)(v) >> 24) & 0xff)

#define JIT64_ARGS(v)                                                          \
    (uint8_t)((uint64_t)(v) & 0xff), \
    (uint8_t)(((uint64_t)(v) >> 8) & 0xff), \
    (uint8_t)(((uint64_t)(v) >> 16) & 0xff), \
    (uint8_t)(((uint64_t)(v) >> 24) & 0xff), \
    (uint8_t)(((uint64_t)(v) >> 32) & 0xff), \
    (uint8_t)(((uint64_t)(v) >> 40) & 0xff), \
    (uint8_t)(((uint64_t)(v) >> 48) & 0xff), \
    (uint8_t)(((uint64_t)(v) >> 56) & 0xff)

static void jit_emit_call(struct jit *j, void *func)
{
    struct jit_call_patch *patch = jit_call_append(&j->calls);

    jit_emit(j, 1, 0xe8); /* call imm32 */
    patch->target = func;
    patch->patch = jit_code_len(&j->code);
    jit_emit(j, 4, JIT32_ARGS(0xbebacafe)); /* imm32 (to be patched) */
}

static void jit_init(struct jit *j, void *append_bytes)
{
    jit_code_init(&j->code);
    jit_call_init(&j->calls);

    j->append_bytes = append_bytes;
    j->size_for_unmapping = 0;
    j->mapped = NULL;

    /* Prologue */
    jit_emit(j, 1, 0x55);                   /* push rbp */
    jit_emit(j, 3, 0x48, 0x89, 0xe5);       /* mov rbp, rsp */
    jit_emit(j, 4, 0x48, 0x83, 0xec, 0x20); /* sub rsp, 0x20 */
    jit_emit(j, 4, 0x48, 0x89, 0x7d, 0xe8); /* mov -0x18(rbp), rdi */
    jit_emit(j, 4, 0x48, 0x89, 0x75, 0xe0); /* mov -0x20(rbp), rsi */
    jit_emit(j, 7, 0xc7, 0x45, 0xfc, JIT32_ARGS(0)); /* mov -0x4(rbp), 0 */
}

static void jit_emit_append_const_strz(struct jit *j, const char *str)
{
    const size_t len = strlen(str);

    jit_emit(j, 4, 0x48, 0x8b, 0x45, 0xe0);       /* mov rax, -0x20(rbp) */
    jit_emit(j, 3, 0x48, 0x89, 0xc2);             /* mov rdx, rax */
    jit_emit(j, 5, 0xbe, JIT32_ARGS(len));        /* mov esi, len */
    jit_emit(j, 10, 0x48, 0xbf, JIT64_ARGS(str)); /* mov rdi, str */
    jit_emit_call(j, j->append_bytes);
    jit_emit(j, 3, 0x09, 0x45, 0xfc); /* or -0x4(rbp), eax */
}

static void jit_emit_alignment_nops(struct jit *j)
{
    const size_t n_nops = jit_code_len(&j->code) & 31;

    for (size_t i = 0; i < n_nops; i++)
        jit_emit(j, 1, 0x90); /* nop */
}

/*
 * Routines has_zero() and has_value() are from
 * https://graphics.stanford.edu/~seander/bithacks.html#ZeroInWord
 */
static inline uint64_t has_zero(uint64_t v)
{
    return (v - 0x0101010101010101UL) & ~v & 0x8080808080808080UL;
}

static inline uint64_t has_value(uint64_t x, char n)
{
    return has_zero(x ^ (~0UL / 255 * (uint64_t)n));
}

static char escape_as(char chr)
{
    static const char escaped[] = {'"', '\\', 'b', 'f', 'n', 'r', 't', 't'};
    uint64_t mask = has_value(0x225c080c0a0d0909UL, chr);
    return mask == 0 ? 0 : escaped[__builtin_clzl(mask) / 8];
}

static const int jit_aux_func_encode_string;
static void jit_emit_aux_func_encode_string(struct jit *j)
{
    /* FIXME: arguments to append_bytes! */

    j->aux_func_encode_str_loc = jit_code_len(&j->code);

    jit_emit(j, 2, 0x41, 0x54); /* push %r12 */
    jit_emit(j, 1, 0x55); /* push %rbp */
    jit_emit(j, 3, 0x48, 0x89, 0xfd); /* mov %rdi,%rbp */
    jit_emit(j, 1, 0x53); /* push %rbx */
    jit_emit(j, 3, 0x0f, 0xbe, 0x3f); /* movsbl (%rdi),%edi */
    jit_emit(j, 3, 0x40, 0x84, 0xff); /* test %dil,%dil */
    jit_emit(j, 2, 0x74, 0x71); /* je 14c00 <foobar2000+0x80> */
    jit_emit(j, 3, 0x49, 0x89, 0xeb); /* mov %rbp,%r11 */
    jit_emit(j, 3, 0x45, 0x31, 0xe4); /* xor %r12d,%r12d */
    jit_emit(j, 2, 0xeb, 0x27); /* jmp 14bbe <foobar2000+0x3e> */

    jit_emit_alignment_nops(j);
    
    jit_emit_call(j, j->append_bytes);

    jit_emit(j, 3, 0x41, 0x09, 0xc4); /* or %eax,%r12d */

    jit_emit_call(j, j->append_bytes);

    jit_emit(j, 3, 0x48, 0x89, 0xdd); /* mov %rbx,%rbp */
    jit_emit(j, 3, 0x41, 0x09, 0xc4); /* or %eax,%r12d */
    jit_emit(j, 3, 0x0f, 0xbe, 0x3b); /* movsbl (%rbx),%edi */
    jit_emit(j, 3, 0x49, 0x89, 0xdb); /* mov %rbx,%r11 */
    jit_emit(j, 3, 0x40, 0x84, 0xff); /* test %dil,%dil */
    jit_emit(j, 2, 0x74, 0x25); /* je 14be3 <foobar2000+0x63> */

/*14bbe:*/
    jit_emit_call(j, escape_as);

    jit_emit(j, 4, 0x49, 0x8d, 0x5b, 0x01); /* lea 0x1(%r11),%rbx */
    jit_emit(j, 2, 0x84, 0xc0); /* test %al,%al */
    jit_emit(j, 2, 0x74, 0xe8); /* je 14bb3 <foobar2000+0x33> */
    jit_emit(j, 3, 0x4c, 0x39, 0xdd); /* cmp %r11,%rbp */
    jit_emit(j, 2, 0x75, 0xd0); /* jne 14ba0 <foobar2000+0x20> */

    jit_emit_call(j, j->append_bytes);

    jit_emit(j, 3, 0x0f, 0xbe, 0x3b); /* movsbl (%rbx),%edi */
    jit_emit(j, 3, 0x41, 0x09, 0xc4); /* or %eax,%r12d */
    jit_emit(j, 3, 0x49, 0x89, 0xdb); /* mov %rbx,%r11 */
    jit_emit(j, 3, 0x40, 0x84, 0xff); /* test %dil,%dil */
    jit_emit(j, 2, 0x75, 0xdb); /* jne 14bbe <foobar2000+0x3e> */

/* 14be3: */
    jit_emit(j, 3, 0x48, 0x39, 0xdd); /* cmp %rbx,%rbp */
    jit_emit(j, 2, 0x74, 0x08); /* je 14bf0 <foobar2000+0x70> */

    jit_emit_call(j, j->append_bytes);

    jit_emit(j, 3, 0x41, 0x09, 0xc4); /* or %eax,%r12d */

/* 14bf0: */
    jit_emit(j, 1, 0x5b); /* pop %rbx */
    jit_emit(j, 1, 0x5d); /* pop %rbp */
    jit_emit(j, 3, 0x44, 0x89, 0xe0); /* mov %r12d,%eax */
    jit_emit(j, 2, 0x41, 0x5c); /* pop %r12 */
    jit_emit(j, 1, 0xc3); /* retq */

    jit_emit_alignment_nops(j);

/* 14c00: */
    jit_emit(j, 3, 0x45, 0x31, 0xe4); /* xor %r12d,%r12d */
    jit_emit(j, 1, 0x5d); /* pop %rbp */
    jit_emit(j, 3, 0x44, 0x89, 0xe0); /* mov %r12d,%eax */
    jit_emit(j, 2, 0x41, 0x5c); /* pop %r12 */
    jit_emit(j, 1, 0xc3); /* retq */
}

static void *jit_generate(struct jit *j)
{
    const struct jit_call_patch *patch;
    void *mapped;

    assert(j->size_for_unmapping == 0);

    /* Epilogue */
    jit_emit(j, 3, 0x8b, 0x45, 0xfc); /* mov eax, -0x4(rbp) */
    jit_emit(j, 1, 0xc9);             /* leave */
    jit_emit(j, 1, 0xc3);             /* ret */

    jit_emit_alignment_nops(j);

    jit_emit_aux_func_encode_string(j);

    mapped = try_map_near_symbol(j->append_bytes, jit_code_len(&j->code),
                                 PROT_READ | PROT_WRITE);
    if (mapped == NULL)
        goto out;

    memcpy(mapped, j->code.base.base, jit_code_len(&j->code));
    LWAN_ARRAY_FOREACH (&j->calls, patch) {
        int32_t imm32 = (int32_t)(patch->target - mapped - patch->patch - 4);
        memcpy((char *)mapped + patch->patch, &imm32, sizeof(imm32));
    }

    if (mprotect(mapped, jit_code_len(&j->code), PROT_READ | PROT_EXEC) < 0) {
        munmap(mapped, jit_code_len(&j->code));
        mapped = NULL;
    }

out:
    if (mapped) {
        j->size_for_unmapping = jit_code_len(&j->code);

        printf("TRY IN GDB: disassemble %p, %p\n", mapped,
               (char *)mapped + jit_code_len(&j->code));
    } else {
        j->size_for_unmapping = 0;
    }
    j->mapped = mapped;

    jit_code_reset(&j->code);
    jit_call_reset(&j->calls);

    return mapped;
}

void jit_destroy(struct jit *j)
{
    jit_code_reset(&j->code);
    jit_call_reset(&j->calls);
    if (j->size_for_unmapping)
        munmap(j->mapped, j->size_for_unmapping);
}


void jit_emit_json_value(struct jit *j, const struct json_obj_descr *descr)
{
    switch (descr->type) {
    case JSON_TOK_STRING:
        jit_emit_append_const_strz(j, "\"");
//        jit_emit_json_escaped_str(j, descr);
        jit_emit_append_const_strz(j, "\"");
        break;
    default:
        printf("can't generate code for type %d\n", descr->type);
        exit(1);
    }
}

void jit_emit_json_object(struct jit *j, const struct json_obj_descr *descrs, size_t n_descr)
{
    jit_emit_append_const_strz(j, "{");

    for (size_t i = 0; i < n_descr; i++) {
        const struct json_obj_descr *descr = &descrs[i];

        jit_emit_append_const_strz(j, descr->field_name + descr->field_name_len);
        jit_emit_json_value(j, descr);

        if (i < n_descr - 1)
            jit_emit_append_const_strz(j, ",");
    }

    jit_emit_append_const_strz(j, "}");
}

static int user_provided_callback(const void *buf, size_t len, void *user_data)
{
    write(1, buf, len);
    return 42;
}

struct hello_world_json {
    const char *message;
};
static const struct json_obj_descr hello_world_json_desc[] = {
    JSON_OBJ_DESCR_PRIM(struct hello_world_json, message, JSON_TOK_STRING),
};

int main()
{
    struct jit j;
    int (*func)(const void *ptr, void *user_data);

    jit_init(&j, user_provided_callback);
    jit_emit_json_object(&j, hello_world_json_desc, 1);
    func = jit_generate(&j);

    asm volatile ("int3");

    if (!func) {
        printf("Could not generate code\n");
        return 1;
    }

    int r = func(NULL, NULL);
    printf("Function returned %d\n", r);

    jit_destroy(&j);
}