Quick & dirty thread switching support

src/analyzer/main.cpp

@@ -48,7 +48,17 @@ namespace
 			}
 			else
 			{
-				win_emu.emu().start_from_ip();
+				while (true)
+				{
+					win_emu.emu().start_from_ip();
+					if (win_emu.switch_thread)
+					{
+						win_emu.perform_thread_switch();
+						continue;
+					}
+
+					break;
+				}
 			}
 		}
 		catch (...)

src/windows-emulator/process_context.hpp

@@ -201,6 +201,8 @@ public:
 
 	uint32_t id{};
 
+	std::optional<handle> await_object{};
+
 	std::optional<emulator_allocator> gs_segment;
 	std::optional<emulator_object<TEB>> teb;
 

src/windows-emulator/syscalls.cpp

@@ -1967,6 +1967,10 @@ namespace
 				{
 					write_attribute(c.emu, attribute, thread->teb->ptr());
 				}
+				else
+				{
+					printf("Unsupported thread attribute type: %llX\n", type);
+				}
 			}, i);
 		}
 
@@ -1977,6 +1981,38 @@ namespace
 	{
 		return FALSE;
 	}
+
+	NTSTATUS handle_NtWaitForSingleObject(const syscall_context& c, const uint64_t handle_value,
+	                                      const BOOLEAN alertable,
+	                                      const emulator_object<LARGE_INTEGER> timeout)
+	{
+		if (timeout.value())
+		{
+			puts("NtWaitForSingleObject timeout not supported yet!");
+			return STATUS_NOT_SUPPORTED;
+		}
+
+		if (alertable)
+		{
+			puts("Alertable NtWaitForSingleObject not supported yet!");
+			return STATUS_NOT_SUPPORTED;
+		}
+
+		handle h{};
+		h.bits = handle_value;
+
+		if (h.value.type != handle_types::thread)
+		{
+			puts("NtWaitForSingleObject only supported with thread handles yet!");
+			return STATUS_NOT_SUPPORTED;
+		}
+
+		c.win_emu.current_thread().await_object = h;
+		c.win_emu.switch_thread = true;
+		c.emu.stop();
+
+		return STATUS_WAIT_0;
+	}
 }
 
 void syscall_dispatcher::setup(const exported_symbols& ntdll_exports, const exported_symbols& win32u_exports)
@@ -2077,6 +2113,7 @@ void syscall_dispatcher::add_handlers()
 	add_handler(NtQueryInformationFile);
 	add_handler(NtCreateThreadEx);
 	add_handler(NtQueryDebugFilterState);
+	add_handler(NtWaitForSingleObject);
 
 #undef add_handler
 

src/windows-emulator/windows_emulator.cpp

@@ -4,6 +4,8 @@
 
 #include <unicorn_x64_emulator.hpp>
 
+constexpr auto MAX_INSTRUCTIONS_PER_TIME_SLICE = 100;
+
 namespace
 {
 	template <typename T>
@@ -438,6 +440,28 @@ namespace
 
 		switch_to_thread(emu, context, *thread);
 	}
+
+	void switch_to_next_thread(x64_emulator& emu, process_context& context)
+	{
+		bool next_thread = false;
+
+
+		for (auto& thread : context.threads)
+		{
+			if (next_thread)
+			{
+				switch_to_thread(emu, context, thread.second);
+				return;
+			}
+
+			if (&thread.second == context.active_thread)
+			{
+				next_thread = true;
+			}
+		}
+
+		switch_to_thread(emu, context, context.threads.begin()->second);
+	}
 }
 
 emulator_thread::emulator_thread(x64_emulator& emu, const process_context& context,
@@ -448,8 +472,8 @@ emulator_thread::emulator_thread(x64_emulator& emu, const process_context& conte
 	  , stack_size(page_align_up(std::max(stack_size, STACK_SIZE)))
 	  , start_address(start_address)
 	  , argument(argument)
-	  , last_registers(context.default_register_set)
 	  , id(id)
+	  , last_registers(context.default_register_set)
 {
 	this->stack_base = emu.allocate_memory(this->stack_size, memory_permission::read_write);
 
@@ -548,6 +572,13 @@ void windows_emulator::setup_process(const std::filesystem::path& application,
 	switch_to_thread(emu, context, main_thread_id);
 }
 
+void windows_emulator::perform_thread_switch()
+{
+	this->logger.print(color::green, "Performing thread switch...\n");
+	switch_to_next_thread(this->emu(), this->process());
+	this->switch_thread = false;
+}
+
 void windows_emulator::setup_hooks()
 {
 	this->emu().hook_instruction(x64_hookable_instructions::syscall, [&]
@@ -619,6 +650,16 @@ void windows_emulator::setup_hooks()
 
 		                                  ++process.executed_instructions;
 
+		                                  auto& thread = this->current_thread();
+		                                  if (thread.executed_instructions == MAX_INSTRUCTIONS_PER_TIME_SLICE)
+		                                  {
+			                                  this->switch_thread = true;
+			                                  this->emu().stop();
+		                                  }
+
+		                                  ++thread.executed_instructions;
+		                                  thread.executed_instructions %= MAX_INSTRUCTIONS_PER_TIME_SLICE;
+
 		                                  process.previous_ip = process.current_ip;
 		                                  process.current_ip = this->emu().read_instruction_pointer();
 

src/windows-emulator/windows_emulator.hpp

@@ -77,6 +77,9 @@ public:
 	bool verbose_calls{false};
 	bool buffer_stdout{false};
 	bool fuzzing{false};
+	bool switch_thread{false};
+
+	void perform_thread_switch();
 
 private:
 	std::unique_ptr<x64_emulator> emu_{};