From 048ae776254fb14bfe1b0c4a1f727d1dc01e3474 Mon Sep 17 00:00:00 2001
From: Andrea Fioraldi <andreafioraldi@gmail.com>
Date: Mon, 4 Oct 2021 10:42:37 +0200
Subject: [PATCH] Fallback on start cpustate when current_cpu is null

---
 cpu.c             | 17 ++++++++++++++---
 linux-user/main.c |  2 +-
 2 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/cpu.c b/cpu.c
index dbfdf9a473..f9652523ee 100644
--- a/cpu.c
+++ b/cpu.c
@@ -63,6 +63,8 @@ struct libafl_hook {
 
 struct libafl_hook* libafl_qemu_hooks = NULL;
 
+CPUArchState *libafl_qemu_env;
+
 void libafl_helper_table_add(TCGHelperInfo* info);
 
 static GByteArray *libafl_qemu_mem_buf = NULL;
@@ -79,7 +81,10 @@ int libafl_qemu_write_reg(int reg, uint8_t* val)
 {
     CPUState *cpu = current_cpu;
     if (!cpu) {
-        return 0;
+        cpu = env_cpu(libafl_qemu_env);
+        if (!cpu) {
+            return 0;
+        }
     }
 
     CPUClass *cc = CPU_GET_CLASS(cpu);
@@ -93,7 +98,10 @@ int libafl_qemu_read_reg(int reg, uint8_t* val)
 {
     CPUState *cpu = current_cpu;
     if (!cpu) {
-        return 0;
+        cpu = env_cpu(libafl_qemu_env);
+        if (!cpu) {
+            return 0;
+        }
     }
 
     if (libafl_qemu_mem_buf == NULL) {
@@ -116,7 +124,10 @@ int libafl_qemu_num_regs(void)
 {
     CPUState *cpu = current_cpu;
     if (!cpu) {
-        return 0;
+        cpu = env_cpu(libafl_qemu_env);
+        if (!cpu) {
+            return 0;
+        }
     }
 
     CPUClass *cc = CPU_GET_CLASS(cpu);
diff --git a/linux-user/main.c b/linux-user/main.c
index 86eb0936db..dcb969d17e 100644
--- a/linux-user/main.c
+++ b/linux-user/main.c
@@ -630,7 +630,7 @@ uint64_t libafl_load_addr(void);
 int libafl_qemu_main(void);
 int libafl_qemu_run(void);
 
-static CPUArchState *libafl_qemu_env;
+extern CPUArchState *libafl_qemu_env;
 
 struct image_info libafl_image_info;