<commit_before>#include "unicorn.h"
#include "log.h"

#include <cstring>

#include <map>
#include <vector>
#include <unordered_set>

#define PAGE_SIZE 0x1000
#define PAGE_SHIFT 12

#define MAX_REG_SIZE 0x2000 // hope it's big enough

extern "C" void x86_reg_update(uc_engine *uc, uint8_t *buf, int save);
extern "C" void mips_reg_update(uc_engine *uc, uint8_t *buf, int save);

typedef enum taint: uint8_t {
	TAINT_NONE = 0,
	TAINT_DIRTY = 1,
	TAINT_SYMBOLIC = 2,
} taint_t;

typedef enum stop {
	STOP_NORMAL=0,
	STOP_STOPPOINT,
	STOP_SYMBOLIC,
	STOP_ERROR,
	STOP_SYSCALL,
	STOP_EXECNONE,
	STOP_ZEROPAGE,
	STOP_NOSTART,
	STOP_SEGFAULT,
} stop_t;

typedef taint_t PageBitmap[PAGE_SIZE];
typedef std::map<uint64_t, std::pair<char *, uint64_t>> PageCache;
std::map<uint64_t, PageCache*> global_cache;


typedef struct mem_access {
	uint64_t address;
	uint8_t value[8]; // assume size of any memory write is no more than 8
	int size;
	int clean; // save current page bitmap
} mem_access_t; // actually it should be `mem_write_t` :)

typedef struct mem_update {
	uint64_t address, length;
	struct mem_update *next;
} mem_update_t;

static void hook_mem_read(uc_engine *uc, uc_mem_type type, uint64_t address, int size, int64_t value, void *user_data);
static void hook_mem_write(uc_engine *uc, uc_mem_type type, uint64_t address, int size, int64_t value, void *user_data);
static bool hook_mem_unmapped(uc_engine *uc, uc_mem_type type, uint64_t address, int size, int64_t value, void *user_data);
static bool hook_mem_prot(uc_engine *uc, uc_mem_type type, uint64_t address, int size, int64_t value, void *user_data);
static void hook_block(uc_engine *uc, uint64_t address, int32_t size, void *user_data);

class State {
private:
	uc_engine *uc;
	PageCache *page_cache;
	bool hooked;

	void (*uc_reg_update)(uc_engine *uc, void *buf, int save);
	uint8_t tmp_reg[MAX_REG_SIZE];

	std::vector<mem_access_t> mem_writes;
	std::map<uint64_t, taint_t *> active_pages;
	std::unordered_set<uint64_t> stop_points;

public:
	std::vector<uint64_t> bbl_addrs;
	uint64_t cur_steps, max_steps;
	uc_hook h_read, h_write, h_block, h_prot, h_unmap;
	stop_t stop_reason;
	bool ignore_next_block;
	bool ignore_next_selfmod;
	uint64_t cur_address;
	int32_t cur_size;

	State(uc_engine *_uc, uint64_t cache_key):uc(_uc)
	{
		hooked = false;
		h_read = h_write = h_block = h_prot = 0;
		max_steps = cur_steps = 0;
		stop_reason = STOP_NOSTART;
		ignore_next_block = false;
		ignore_next_selfmod = false;

		auto it = global_cache.find(cache_key);
		if (it == global_cache.end()) {
				page_cache = new PageCache();
				global_cache[cache_key] = page_cache;
		} else {
				page_cache = it->second;
		}

		switch (*(uc_arch *)uc) { // tricky:)
				case UC_ARCH_X86:
						*(void **)&uc_reg_update = (void *)x86_reg_update;
						break;
				case UC_ARCH_MIPS:
						*(void **)&uc_reg_update = (void *)mips_reg_update;
						break;
				default:
						*(void **)&uc_reg_update = (void *)0xdeadbeef;
						LOG_E("unsupported arch");
		};
	}
	
	/*
	 * HOOK_MEM_WRITE is called before checking if the address is valid. so we might
	 * see uninitialized pages. Using HOOK_MEM_PROT is too late for tracking taint.
	 * so we don't have to use HOOK_MEM_PROT to track dirty pages.
	 *
	 * syscall is not hooked here, because we always stop in  any syscall, and we
	 * prefer to deal with different archs in python.
	 */
	void hook() {
		if (hooked) {
			LOG_D("already hooked");
			return ;
		}
		uc_err err;
		err = uc_hook_add(uc, &h_read, UC_HOOK_MEM_READ, (void *)hook_mem_read, this, 1, 0);

		err = uc_hook_add(uc, &h_write, UC_HOOK_MEM_WRITE, (void *)hook_mem_write, this, 1, 0);

		err = uc_hook_add(uc, &h_block, UC_HOOK_BLOCK, (void *)hook_block, this, 1, 0);

		err = uc_hook_add(uc, &h_prot, UC_HOOK_MEM_PROT, (void *)hook_mem_prot, this, 1, 0);

		err = uc_hook_add(uc, &h_unmap, UC_HOOK_MEM_UNMAPPED, (void *)hook_mem_unmapped, this, 1, 0);

		hooked = true;
	}

