Add premium mode. Closes #3.

activation_key.txt

@@ -0,0 +1 @@
+THIS IS A TEST. JUST A TEST. NOTHING BUT A TEST. YOU'VE BEEN WARNED.

debug_docker/Dockerfile

@@ -6,7 +6,7 @@ FROM debian:bullseye
 
 
 
-RUN apt update -y && apt upgrade -y && apt install -y build-essential wget cmake tar gdb libc6-dbg python3 file
+RUN apt update -y && apt upgrade -y && apt install -y build-essential wget cmake tar gdb libc6-dbg python3 file strace
 
 
 ############### INSTALL FNETD
@@ -41,6 +41,7 @@ RUN make
 WORKDIR /
 
 RUN cp /home/pwn/build/vuln /home/pwn/vuln
+RUN cp /home/pwn/source/activation_key.txt /home/pwn/activation_key.txt
 
 RUN chmod 0755 /home/pwn/vuln
 

tests/common.py

@@ -43,39 +43,45 @@ def instr_r(opcode, reg1, reg2):
     return bytes([opcode, reg1, 0, 0, reg2, 0, 0, 0])
 
 
-def exec_program(program: bytes, debug: bool = False) -> int | None:
+def exec_program(program: bytes, assert_f, debug: bool = False) -> int | None:
     if debug:
         context.log_level = 'debug'
+    else:
+        context.log_level = 'warn'
+
     with remote("localhost", 1337, fam="ipv4") as p:
-        p.recvuntil(b"Password: ", timeout=1)
+        msg = p.recvuntil(b"Password: ", timeout=1)
+        assert_f(msg != b'')
         p.sendline(b"1234")
 
-        msg = p.recvuntil(b"always a Surprise)", timeout=1)
+        msg = p.recvuntil(b"always a Surprise)\n", timeout=1)
-        if msg == b'':
+        assert_f(msg == b'Welcome to JIT-aaS (Just In Time - always a Surprise)\n')
-            return None
+
-        print(msg.decode())
+        msg = p.recvuntil(b"? (y/N):", timeout=1)
+        assert_f(msg == b'Do you want to activate the premium version? (y/N):')
+
+        p.sendline(b"N")
+        msg = p.recvuntil(b"Using the demo version!\n")
+        assert_f(msg == b"Using the demo version!\n")
 
         msg = p.recvuntil(b"should it bee?", timeout=1)
