From c2c69cfc528398d9db9363b92f8c50db4008c98f Mon Sep 17 00:00:00 2001
From: Sergej Schumilo <sergej@schumilo.de>
Date: Fri, 21 Jan 2022 20:23:52 +0100
Subject: [PATCH] abort if a configuration was not set or received (via
 GET_HOST / SET_AGENT) or if either was executed twice

---
 nyx/hypercall/configuration.c | 12 ++++++++++++
 nyx/hypercall/hypercall.c     | 11 +++++++++++
 nyx/state/snapshot_state.c    |  4 ++++
 nyx/state/state.c             |  3 +++
 nyx/state/state.h             |  3 +++
 5 files changed, 33 insertions(+)

diff --git a/nyx/hypercall/configuration.c b/nyx/hypercall/configuration.c
index b31baf0f98..3b54415653 100644
--- a/nyx/hypercall/configuration.c
+++ b/nyx/hypercall/configuration.c
@@ -12,6 +12,11 @@ void handle_hypercall_kafl_get_host_config(struct kvm_run *run, CPUState *cpu, u
 		return;
 	}
 
+	if (GET_GLOBAL_STATE()->get_host_config_done){
+		nyx_abort((char*)"KVM_EXIT_KAFL_GET_HOST_CONFIG called twice...");
+		return;
+	}
+
 	memset((void*)&config, 0, sizeof(host_config_t));
 
 	config.host_magic = NYX_HOST_MAGIC;
@@ -21,6 +26,7 @@ void handle_hypercall_kafl_get_host_config(struct kvm_run *run, CPUState *cpu, u
 	config.payload_buffer_size = GET_GLOBAL_STATE()->shared_payload_buffer_size;
 
 	write_virtual_memory(vaddr, (uint8_t*)&config, sizeof(host_config_t), cpu);
+	GET_GLOBAL_STATE()->get_host_config_done = true;
 }
 
 void handle_hypercall_kafl_set_agent_config(struct kvm_run *run, CPUState *cpu, uint64_t hypercall_arg){
@@ -31,6 +37,11 @@ void handle_hypercall_kafl_set_agent_config(struct kvm_run *run, CPUState *cpu,
 		return;
 	}
 
+	if (GET_GLOBAL_STATE()->set_agent_config_done){
+		nyx_abort((char*)"KVM_EXIT_KAFL_SET_AGENT_CONFIG called twice...");
+		return;
+	}
+
 	X86CPU *cpux86 = X86_CPU(cpu);
   	CPUX86State *env = &cpux86->env;
 
@@ -89,4 +100,5 @@ void handle_hypercall_kafl_set_agent_config(struct kvm_run *run, CPUState *cpu,
 		fprintf(stderr, "[QEMU-Nyx] Error: %s - failed (vaddr: 0x%lx)!\n", __func__, vaddr);
 		exit(1);
 	}
+	GET_GLOBAL_STATE()->set_agent_config_done = true;
 }
\ No newline at end of file
diff --git a/nyx/hypercall/hypercall.c b/nyx/hypercall/hypercall.c
index 3fae6e54c1..64703657e2 100644
--- a/nyx/hypercall/hypercall.c
+++ b/nyx/hypercall/hypercall.c
@@ -108,6 +108,12 @@ bool handle_hypercall_kafl_next_payload(struct kvm_run *run, CPUState *cpu, uint
 			synchronization_lock();
 
 		} else {
+
+			if (GET_GLOBAL_STATE()->set_agent_config_done == false){
+				nyx_abort((char*)"KVM_EXIT_KAFL_SET_AGENT_CONFIG was not called...");
+				return false;
+			}
+
 			if(!setup_snapshot_once){ 
 				//pt_reset_bitmap();
 
@@ -193,6 +199,11 @@ static void handle_hypercall_get_payload(struct kvm_run *run, CPUState *cpu, uin
 		return;
 	}
 
+	if (GET_GLOBAL_STATE()->get_host_config_done == false){
+		nyx_abort((char*)"KVM_EXIT_KAFL_GET_HOST_CONFIG was not called...");
+		return;
+	}
+
 	if(hypercall_enabled && !setup_snapshot_once){
 			QEMU_PT_PRINTF(CORE_PREFIX, "Payload Address:\t%lx", hypercall_arg);
 			kvm_arch_get_registers(cpu);	
diff --git a/nyx/state/snapshot_state.c b/nyx/state/snapshot_state.c
index bf45351171..ea0ffe3560 100644
--- a/nyx/state/snapshot_state.c
+++ b/nyx/state/snapshot_state.c
@@ -128,6 +128,10 @@ void deserialize_state(const char* filename_prefix){
 
         assert(apply_capabilities(qemu_get_cpu(0)));
         remap_payload_buffer(nyx_global_state->payload_buffer, ((CPUState *)qemu_get_cpu(0)) );
+
+        /* makes sure that we are allowed to enter the fuzzing loop */
+        nyx_global_state->get_host_config_done = true;
+        nyx_global_state->set_agent_config_done = true;
     }
     else{
         fprintf(stderr, "[QEMU-Nyx]: this feature is currently missing\n");
diff --git a/nyx/state/state.c b/nyx/state/state.c
index 9a2c827f01..9dde7c97f1 100644
--- a/nyx/state/state.c
+++ b/nyx/state/state.c
@@ -107,6 +107,9 @@ void state_init_global(void){
     global_state.pt_trace_mode_force = false;
 
     global_state.num_dirty_pages = 0;
+    
+    global_state.get_host_config_done = false;
+    global_state.set_agent_config_done = false;
 
     global_state.sharedir = sharedir_new();
 
diff --git a/nyx/state/state.h b/nyx/state/state.h
index 3ed5645a9f..7a9367f481 100644
--- a/nyx/state/state.h
+++ b/nyx/state/state.h
@@ -139,6 +139,9 @@ typedef struct qemu_nyx_state_s{
 
     uint32_t num_dirty_pages;
 
+    bool get_host_config_done;
+    bool set_agent_config_done;
+
     /* capabilites */
     uint8_t cap_timeout_detection;
     uint8_t cap_only_reload_mode;