Fix threading to fully work

src/windows-emulator/process_context.hpp

@@ -219,14 +219,26 @@ public:
 
 	std::optional<uint32_t> exit_status{};
 	std::optional<handle> await_object{};
-	std::optional<bool> alerted{};
+	bool waiting_for_alert{false};
+	bool alerted{false};
 	std::optional<std::chrono::steady_clock::time_point> await_time{};
 
+	std::optional<NTSTATUS> pending_status{};
+
 	std::optional<emulator_allocator> gs_segment;
 	std::optional<emulator_object<TEB>> teb;
 
 	std::vector<std::byte> last_registers{};
 
+	void mark_as_ready(NTSTATUS status);
+
+	bool is_await_time_over() const
+	{
+		return this->await_time.has_value() && this->await_time.value() < std::chrono::steady_clock::now();
+	}
+
+	bool is_thread_ready(process_context& context);
+
 	void save(x64_emulator& emu)
 	{
 		this->last_registers = emu.save_registers();
@@ -237,12 +249,20 @@ public:
 		emu.restore_registers(this->last_registers);
 	}
 
-	void setup_if_necessary(x64_emulator& emu, const process_context& context) const
+	void setup_if_necessary(x64_emulator& emu, const process_context& context)
 	{
 		if (!this->executed_instructions)
 		{
 			this->setup_registers(emu, context);
 		}
+
+		if (this->pending_status.has_value())
+		{
+			const auto status = *this->pending_status;
+			this->pending_status = {};
+
+			emu.reg<uint64_t>(x64_register::rax, static_cast<uint64_t>(status));
+		}
 	}
 
 	void serialize(utils::buffer_serializer&) const

src/windows-emulator/syscalls.cpp

@@ -176,38 +176,6 @@ namespace
 		}
 	}
 
-	bool is_object_signaled(process_context& c, const handle h)
-	{
-		const auto type = h.value.type;
-
-		switch (type)
-		{
-		case handle_types::event:
-			{
-				const auto* e = c.events.get(h);
-				if (e)
-				{
-					return e->signaled;
-				}
-
-				break;
-			}
-
-		case handle_types::thread:
-			{
-				const auto* t = c.threads.get(h);
-				if (t)
-				{
-					return t->exit_status.has_value();
-				}
-
-				break;
-			}
-		}
-
-		throw std::runtime_error("Bad object");
-	}
-
 	std::chrono::steady_clock::time_point convert_delay_interval_to_time_point(const LARGE_INTEGER delay_interval)
 	{
 		constexpr auto HUNDRED_NANOSECONDS_IN_ONE_SECOND = 10000000LL;
@@ -2084,8 +2052,6 @@ namespace
 		handle h{};
 		h.bits = handle_value;
 
-		auto& t = c.win_emu.current_thread();
-
 		if (h.value.type != handle_types::thread && h.value.type != handle_types::event)
 		{
 			puts("Unsupported handle type for NtWaitForSingleObject!");
@@ -2093,27 +2059,14 @@ namespace
 			return STATUS_NOT_SUPPORTED;
 		}
 
+		auto& t = c.win_emu.current_thread();
+		t.await_object = h;
+
 		if (timeout.value() && !t.await_time.has_value())
 		{
 			t.await_time = convert_delay_interval_to_time_point(timeout.read());
 		}
 
-		if (is_object_signaled(c.proc, h))
-		{
-			t.await_time = {};
-			t.await_object = {};
-			return STATUS_WAIT_0;
-		}
-
-		if (t.await_time.has_value() && *t.await_time < std::chrono::steady_clock::now())
-		{
-			t.await_time = {};
-			t.await_object = {};
-			return STATUS_TIMEOUT;
-		}
-
-		t.await_object = h;
-		c.retrigger_syscall = true;
 		c.win_emu.switch_thread = true;
 		c.emu.stop();
 
@@ -2153,18 +2106,8 @@ namespace
 		}
 
 		auto& t = c.win_emu.current_thread();
+		t.await_time = convert_delay_interval_to_time_point(delay_interval.read());
 
-		if (!t.await_time.has_value())
-		{
-			t.await_time = convert_delay_interval_to_time_point(delay_interval.read());
-		}
-		else if (*t.await_time < std::chrono::steady_clock::now())
-		{
-			t.await_time = {};
-			return STATUS_SUCCESS;
-		}
-
-		c.retrigger_syscall = true;
 		c.win_emu.switch_thread = true;
 		c.emu.stop();
 
@@ -2197,30 +2140,13 @@ namespace
 	                                         const emulator_object<LARGE_INTEGER> timeout)
 	{
 		auto& t = c.win_emu.current_thread();
+		t.waiting_for_alert = true;
 
 		if (timeout.value() && !t.await_time.has_value())
 		{
 			t.await_time = convert_delay_interval_to_time_point(timeout.read());
 		}
 
-		if (!t.alerted.has_value())
-		{
-			t.alerted = false;
-		}
-		else if (*t.alerted)
-		{
-			t.alerted = {};
-			t.await_time = {};
-			return STATUS_ALERTED;
-		}
-		else if (t.await_time.has_value() && *t.await_time < std::chrono::steady_clock::now())
-		{
-			t.alerted = {};
-			t.await_time = {};
-			return STATUS_TIMEOUT;
-		}
-
-		c.retrigger_syscall = true;
 		c.win_emu.switch_thread = true;
 		c.emu.stop();
 

src/windows-emulator/windows_emulator.cpp

@@ -417,9 +417,23 @@ namespace
 		dispatch_exception_pointers(emu, dispatcher, pointers);
 	}
 