-        if msg == b'':
+        assert_f(msg == b'Now to your next program: How long should it bee?')
-            return None
-        print(msg.decode())
 
         len_msg = str(len(program) // INSTR_LEN).encode()
-        log.info(f"Sending: {len_msg}")
+        log.debug(f"Sending: {len_msg}")
         p.sendline(len_msg)
 
         msg = p.recvuntil(b"Now your program:", timeout=1)
-        if msg == b'':
+        assert_f(msg == b'Now your program:')
-            return None
-        print(msg.decode())
 
-        log.info(f"Sending program: {list(program)}")
+        log.debug(f"Sending program: {list(program)}")
         p.send(program)
 
         msg = p.recvuntil(b"Your program exited with ", timeout=1)
-        if msg == b'':
+        assert_f(msg == b'Your program exited with ')
-            return None
-        print(msg.decode())
 
-        exit_code = int(p.recvuntil(b"!", drop=True, timeout=1))
+        exit_code_msg = p.recvuntil(b"!", drop=True, timeout=1)
+        assert_f(exit_code_msg != b'')
+
+        exit_code = int(exit_code_msg)
         return exit_code

tests/test_add.py

@@ -6,7 +6,7 @@ class BasicInstructionTest(unittest.TestCase):
     def test_addi(self):
         for val in [0, 1, 50, 256, 0x543210, 0xFFFFFFFF]:
             program = instr_i(Opcode.ADDI, Register.A, val)
-            exit_code = exec_program(program)
+            exit_code = exec_program(program, self.assertTrue)
             self.assertIsNotNone(exit_code, "Connection timeout!")
             self.assertEqual(exit_code, val & 0xFF, "Computed the wrong result!")
 
@@ -17,7 +17,7 @@ class BasicInstructionTest(unittest.TestCase):
             val_sum = sum(vals)
             for val in vals:
                 program += instr_i(Opcode.ADDI, Register.A, val)
-            exit_code = exec_program(program)
+            exit_code = exec_program(program, self.assertTrue)
             self.assertIsNotNone(exit_code, "Connection timeout!")
             self.assertEqual(exit_code, val_sum & 0xFF, "Computed the wrong result!")
 
@@ -25,7 +25,7 @@ class BasicInstructionTest(unittest.TestCase):
         for val1, val2 in [(0, 0), (5, 8), (0xFFFFFFFF, 0xFFFFFFFF), (0xFFFFFFFF, 1)]:
             program = instr_i(Opcode.ADDI, Register.A, val1) + instr_i(Opcode.ADDI, Register.B, val2) + instr_r(
                 Opcode.ADD, Register.A, Register.B)
-            exit_code = exec_program(program)
+            exit_code = exec_program(program, self.assertTrue)
             self.assertIsNotNone(exit_code, "Connection timeout!")
             self.assertEqual(exit_code, (val1 + val2) & 0xFF, "Computed the wrong result!")
 

vuln.c

@@ -1,3 +1,4 @@
+#include <fcntl.h>
 #include <stdbool.h>
 #include <stdint.h>
 #include <stdio.h>
@@ -8,6 +9,9 @@
 #include <unistd.h>
 
 #define MAX_PROGRAM_LEN 0x1000
+#define ACTIVATION_KEY_LEN 0x80
+
+static char activation_key[ACTIVATION_KEY_LEN] = {0};
 
 typedef enum Opcode { ADD = 0, ADDI = 1, SUB = 2, COPY = 3, LOADI = 4, COUNT_OPCODES } Opcode;
 
@@ -242,6 +246,11 @@ uint8_t run_jit(Instruction *program, size_t len) {
         exit(EXIT_FAILURE);
     }
 
+    if (premium_activated) {
+        // premium version does not use sandbox
+        exec_func_t exec_func = (exec_func_t)code;
+        return exec_func();
+    } else {
         int child_pid = fork();
         switch (child_pid) {
         case -1:
@@ -278,6 +287,47 @@ uint8_t run_jit(Instruction *program, size_t len) {
 
         return WEXITSTATUS(wstatus);
     }
+}
+
+void check_premium() {
+    printf("Do you want to activate the premium version? (y/N):");
+    char answer[3] = {0};
+    if (fgets(answer, sizeof(answer), stdin) == NULL) {
+        exit(EXIT_FAILURE);
+    }
+
+    switch (answer[0]) {
+    case 'y':
+        break;
+    default:
+        premium_activated = false;
+        return;
+    }
+
+    int secret_fd = open("activation_key.txt", O_RDONLY);
+    if (secret_fd < 0) {
+        puts("Cannot open activation key file!");
+        exit(EXIT_FAILURE);
+    }
+
+    read(secret_fd, activation_key, sizeof(activation_key));
+    close(secret_fd);
+
+    printf("Then please enter your activation key:");
+    char buf[ACTIVATION_KEY_LEN] = {0};
+
+    ssize_t read_bytes = read(0, buf, sizeof(buf));
+    if (read_bytes <= 0) {
+        puts("Cannot read activation key from user!");
+        exit(EXIT_FAILURE);
+    }
+
+    if (buf[read_bytes - 1] == '\n') {
+        buf[read_bytes - 1] = 0;
+    }
+
+    premium_activated = memcmp(activation_key, buf, sizeof(activation_key)) == 0;
+}
 
 int main() {
     // TODO: signal handlers? SIGCHILD? seccomp?
@@ -289,12 +339,18 @@ int main() {
     // TODO: better pun, add reference to pop-culture
     puts("Welcome to JIT-aaS (Just In Time - always a Surprise)");
 
+    check_premium();
+    if (premium_activated) {
+        puts("Using premium version! No sandbox for you!");
+    } else {
+        puts("Using the demo version!");
+    }
+
     Instruction *program;
     size_t program_len;
     int exit_code;
 
     while (true) {
-        // TODO: check for password and enable premium mode
         program = get_program(&program_len);
         if (!validate_program(program, program_len)) {
             puts("Your program is not valid. You possibly use invalid opcodes or registers!");