refactor 'redqueen trace' to separate redqueen_trace.c

nyx/auxiliary_buffer.c

@@ -103,7 +103,7 @@ void check_auxiliary_config_buffer(auxilary_buffer_t* auxilary_buffer, auxilary_
 #ifdef SUPPORT_COMPILE_TIME_REDQUEEN
         GET_GLOBAL_STATE()->pt_trace_mode_force = true;		  
 #endif
-        redqueen_set_trace_mode(GET_GLOBAL_STATE()->redqueen_state);
+        redqueen_set_trace_mode();
       }
     }
     else {
@@ -112,7 +112,7 @@ void check_auxiliary_config_buffer(auxilary_buffer_t* auxilary_buffer, auxilary_
 #ifdef SUPPORT_COMPILE_TIME_REDQUEEN
         GET_GLOBAL_STATE()->pt_trace_mode_force = false;
 #endif
-		    redqueen_unset_trace_mode(GET_GLOBAL_STATE()->redqueen_state);
+        redqueen_unset_trace_mode();
       }
     }
 

nyx/file_helper.c

@@ -56,7 +56,6 @@ void parse_address_file(char* path, size_t* num_addrs, uint64_t** addrs){
 
 int re_fd = 0;
 int se_fd = 0;
-int trace_fd = 0;
 
 void write_re_result(char* buf){
   int unused __attribute__((unused));
@@ -65,20 +64,7 @@ void write_re_result(char* buf){
 	unused = write(re_fd, buf, strlen(buf));
 }
 
-void write_trace_result(redqueen_trace_t* trace_state){
+void fsync_redqueen_files(void){
-	//int fd;
-  int unused __attribute__((unused));
-	if (!trace_fd)
-		trace_fd = open(redqueen_workdir.pt_trace_results, O_WRONLY | O_CREAT | O_APPEND, S_IRWXU);
-  redqueen_trace_write_file(trace_state, trace_fd);
-  //unused = write(trace_fd, buf, strlen(buf));
-	//close(fd);
-}
-
-void fsync_all_traces(void){
-  if (!trace_fd){
-    fsync(trace_fd);
-  }
   if (!se_fd){
     fsync(se_fd);
   }
@@ -96,13 +82,6 @@ void write_se_result(char* buf){
 	//close(fd);
 }
 
-void delete_trace_files(void){
-  int unused __attribute__((unused));
-	if (!trace_fd)
-		trace_fd = open(redqueen_workdir.pt_trace_results, O_WRONLY | O_CREAT | O_APPEND, S_IRWXU);
-	unused = ftruncate(trace_fd, 0);
-}
-
 void delete_redqueen_files(void){
   int unused __attribute__((unused));
 	if (!re_fd)

nyx/file_helper.h

@@ -1,7 +1,8 @@
 #include <stdio.h>
 #include <stdint.h>
 #include <stdlib.h>
-#include "redqueen_trace.h"
+
+#pragma once
 
 //doesn't take ownership of path, num_addrs or addrs
 void parse_address_file(char* path, size_t* num_addrs, uint64_t** addrs);
@@ -12,14 +13,9 @@ void write_re_result(char* buf);
 //doesn't take ownership of buf
 void write_se_result(char* buf);
 
-//doesn't take ownership of buf
-void write_trace_result(redqueen_trace_t* trace_state);
-
 //doesn' take ownership of buf
 void write_debug_result(char* buf);
 
 void delete_redqueen_files(void);
 
-void delete_trace_files(void);
+void fsync_redqueen_files(void);
-
-void fsync_all_traces(void);

nyx/hypercall/hypercall.c

@@ -714,19 +714,6 @@ void pt_set_disable_patches_pending(CPUState *cpu){
 	GET_GLOBAL_STATE()->patches_disable_pending = true;
 }
 
-void pt_enable_rqi_trace(CPUState *cpu){
-	if (GET_GLOBAL_STATE()->redqueen_state){
-		redqueen_set_trace_mode(GET_GLOBAL_STATE()->redqueen_state);
-	}
-}
-
-void pt_disable_rqi_trace(CPUState *cpu){
-	if (GET_GLOBAL_STATE()->redqueen_state){
-		redqueen_unset_trace_mode(GET_GLOBAL_STATE()->redqueen_state);
-		return;
-	}
-}
-
 static void handle_hypercall_kafl_dump_file(struct kvm_run *run, CPUState *cpu, uint64_t hypercall_arg)
 {
 	kafl_dump_file_t file_obj;

nyx/hypercall/hypercall.h

@@ -130,8 +130,6 @@ void pt_enable_rqo(CPUState *cpu);
 void pt_disable_rqo(CPUState *cpu);
 void pt_enable_rqi(CPUState *cpu);
 void pt_disable_rqi(CPUState *cpu);
-void pt_enable_rqi_trace(CPUState *cpu);
-void pt_disable_rqi_trace(CPUState *cpu);
 void pt_set_redqueen_instrumentation_mode(CPUState *cpu, int redqueen_instruction_mode);
 void pt_set_redqueen_update_blacklist(CPUState *cpu, bool newval);
 void pt_set_enable_patches_pending(CPUState *cpu);

nyx/pt.c

@@ -163,7 +163,7 @@ void alt_bitmap_add(uint64_t from, uint64_t to)
 {
 	uint64_t transition_value;
 
-	if (GET_GLOBAL_STATE()->redqueen_state->trace_mode) {
+	if (GET_GLOBAL_STATE()->trace_mode) {
 		if(alt_bitmap) {
 			transition_value = mix_bits(to)^(mix_bits(from)>>1);
 			alt_bitmap[transition_value & (alt_bitmap_size-1)]++;
@@ -233,8 +233,8 @@ int pt_enable(CPUState *cpu, bool hmp_mode){
 	if(!fast_reload_set_bitmap(get_fast_reload_snapshot())){
 		coverage_bitmap_reset();
 	}
-	if (GET_GLOBAL_STATE()->redqueen_state->trace_mode) {
+	if (GET_GLOBAL_STATE()->trace_mode) {
-		delete_trace_files();
+		redqueen_trace_reset();
 		alt_bitmap_reset();
 	}
 	pt_truncate_pt_dump_file();

nyx/redqueen.c

@@ -49,7 +49,6 @@ redqueen_t* new_rq_state(CPUState *cpu,  page_cache_t* page_cache){
 
 	res->cpu = cpu;
 	res->intercept_mode = false;
-	res->trace_mode = false;
 	res->singlestep_enabled = false;
   	res->hooks_applied = 0;
   	res->page_cache = page_cache;
@@ -225,43 +224,6 @@ void redqueen_callback(void* opaque, disassembler_mode_t mode, uint64_t start_ad
   }
 }
 
-
-
-static void redqueen_trace_enabled(redqueen_t* self){
-	int unused __attribute__((unused));
-	if(self->trace_mode){
-
-    //libxdc_enable_tracing(GET_GLOBAL_STATE()->decoder);
-    libxdc_enable_tracing(GET_GLOBAL_STATE()->decoder);
-	  libxdc_register_edge_callback(GET_GLOBAL_STATE()->decoder, (void (*)(void*, disassembler_mode_t, uint64_t, uint64_t))&redqueen_trace_register_transition, self->trace_state);
-		//redqueen_trace_register_transition(self->trace_state, INIT_TRACE_IP, ip);
-    //last_ip = ip;
-	} 
-}
-
-static void redqueen_trace_disabled(redqueen_t* self){
-  int unused __attribute__((unused));
-	if(self->trace_mode){
-    libxdc_disable_tracing(GET_GLOBAL_STATE()->decoder);
-
-		//redqueen_trace_register_transition(self->trace_state, last_ip, ip);
-		//edqueen_trace_register_transition(self->trace_state, ip, INIT_TRACE_IP);
-	}
-}
-
-void redqueen_set_trace_mode(redqueen_t* self){
-	self->trace_mode = true;
-  redqueen_trace_enabled(self);
-}
-
-void redqueen_unset_trace_mode(redqueen_t* self){
-	//write_trace_result(self->trace_state);
-	//redqueen_trace_reset(self->trace_state);
-  redqueen_trace_disabled(self);
-
-	self->trace_mode = false;
-}
-
 void destroy_rq_state(redqueen_t* self){
 	redqueen_trace_free(self->trace_state);
 	kh_destroy(RQ, self->lookup);

nyx/redqueen.h

@@ -70,7 +70,6 @@ KHASH_MAP_INIT_INT64(RQ, uint32_t)
 typedef struct redqueen_s{
 	khash_t(RQ) *lookup;
 	bool intercept_mode;
-	bool trace_mode;
 	bool singlestep_enabled;
 	int hooks_applied;
 	CPUState *cpu;
@@ -109,10 +108,6 @@ void enable_rq_intercept_mode(redqueen_t* self);
 void disable_rq_intercept_mode(redqueen_t* self);
 
 
-void redqueen_register_transition(redqueen_t* self, uint64_t ip, uint64_t transition_val);
-void redqueen_set_trace_mode(redqueen_t* self);
-void redqueen_unset_trace_mode(redqueen_t* self);
-
 void set_se_instruction(redqueen_t* self, uint64_t addr);
 
 void dump_se_registers(redqueen_t* self);

nyx/redqueen_trace.c

@@ -2,17 +2,29 @@
 #include <unistd.h>
 #include <stdio.h>
 #include <assert.h>
+
 #include "redqueen_trace.h"
+#include "redqueen.h"
+#include "state/state.h"
+
 
 void alt_bitmap_add(uint64_t from, uint64_t to);
 
 /* write full trace of edge transitions rather than sorted list? */
 //#define KAFL_FULL_TRACES
 
-#ifdef KAFL_FULL_TRACES
+int trace_fd = 0;
-#include "redqueen.h"
+
-extern int trace_fd;
+static int reset_trace_fd(void) {
-#endif
+	if (trace_fd)
+		close(trace_fd);
+	trace_fd = open(redqueen_workdir.pt_trace_results, O_WRONLY | O_CREAT | O_TRUNC, 0644);
+	if (trace_fd < 0) {
+		fprintf(stderr, "Failed to initiate trace output: %s\n", strerror(errno));
+		assert(0);
+	}
+	return trace_fd;
+}
 
 redqueen_trace_t* redqueen_trace_new(void){
 	redqueen_trace_t* self = malloc(sizeof(redqueen_trace_t));
@@ -23,7 +35,8 @@ redqueen_trace_t* redqueen_trace_new(void){
 	return self;
 }
 
-void redqueen_trace_reset(redqueen_trace_t* self){
+static void redqueen_state_reset(void){
+	redqueen_trace_t *self = GET_GLOBAL_STATE()->redqueen_state->trace_state;
 	kh_destroy(RQ_TRACE, self->lookup);
 	self->lookup = kh_init(RQ_TRACE);
 	self->num_ordered_transitions = 0;
@@ -43,9 +56,7 @@ void redqueen_trace_register_transition(redqueen_trace_t* self, disassembler_mod
 	if (from != exit_ip && to != exit_ip)
 		alt_bitmap_add(from, to);
 #ifdef KAFL_FULL_TRACES
-	extern int trace_fd;
+	assert(trace_fd >= 0);
-	if (!trace_fd)
-		trace_fd = open(redqueen_workdir.pt_trace_results, O_WRONLY | O_CREAT | O_APPEND, S_IRWXU);
 	dprintf(trace_fd, "%lx,%lx\n", from, to);
 	return;
 #endif
@@ -62,33 +73,60 @@ void redqueen_trace_register_transition(redqueen_trace_t* self, disassembler_mod
 	}
 }	
 
-void redqueen_trace_write_file(redqueen_trace_t* self, int fd){
+static void redqueen_trace_write(void){
 #ifdef KAFL_FULL_TRACES
 	return;
 #endif
+	redqueen_trace_t *self = GET_GLOBAL_STATE()->redqueen_state->trace_state;
+	assert(trace_fd >= 0);
 	for(size_t i = 0; i < self->num_ordered_transitions; i++){
 		khiter_t k;
 		uint128_t key = self->ordered_transitions[i];
 		k = kh_get(RQ_TRACE, self->lookup, key); 
 		assert(k != kh_end(self->lookup));
-		dprintf(fd, "%lx,%lx,%lx\n",  (uint64_t)(key>>64), (uint64_t)key, kh_value(self->lookup, k) );
+		dprintf(trace_fd, "%lx,%lx,%lx\n",  (uint64_t)(key>>64), (uint64_t)key, kh_value(self->lookup, k) );
 	}
 }
 
+void redqueen_trace_reset(void){
+	redqueen_state_reset();
+	reset_trace_fd();
+}
+
+void redqueen_trace_flush(void){
+	redqueen_trace_write();
+	if (trace_fd)
+		fsync(trace_fd);
+}
+
+void redqueen_set_trace_mode(void){
+	GET_GLOBAL_STATE()->trace_mode = true;
+	libxdc_enable_tracing(GET_GLOBAL_STATE()->decoder);
+	libxdc_register_edge_callback(GET_GLOBAL_STATE()->decoder,
+			(void (*)(void*, disassembler_mode_t, uint64_t, uint64_t))&redqueen_trace_register_transition,
+			GET_GLOBAL_STATE()->redqueen_state->trace_state);
+}
+
+void redqueen_unset_trace_mode(void){
+    libxdc_disable_tracing(GET_GLOBAL_STATE()->decoder);
+	GET_GLOBAL_STATE()->trace_mode = false;
+}
 
 #ifdef DEBUG_MAIN
 int main(int argc, char** argv){
 
 	redqueen_trace_t* rq_obj = redqueen_trace_new();
 
+	reset_trace_fd();
+
 	for (uint64_t j = 0; j < 0x5; j++){
 		redqueen_trace_register_transition(rq_obj, 0xBADF, 0xC0FFEE);
 		redqueen_trace_register_transition(rq_obj, 0xBADBEEF, 0xC0FFEE);
 		for (uint64_t i = 0; i < 0x10000; i++){
 			redqueen_trace_register_transition(rq_obj, 0xBADBEEF, 0xC0FFEE);
 		}
-		redqueen_trace_write_file(rq_obj, STDOUT_FILENO);
+		redqueen_trace_write(rq_obj, STDOUT_FILENO);
-		redqueen_trace_reset(rq_obj);
+		redqueen_state_reset();
 	}
 
 	redqueen_trace_free(rq_obj);

nyx/redqueen_trace.h

@@ -37,7 +37,11 @@ typedef struct redqueen_trace_s{
 } redqueen_trace_t;
 
 redqueen_trace_t* redqueen_trace_new(void);
-void redqueen_trace_reset(redqueen_trace_t* self);
 void redqueen_trace_free(redqueen_trace_t* self);
 void redqueen_trace_register_transition(redqueen_trace_t* self, disassembler_mode_t mode, uint64_t from, uint64_t to);
-void redqueen_trace_write_file(redqueen_trace_t* self, int fd);
+
+void redqueen_set_trace_mode(void);
+void redqueen_unset_trace_mode(void);
+
+void redqueen_trace_flush(void);
+void redqueen_trace_reset(void);

nyx/state/state.c

@@ -91,6 +91,7 @@ void state_init_global(void){
     global_state.in_fuzzing_mode = false;
     global_state.in_reload_mode = true;
     global_state.starved = false;
+    global_state.trace_mode = false;
     global_state.shutdown_requested = false;
     global_state.cow_cache_full = false;
 

nyx/state/state.h

@@ -133,6 +133,7 @@ typedef struct qemu_nyx_state_s{
     bool in_fuzzing_mode;
     bool in_reload_mode; 
     bool starved;
+    bool trace_mode;
 
     bool shutdown_requested;
     bool cow_cache_full;

nyx/synchronization.c

@@ -277,12 +277,12 @@ void synchronization_lock(void){
 
 	//last_timeout = false;
 
-	if(unlikely(GET_GLOBAL_STATE()->in_redqueen_reload_mode || GET_GLOBAL_STATE()->redqueen_state->trace_mode)){
+	if(unlikely(GET_GLOBAL_STATE()->in_redqueen_reload_mode)) {
-		if(GET_GLOBAL_STATE()->redqueen_state->trace_mode){
+			fsync_redqueen_files();		
-			write_trace_result(GET_GLOBAL_STATE()->redqueen_state->trace_state);
-			redqueen_trace_reset(GET_GLOBAL_STATE()->redqueen_state->trace_state);
 	}
-		fsync_all_traces();		
+
+	if (unlikely(GET_GLOBAL_STATE()->trace_mode)) {
+		redqueen_trace_flush();
 	}
 
 	interface_send_char(NYX_INTERFACE_PING);