Add tasknumner as param

system/main.c

@@ -27,14 +27,12 @@
 #include "sysemu/runstate.h"
 #include "sysemu/sysemu.h"
 #include "migration/snapshot.h"
+#include <math.h>
 
 #ifdef CONFIG_SDL
 #include <SDL.h>
 #endif
 
-
-
-
 int snapshot_save(const char *name);
 int snapshot_load(const char *name);
 
@@ -71,17 +69,17 @@ int (*qemu_main)(void) = qemu_default_main;
 #include "exec/cpu-common.h"
 void libafl_qemu_set_native_breakpoint(vaddr);
 void libafl_qemu_remove_native_breakpoint(vaddr);
-int libafl_qemu_write_reg(CPUState* cpu, int reg, uint8_t* val);
-int libafl_qemu_read_reg(CPUState* cpu, int reg, uint8_t* val);
-CPUState* libafl_qemu_current_cpu(void);
-int libafl_qemu_num_regs(CPUState* cpu);
+int libafl_qemu_write_reg(CPUState *cpu, int reg, uint8_t *val);
+int libafl_qemu_read_reg(CPUState *cpu, int reg, uint8_t *val);
+CPUState *libafl_qemu_current_cpu(void);
+int libafl_qemu_num_regs(CPUState *cpu);
 int libafl_qemu_num_cpus(void);
