I was recently thinking about how hard it would be to make a JIT compiler. The first question is, how would I actually generate code? As in, actually get arbitrary machine code put into memory at run time to execute?

Turns out it's not that hard.


#include <cstdio>
#include <sys/mman.h>
#include <sys/types.h>
#include <cstring>

int main(int argc, const char * argv[]){


     // What we want to execute. in amd64, 0x90 is a no-op and 0xC3 is ret.
     unsigned char execute[] = {0x90, 0x90, 0x90, 0x90, 0x90, 0x90, 0xC3};

     // Get a proper page to use.
     void *lPage = mmap(nullptr, sizeof(execute), PROT_READ|PROT_WRITE, MAP_ANON|MAP_SHARED, -1, 0);

     // Put our machine code there.
     memcpy(lPage, execute, sizeof(execute));

     // We need execute rights on this memory.
     mprotect(lPage, sizeof(execute), PROT_READ|PROT_EXEC);

     // Tell C++ that this addess is executable, and acts like a function with no args or any return value
     void (*func)(void);
     func = (void (*)(void))lPage;

     // Call it. Segfaults if we failed, so fingers crossed!
     func();

     // We must not have failed! Hurray!
     printf("Success!\n");

     return 0;
}

Gist not embedding right? link directly to Gist.

Note that all these snippets assume you have a Unix-like environment, an amd64 CPU, and are compiling for 64 bits.

Here, I'm making up an array of raw bytes. They are all NOPs (no operations, the CPU sees this and does nothing about it), and finally a 'ret' statement. As long as we are in 64-bits or a have a normal calling convention, 'ret' is just like the proper keyword 'return'.