	void unhook() {
		if (!hooked)
			return ;

		uc_err err;
		err = uc_hook_del(uc, h_read);
		err = uc_hook_del(uc, h_write);
		err = uc_hook_del(uc, h_block);
		err = uc_hook_del(uc, h_prot);
		err = uc_hook_del(uc, h_unmap);

		hooked = false;
		h_read = h_write = h_block = h_prot = h_unmap = 0;
	}

	~State() {
		for (auto it = active_pages.begin(); it != active_pages.end(); it++) {
			// only poor guys consider about memory leak :(
			LOG_D("delete active page %#lx", it->first);
			// delete should use the bracket operator since PageBitmap is an array typedef
			delete[] it->second;
		}
		active_pages.clear();
	}

	uc_err start(uint64_t pc, uint64_t step = 1) {
		stop_reason = STOP_NOSTART;
		max_steps = step;
		cur_steps = -1;
		return uc_emu_start(uc, pc, 0, 0, 0);
	}

	void stop(stop_t reason) {
		const char *msg = NULL;
		switch (reason) {
			case STOP_NORMAL:
				msg = "reached maximum steps";
				break;
			case STOP_STOPPOINT:
				msg = "hit a stop point";
				break;
			case STOP_SYMBOLIC:
				msg = "read symbolic data";
				break;
			case STOP_ERROR:
				msg = "something wrong";
				break;
			case STOP_SYSCALL:
				msg = "unable to handle syscall";
				commit();
				uc_reg_update(uc, tmp_reg, 1);
				break;
			case STOP_ZEROPAGE:
				msg = "accessing zero page";
				break;
			case STOP_EXECNONE:<commit_msg>dont crash by jumping to 0 in unicorn<commit_after>#include "unicorn.h"
#include "log.h"

#include <cstring>

#include <map>
#include <vector>
#include <unordered_set>

#define PAGE_SIZE 0x1000
#define PAGE_SHIFT 12

#define MAX_REG_SIZE 0x2000 // hope it's big enough

extern "C" void x86_reg_update(uc_engine *uc, uint8_t *buf, int save);
extern "C" void mips_reg_update(uc_engine *uc, uint8_t *buf, int save);

typedef enum taint: uint8_t {
	TAINT_NONE = 0,
	TAINT_DIRTY = 1,
	TAINT_SYMBOLIC = 2,
} taint_t;

typedef enum stop {
	STOP_NORMAL=0,
	STOP_STOPPOINT,
	STOP_SYMBOLIC,
	STOP_ERROR,
	STOP_SYSCALL,
	STOP_EXECNONE,
	STOP_ZEROPAGE,
	STOP_NOSTART,
	STOP_SEGFAULT,
} stop_t;

typedef taint_t PageBitmap[PAGE_SIZE];
typedef std::map<uint64_t, std::pair<char *, uint64_t>> PageCache;
std::map<uint64_t, PageCache*> global_cache;


typedef struct mem_access {
	uint64_t address;
	uint8_t value[8]; // assume size of any memory write is no more than 8
	int size;
	int clean; // save current page bitmap
} mem_access_t; // actually it should be `mem_write_t` :)

typedef struct mem_update {
	uint64_t address, length;
	struct mem_update *next;
} mem_update_t;

static void hook_mem_read(uc_engine *uc, uc_mem_type type, uint64_t address, int size, int64_t value, void *user_data);
static void hook_mem_write(uc_engine *uc, uc_mem_type type, uint64_t address, int size, int64_t value, void *user_data);
static bool hook_mem_unmapped(uc_engine *uc, uc_mem_type type, uint64_t address, int size, int64_t value, void *user_data);
static bool hook_mem_prot(uc_engine *uc, uc_mem_type type, uint64_t address, int size, int64_t value, void *user_data);
static void hook_block(uc_engine *uc, uint64_t address, int32_t size, void *user_data);

class State {
private:
	uc_engine *uc;
	PageCache *page_cache;
	bool hooked;

	void (*uc_reg_update)(uc_engine *uc, void *buf, int save);
	uint8_t tmp_reg[MAX_REG_SIZE];

	std::vector<mem_access_t> mem_writes;
	std::map<uint64_t, taint_t *> active_pages;
	std::unordered_set<uint64_t> stop_points;

public:
	std::vector<uint64_t> bbl_addrs;
	uint64_t cur_steps, max_steps;
	uc_hook h_read, h_write, h_block, h_prot, h_unmap;
	stop_t stop_reason;
	bool ignore_next_block;
	bool ignore_next_selfmod;
	uint64_t cur_address;
	int32_t cur_size;

	State(uc_engine *_uc, uint64_t cache_key):uc(_uc)
	{
		hooked = false;
		h_read = h_write = h_block = h_prot = 0;
		max_steps = cur_steps = 0;
		stop_reason = STOP_NOSTART;
		ignore_next_block = false;
		ignore_next_selfmod = false;

		auto it = global_cache.find(cache_key);
		if (it == global_cache.end()) {
				page_cache = new PageCache();
				global_cache[cache_key] = page_cache;
		} else {
				page_cache = it->second;
		}

		switch (*(uc_arch *)uc) { // tricky:)
				case UC_ARCH_X86:
						*(void **)&uc_reg_update = (void *)x86_reg_update;
						break;
				case UC_ARCH_MIPS:
						*(void **)&uc_reg_update = (void *)mips_reg_update;
						break;
				default:
						*(void **)&uc_reg_update = (void *)0xdeadbeef;
						LOG_E("unsupported arch");
		};
	}
	
	/*
	 * HOOK_MEM_WRITE is called before checking if the address is valid. so we might
	 * see uninitialized pages. Using HOOK_MEM_PROT is too late for tracking taint.
	 * so we don't have to use HOOK_MEM_PROT to track dirty pages.
	 *
	 * syscall is not hooked here, because we always stop in  any syscall, and we
	 * prefer to deal with different archs in python.
	 */
	void hook() {
		if (hooked) {
			LOG_D("already hooked");
			return ;
		}
		uc_err err;
		err = uc_hook_add(uc, &h_read, UC_HOOK_MEM_READ, (void *)hook_mem_read, this, 1, 0);

		err = uc_hook_add(uc, &h_write, UC_HOOK_MEM_WRITE, (void *)hook_mem_write, this, 1, 0);

		err = uc_hook_add(uc, &h_block, UC_HOOK_BLOCK, (void *)hook_block, this, 1, 0);

		err = uc_hook_add(uc, &h_prot, UC_HOOK_MEM_PROT, (void *)hook_mem_prot, this, 1, 0);

		err = uc_hook_add(uc, &h_unmap, UC_HOOK_MEM_UNMAPPED, (void *)hook_mem_unmapped, this, 1, 0);

		hooked = true;
	}

	void unhook() {
		if (!hooked)
			return ;

		uc_err err;
		err = uc_hook_del(uc, h_read);
		err = uc_hook_del(uc, h_write);
		err = uc_hook_del(uc, h_block);
		err = uc_hook_del(uc, h_prot);
		err = uc_hook_del(uc, h_unmap);

		hooked = false;
		h_read = h_write = h_block = h_prot = h_unmap = 0;
	}

	~State() {
		for (auto it = active_pages.begin(); it != active_pages.end(); it++) {
			// only poor guys consider about memory leak :(
			LOG_D("delete active page %#lx", it->first);
			// delete should use the bracket operator since PageBitmap is an array typedef
			delete[] it->second;
		}
		active_pages.clear();
	}

	uc_err start(uint64_t pc, uint64_t step = 1) {
		stop_reason = STOP_NOSTART;
		max_steps = step;
		cur_steps = -1;

		// error if pc is 0
		if (pc == 0) {
		  stop_reason == STOP_ZEROPAGE;
		  return UC_ERR_MAP;
		}

		return uc_emu_start(uc, pc, 0, 0, 0);
	}

	void stop(stop_t reason) {
		const char *msg = NULL;
		switch (reason) {
			case STOP_NORMAL:
				msg = "reached maximum steps";
				break;
			case STOP_STOPPOINT:
				msg = "hit a stop point";
				break;
			case STOP_SYMBOLIC:
				msg = "read symbolic data";
				break;
			case STOP_ERROR:
				msg = "something wrong";
				break;
			case STOP_SYSCALL:
				msg = "unable to handle syscall";
				commit();
				uc_reg_update(uc, tmp_reg, 1);
				break;
			case STOP_ZEROPAGE:
				msg = "accessing zero page";
				break;
			case STOP_EXECNONE:
{'fraction_non_alphanumeric': 0.11317219411654487, 'fraction_numerical': 0.018456061457747797, 'mean_word_length': 4.174503150751333, 'pattern_counts': {'":': 0, '<': 25, '<?xml version=': 0, '>': 31, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2234610', 'n_tokens_mistral': 4735, 'n_tokens_neox': 4126, 'n_words': 1110}