-	void switch_to_thread(x64_emulator& emu, process_context& context, emulator_thread& thread)
+	bool switch_to_thread(const logger& logger, x64_emulator& emu, process_context& context, emulator_thread& thread)
 	{
+		if (thread.exit_status.has_value() || !thread.is_thread_ready(context))
+		{
+			return false;
+		}
+
 		auto* active_thread = context.active_thread;
+
+		if (active_thread == &thread)
+		{
+			return true;
+		}
+
+		logger.print(color::green, "Performing thread switch...\n");
+
+
 		if (active_thread)
 		{
 			active_thread->save(emu);
@@ -430,9 +444,11 @@ namespace
 
 		thread.restore(emu);
 		thread.setup_if_necessary(emu, context);
+
+		return true;
 	}
 
-	void switch_to_thread(x64_emulator& emu, process_context& context, const handle thread_handle)
+	bool switch_to_thread(const logger& logger, x64_emulator& emu, process_context& context, const handle thread_handle)
 	{
 		auto* thread = context.threads.get(thread_handle);
 		if (!thread)
@@ -440,26 +456,23 @@ namespace
 			throw std::runtime_error("Bad thread handle");
 		}
 
-		switch_to_thread(emu, context, *thread);
+		return switch_to_thread(logger, emu, context, *thread);
 	}
 
-	void switch_to_next_thread(x64_emulator& emu, process_context& context)
+	bool switch_to_next_thread(const logger& logger, x64_emulator& emu, process_context& context)
 	{
-		//cleanup_threads(context);
-
 		bool next_thread = false;
 
 		for (auto& thread : context.threads)
 		{
 			if (next_thread)
 			{
-				if (thread.second.exit_status.has_value())
+				if (switch_to_thread(logger, emu, context, thread.second))
 				{
-					continue;
+					return true;
 				}
 
-				switch_to_thread(emu, context, thread.second);
-				return;
+				continue;
 			}
 
 			if (&thread.second == context.active_thread)
@@ -468,7 +481,50 @@ namespace
 			}
 		}
 
-		switch_to_thread(emu, context, context.threads.begin()->second);
+		for (auto& thread : context.threads)
+		{
+			if (switch_to_thread(logger, emu, context, thread.second))
+			{
+				return true;
+			}
+		}
+
+		return false;
+	}
+
+	bool is_object_signaled(process_context& c, const handle h)
+	{
+		const auto type = h.value.type;
+
+		switch (type)
+		{
+		default:
+			break;
+
+		case handle_types::event:
+			{
+				const auto* e = c.events.get(h);
+				if (e)
+				{
+					return e->signaled;
+				}
+
+				break;
+			}
+
+		case handle_types::thread:
+			{
+				const auto* t = c.threads.get(h);
+				if (t)
+				{
+					return t->exit_status.has_value();
+				}
+
+				break;
+			}
+		}
+
+		throw std::runtime_error("Bad object");
 	}
 }
 
@@ -504,6 +560,70 @@ emulator_thread::emulator_thread(x64_emulator& emu, const process_context& conte
 	});
 }
 
+void emulator_thread::mark_as_ready(const NTSTATUS status)
+{
+	this->pending_status = status;
+	this->await_time = {};
+	this->await_object = {};
+
+	// TODO: Find out if this is correct
+	if (this->waiting_for_alert)
+	{
+		this->alerted = false;
+	}
+
+	this->waiting_for_alert = false;
+}
+
+bool emulator_thread::is_thread_ready(process_context& context)
+{
+	if (this->waiting_for_alert)
+	{
+		if (this->alerted)
+		{
+			this->mark_as_ready(STATUS_ALERTED);
+			return true;
+		}
+		if (this->is_await_time_over())
+		{
+			this->mark_as_ready(STATUS_TIMEOUT);
+			return true;
+		}
+
+		return false;
+	}
+
+	if (this->await_object.has_value())
+	{
+		if (is_object_signaled(context, *this->await_object))
+		{
+			this->mark_as_ready(STATUS_WAIT_0);
+			return true;
+		}
+
+		if (this->is_await_time_over())
+		{
+			this->mark_as_ready(STATUS_TIMEOUT);
+			return true;
+		}
+
+		return false;
+	}
+
+	if (this->await_time.has_value())
+	{
+		if (this->is_await_time_over())
+		{
+			this->mark_as_ready(STATUS_SUCCESS);
+			return true;
+		}
+
+		return false;
+	}
+
+	return true;
+}
+
 void emulator_thread::setup_registers(x64_emulator& emu, const process_context& context) const
 {
 	setup_stack(emu, this->stack_base, this->stack_size);
@@ -577,14 +697,17 @@ void windows_emulator::setup_process(const std::filesystem::path& application,
 	context.default_register_set = emu.save_registers();
 
 	const auto main_thread_id = context.create_thread(emu, context.executable->entry_point, 0, 0);
-	switch_to_thread(emu, context, main_thread_id);
+	switch_to_thread(this->logger, 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;
+	while (!switch_to_next_thread(this->logger, this->emu(), this->process()))
+	{
+		// TODO: Optimize that
+		std::this_thread::sleep_for(1ms);
+	}
 }
 
 void windows_emulator::setup_hooks()