The kind of funky thing is that we not only need to mark the memory we want to execute as executable (which makes sense, here it's the call to mprotect()), we can't do that on just any memory. Normally, all mapped memory is read/write, but not executable.

In Unix/Posix, we can ask for a memory page with mmap(). This ensure we get a whole page, and assures that the address returned meets a bunch of special rules that we aren't too concerned with the details of. The important part is that addresses returned by mmap can be mprotect'd to arbitrary access usage. Conveniently, we can mmap a page for read/write, copy our machine code to it, and then mark it as executable without too much hassle.

All we have to do then is explain to our C++ compiler that the address[font=courier]lPage[/font] can be called like a function (which it kind of is). Interestingly, you NEED a C-style cast here. C++'s wonderful casts simply don't allow you to cross the data/instruction barrier this way.

But, that's not really a compiler of any sort, it's just injecting arbitrary code into a program.

Well, the array of chars that is our machine code could be modified. Say we want to make up some machine code that adds two arbitrary numbers, but we don't want to load the numbers, we want them written into the machine code itself once they are known.

It would look something like this:


#include <cstdio>
#include <sys/mman.h>
#include <sys/types.h>
#include <cstring>
#include <cctype>
#include <cstdlib>


int main(int argc, const char * argv[]){
     char a = '\0';
     char b = '\0';
     do{
     printf("Enter a number.\n");
          a = getc(stdin);
     } while(!isdigit(a));

     printf("OK, using %c.\n", a);

     do{
     printf("Enter a second number.\n");
          b = getc(stdin);
     } while(!isdigit(b));

     printf("OK, using %c.\n", b);

     printf("Running memory test.\n");

     // What we want to execute.
     unsigned char nop[] = {0x90, 0x90, 0x90, 0x90, 0x90, 0x90, 0xC3};

     void *lPage = mmap(nullptr, 0xFF, PROT_READ|PROT_WRITE, MAP_ANON|MAP_SHARED, -1, 0);

     if((lPage==NULL)||(lPage==(void *)(~NULL))){
          printf("Memory map failed.\n");
          return EXIT_FAILURE;
     }

     memcpy(lPage, nop, sizeof(nop));

     mprotect(lPage, sizeof(nop), PROT_READ|PROT_EXEC);

     void (*func)(void);
     func = (void (*)(void))lPage;

     func();

     printf("Successful memory test. Performing addition:\n");

     char ac[] = {a, 0};
     char bc[] = {b, 0};

     unsigned char adds[] = {
          0x48, 0xC7, 0xC0, (unsigned char)atoi(ac), /*put a in rax.*/
          0, 0, 0, /*align*/
          0x48, 0x83, 0xC0, (unsigned char)atoi(bc), /*Add b to rax, storing result in rax.*/
          0xC3, /*return. Return values are in rax in x86_64.*/
     };
     // Change lPage's protection to be read/write so we can change its contents.
     mprotect(lPage, sizeof(adds), PROT_READ|PROT_WRITE);

     memcpy(lPage, adds, sizeof(adds));

     mprotect(lPage, sizeof(adds), PROT_READ|PROT_EXEC);

     int (*ifunc)(void);
     ifunc = (int (*)(void))lPage;

     int r = ifunc();

     printf("In program-modified machine code, %c + %c = %i.\n", a, b, r);

     return EXIT_SUCCESS;
}

Gist not embedding right? Try a old fashioned link.

So that's actually much cooler. Now, we are generating machine code on the fly!

Of note, we reused the same page from the single mmap() call. You may wonder why we don't just mark the page as read/write/execute all at once. Well, that's what a lot of viruses and malicious programs do! Additionally, having something be bother write and execute is just begging for major security issues (as though this whole process wasn't from the start). Systems like SELinux specifically target programs that have mapped pages that are both writeable and executable at once. It's an issue with some very Windows-centric programs that have been ported in the past.
Coincidentally, if someone ever tells you that Linux being more secure than Windows is not true, point this out to them. At the very least, it makes for an excellent counterargument.


...OK, Back on track!

You know what would be even cooler than just writing our values into the machine code? If we made the code's behaviour even more dynamic. Just modifying data is cool and all, but we could have just coded in addresses and used pointers in our machine code (that also would have been kind of nifty, given that now our machine code would have embedded instance- specific addresses...). What if we actually change both instructions and data to generate our code?


#include <cstdio>
#include <sys/mman.h>
#include <sys/types.h>
#include <cstring>
#include <cctype>
#include <cstdlib>


int main(int argc, const char * argv[]){

     bool Add = true;

     char a = '\0';
     char b = '\0';
     do{
     printf("Enter a number.\n");
          a = getc(stdin);
     } while(!isdigit(a));

     printf("OK, using %c.\n", a);

          char t = '\0';
     do{
          t = getc(stdin);
          printf("Press 'a' to add, 's' to subtract.\n");
     }while(t!='a' && t!='s');

     if(t=='s')
          Add = false;
     do{
     printf("Enter a second number.\n");
          b = getc(stdin);
     } while(!isdigit(b));

     printf("OK, using %c.\n", b);

     printf("Running memory test.\n");

     // What we want to execute.
     unsigned char nop[] = {0x90, 0x90, 0x90, 0x90, 0x90, 0x90, 0xC3};

     void *lPage = mmap(nullptr, 0xFF, PROT_READ|PROT_WRITE, MAP_ANON|MAP_SHARED, -1, 0);

     if((lPage==NULL)||(lPage==(void *)(~NULL))){
          printf("Memory map failed.\n");
          return EXIT_FAILURE;
     }

     memcpy(lPage, nop, sizeof(nop));

     mprotect(lPage, sizeof(nop), PROT_READ|PROT_EXEC);

     void (*func)(void);
     func = (void (*)(void))lPage;

     func();

     printf("Successful memory test. Performing %s:\n", Add?"addition":"subtraction");

     char ac[] = {a, 0};
     char bc[] = {b, 0};

     unsigned char adds[] = {
          0x48, 0xC7, 0xC0, (unsigned char)atoi(ac), /*put a in rax.*/
          0, 0, 0, /*align*/
          0x48, 0x83, (unsigned char)(Add?0xC0:0xE8), (unsigned char)atoi(bc), /*Add:Subtract b to rax, storing result in rax.*/
          0xC3, /*return. Return values are in rax in x86_64.*/
     };
     // Change lPage's protection to be read/write so we can change its contents.
     mprotect(lPage, sizeof(adds), PROT_READ|PROT_WRITE);

     memcpy(lPage, adds, sizeof(adds));

     mprotect(lPage, sizeof(adds), PROT_READ|PROT_EXEC);

     int (*ifunc)(void);
     ifunc = (int (*)(void))lPage;

     int r = ifunc();

     printf("In program-modified machine code, %c %c %c = %i.\n", a, Add?'+':'-', b, r);

     return EXIT_SUCCESS;
}

Gist not embedding right? Try a old fashioned link.

Now that's much more like it.

So, what did I learn from this adventure?

Dynamic code generation and execution is frighteningly easy. I didn't expect this to be so simple, or to work so easily. Of course, this is bordering on the kind of black magic that could destroy any project. It's ridiculous and completely unnecessary. Don't actually do this unless it is the intended product of your program.

...But it's also really fun to do!