-CPUState* libafl_qemu_get_cpu(int cpu_index);
+CPUState *libafl_qemu_get_cpu(int cpu_index);
 int64_t icount_get_raw(void);
 //========= Instrumentation end
 int main(int argc, char **argv)
 {
-    
+
     int input_size;
     printf("argc: %d\n", argc);
     //========= Instrumentation start
@@ -94,65 +92,93 @@ int main(int argc, char **argv)
     hwaddr prep = (hwaddr)strtoll(argv[1], NULL, 16);
     hwaddr start = (hwaddr)strtoll(argv[2], NULL, 16);
     hwaddr end = (hwaddr)strtoll(argv[3], NULL, 16);
-    input_size = atoi(argv[4]);   
-    char* output_path = argv[5]; 
+    input_size = atoi(argv[4]);
+    char *output_path = argv[5];
+    unsigned int num_tasks = atoi(argv[6]);
     // hwaddr target_addr = (hwaddr) strtoll(argv[1], NULL, 16);
     // vm_start();
     // fix arguments for qemu
-    argv[5]=argv[0];
-    argv=&argv[5];
+    argv[6] = argv[0];
+    argv = &argv[6];
 
-    argc -= 5;
+    argc -= 6;
 
-
-    unsigned long deltas[input_size];
-    u_int32_t inputs[input_size];
-    u_int32_t outputs[input_size];
+    int full_input_room = (int)pow(input_size, num_tasks);
+    printf("Full input room: %d\n", full_input_room);
+    unsigned long *deltas = malloc(full_input_room * sizeof(unsigned long));
+    if (deltas == NULL)
+    {
+        fprintf(stderr, "Memory allocation failed\n");
+        exit(1);
+    }
+    u_int32_t *inputs = malloc(full_input_room * sizeof(u_int32_t));
+    if (inputs == NULL)
+    {
+        fprintf(stderr, "Memory allocation failed\n");
+        exit(1);
+    }
     //========= Instrumentation end
     qemu_init(argc, argv);
     //========= Instrumentation start
 
     libafl_qemu_set_native_breakpoint(prep);
-        //set int in in the vm to i
+    // set int in in the vm to i
     vm_start();
     qemu_main_loop();
 
-    //Now execution is halted at the start of the task we want to measure in order to write the input to a register
+    // Now execution is halted at the start of the task we want to measure in order to write the input to a register
     libafl_qemu_remove_native_breakpoint(prep);
 
     snapshot_save("base");
 
     uint8_t register_in_32b[4];
-    uint8_t reg_tmp_val[4];
-    uint8_t code_output[4];
+    uint8_t reg_tmp_val[num_tasks][4];
     // load input
     // cpu_physical_memory_rw(target_addr, buffer, read_len, true);
-    
-    for (u_int32_t i = 0; i < input_size; i++)
+
+    u_int32_t task_inputs[num_tasks];
+
+    fclose(fopen(output_path, "w"));
+
+    FILE *fptr = fopen(output_path, "a");
+
+    // input of all tasks combined
+    for (long i = 0; i < full_input_room; i++)
     {
-        //load the system in the halted state at the beginning of the task; Write input to register
+        printf("Input: %lu\n", i);
+        for (int j = 0; j < num_tasks; j++)
+        {
+            // from the "global input" i, extract the input bits for the task j
+            task_inputs[j] = (i >> (j * (__builtin_popcount(input_size - 1)))) & (input_size - 1);
+            printf("Task %d: %d\n", j + 1, task_inputs[j]);
+        }
+
+        // load the system in the halted state at the beginning of the task; Write input to register
         snapshot_load("base");
         CPUState *cpu = libafl_qemu_get_cpu(0);
-        if (cpu == NULL) {
+        if (cpu == NULL)
+        {
             printf("Error: CPU is NULL.\n");
-        } 
-        //printf("reg count: %d\n", libafl_qemu_num_regs(cpu));
-
-        //Write i to register format
-        register_in_32b[0] = i & 0xFF; // Least significant byte
-        register_in_32b[1] = (i >> 8) & 0xFF;
-        register_in_32b[2] = (i >> 16) & 0xFF;
-        register_in_32b[3] = (i >> 24) & 0xFF; // Most significant byte
-
-
-        int length = libafl_qemu_read_reg(cpu, 12, reg_tmp_val);
-        if(length != 4) {
-            printf("Error: Could not read register\n");
         }
-        libafl_qemu_write_reg(cpu, 12, register_in_32b);
-        
-        //Read Result to unint32_t (for debugging)
-        //uint32_t res_val = (uint32_t)res_ptr[0] | ((uint32_t)res_ptr[1] << 8) | ((uint32_t)res_ptr[2] << 16) | ((uint32_t)res_ptr[3] << 24);
+        // printf("reg count: %d\n", libafl_qemu_num_regs(cpu));
+        for (int j = 0; j < num_tasks; j++)
+        {
+            // Initialize register_in_32b array
+            memset(register_in_32b, 0, sizeof(register_in_32b));
+
+            // Write i to register format
+            register_in_32b[0] = task_inputs[j] & 0xFF; // Least significant byte
+            register_in_32b[1] = (task_inputs[j] >> 8) & 0xFF;
+            register_in_32b[2] = (task_inputs[j] >> 16) & 0xFF;
+            register_in_32b[3] = (task_inputs[j] >> 24) & 0xFF; // Most significant byte
+
+            int length = libafl_qemu_read_reg(cpu, j + 1, reg_tmp_val[j]);
+            if (length != 4)
+            {
+                printf("Error: Could not read register\n");
+            }
+            libafl_qemu_write_reg(cpu, j + 1, register_in_32b);
+        }
 
         libafl_qemu_set_native_breakpoint(start);
 
@@ -163,40 +189,39 @@ int main(int argc, char **argv)
         libafl_qemu_remove_native_breakpoint(start);
         libafl_qemu_set_native_breakpoint(end);
 
-        //Write back the original value to the register
-        libafl_qemu_write_reg(cpu, 12, reg_tmp_val);
+        // Write back the original value to the register
+        for (int j = 0; j < num_tasks; j++)
+        {
+            libafl_qemu_write_reg(cpu, j + 1, reg_tmp_val[j]);
+        }
 
         unsigned long start_count = icount_get_raw();
-        //printf("Start: %lu\n", start_count);
+        // printf("Start: %lu\n", start_count);
         vm_start();
         qemu_main_loop();
         libafl_qemu_remove_native_breakpoint(end);
 
-        length = libafl_qemu_read_reg(cpu, 11, code_output);
-        if(length != 4) {
-            printf("Error: Could not read register\n");
-        }
-        u_int32_t output = (u_int32_t)code_output[0] | ((u_int32_t)code_output[1] << 8) | ((u_int32_t)code_output[2] << 16) | ((u_int32_t)code_output[3] << 24);
-        //printf("Output: %u\n", output);
-
         unsigned long end_count = icount_get_raw();
-        //printf("End: %lu\n", end_count);
+        // printf("End: %lu\n", end_count);
         inputs[i] = i;
-        outputs[i] = output;
         deltas[i] = end_count - start_count;
 
-        //printf("Delta: %lu\n", deltas[i]);
+        printf("Delta: %lu\n", deltas[i]);
+
+        fprintf(fptr, "%d", inputs[i]);
+        for (int j = 0; j < num_tasks; j++)
+        {
+            fprintf(fptr, ",%d", task_inputs[j]);
+        }
+        fprintf(fptr, ",%lu\n", deltas[i]);
     }
 
     // // Write to serial port
     // qemu_chr_fe_write(serial_chr, data, length);
-    FILE *fptr = fopen(output_path, "w");
-
-    for (int i = 0; i < input_size; i++) {
-        fprintf(fptr, "%d,%lu,%d\n", inputs[i], deltas[i],outputs[i]);
-    }
 
     fclose(fptr);
+    free(deltas);
+    free(inputs);
     // // Write some text to the file
     // fprintf(fptr, "%lu",delta);
     return 0;