From 01934af8bec4485d30bc3294ec2321051f02a80f Mon Sep 17 00:00:00 2001
From: Maier Johannes <johannes2.maier@tum.de>
Date: Thu, 25 Jan 2024 01:24:31 +0100
Subject: [PATCH] Random code addresses (#5)

* Enable random addresses for generated code

* Enable custom stack in sandbox

* Triggering CI?

* Update exploit to work in release mode.
---
 exploit/exploit.py |  9 +++++----
 vuln.c             | 27 +++++++++++++++++++++++----
 2 files changed, 28 insertions(+), 8 deletions(-)

diff --git a/exploit/exploit.py b/exploit/exploit.py
index 79546d0..8e65820 100755
--- a/exploit/exploit.py
+++ b/exploit/exploit.py
@@ -126,9 +126,9 @@ def exec_program(p: pwnlib.tubes.remote.remote, program: bytes) -> int:
 
 def extract_premium_key(is_debug: bool = False, port: int = PORT):
     if is_debug:
-        offset_saved_rip_to_activation_key = 0x396d  # debug mode
+        offset_saved_rip_to_activation_key = 0x38bc  # debug mode
     else:
-        offset_saved_rip_to_activation_key = 0x2832  # release mode
+        offset_saved_rip_to_activation_key = 0x37a6  # release mode
 
     premium_key = b""
 
@@ -216,8 +216,9 @@ def get_flag(p: pwnlib.tubes.remote.remote, is_debug: bool = False):
 
 if __name__ == "__main__":
     context.log_level = 'warn'
+    debug = False
 
-    premium_key = extract_premium_key(is_debug=True)
+    premium_key = extract_premium_key(is_debug=debug)
 
     p = connect(premium_key, True)
-    print(get_flag(p, is_debug=True))
+    print(get_flag(p, is_debug=debug))
diff --git a/vuln.c b/vuln.c
index c6b8ba8..332f63c 100644
--- a/vuln.c
+++ b/vuln.c
@@ -158,8 +158,25 @@ int init_seccomp() {
     return prctl(PR_SET_NO_NEW_PRIVS, 1, 0, 0, 0) || prctl(PR_SET_SECCOMP, SECCOMP_MODE_FILTER, &prog);
 }
 
+size_t gen_random_address() {
+    int rand_fd = open("/dev/random", O_RDONLY);
+    size_t random_addr;
+    read(rand_fd, &random_addr, sizeof(random_addr));
+    close(rand_fd);
+
+    return random_addr & 0xFFFFFFFF000;
+}
+
 void exec_code(uint8_t *code) {
-    exec_func_t exec_func = (exec_func_t)code;
+    register exec_func_t exec_func = (exec_func_t)code;
+
+    // allocate a new stack at a random address => no stale addresses left (remember: stack grows down!)
+    void *custom_stack_addr = (void *)gen_random_address() + 0x1000;
+    if (mmap(custom_stack_addr - 0x1000, 0x1000, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS | MAP_FIXED, -1, 0) == MAP_FAILED) {
+        puts("Cannot allocate custom stack!");
+        exit(EXIT_FAILURE);
+    }
+
     close(0);
     close(1);
     close(2);
@@ -167,6 +184,7 @@ void exec_code(uint8_t *code) {
         puts("Cannot enable seccomp jail!");
         exit(EXIT_FAILURE);
     }
+    __asm__ __volatile__("mov %0, %%rsp" : : "r"(custom_stack_addr) : /*"rsp"*/);
     uint8_t res = exec_func();
     _exit(res);
 }
@@ -260,10 +278,11 @@ uint8_t run_jit(Instruction *program, size_t len) {
     // page alignment
     size_t allocated_code_len = (expected_code_len + 0xFFF) & ~0xFFF;
 
-    // TODO: maybe randomly choose address to make exploitation harder
+    size_t random_addr = gen_random_address();
+
     // allocate memory for context and code
-    uint8_t *code = (uint8_t *)mmap(NULL, allocated_code_len, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
-    if (code == (void *)-1) {
+    uint8_t *code = (uint8_t *)mmap((void *)random_addr, allocated_code_len, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS | MAP_FIXED, -1, 0);
+    if (code == MAP_FAILED) {
         puts("Cannot mmap memory for code.");
         exit(EXIT_FAILURE);
     }