jit.c

#include <stdio.h>
#include <string.h>
#include <stdint.h>
#include <sys/mman.h>
#include <fcntl.h>
#include <unistd.h>

const char TEXTSENTINEL[] = "__TEXTSENTINEL";
size_t assemble(const char* asmsrc, uint8_t* execbuf, size_t execbuflen) {
    char asmpath[] = "/tmp/asm-XXXXXX"; mkstemp(asmpath); {
        char asmbuf[1024]; snprintf(asmbuf, sizeof(asmbuf),
                                    ".asciz \"%s\"\n"
                                    "%s",
                                    TEXTSENTINEL, asmsrc);
        FILE *asmfp = fopen(asmpath, "wb");
        fwrite(asmbuf, sizeof(char), strlen(asmbuf), asmfp); fclose(asmfp);
    }
    uint8_t objbuf[1024]; size_t objlen; {
        char as[1024]; snprintf(as, sizeof(as), "as %s -o -", asmpath);
        FILE *asfp = popen(as, "r");
        for (objlen = 0; objlen < sizeof(objbuf); objlen++) {
            int ch; if ((ch = fgetc(asfp)) == EOF) { break; }
            objbuf[objlen] = ch;
        } pclose(asfp);
    }
    for (size_t i = 0; i < objlen; i++) {
        if (memcmp(&objbuf[i], TEXTSENTINEL, sizeof(TEXTSENTINEL)) == 0) {
            size_t execlen = objlen - i - sizeof(TEXTSENTINEL);
            if (execlen > execbuflen) execlen = execbuflen;
            memcpy(execbuf, &objbuf[i] + sizeof(TEXTSENTINEL),
                   execlen); return execlen;
        }
    } return 0;
}

typedef uint64_t (*fn_ptr_t)();
int main() {
    char execpath[] = "/tmp/exec-XXXXXX"; mkstemp(execpath);
    int fd = open(execpath, O_RDWR | O_CREAT, 0600); ftruncate(fd, 1024);
    uint8_t* execbuf = mmap(0, 1024, PROT_READ | PROT_WRITE | PROT_EXEC,
                            MAP_SHARED, fd, 0);
    assemble("addq %rsi, %rdi\n"
             "movq %rdi, %rax\n"
             "ret",
             execbuf, 1024);
    fn_ptr_t f = (fn_ptr_t) execbuf;
    printf("1 + 2 = %llu\n", f(1, 2));
}