CoopAllTheThings/host/src/injection_panel.cpp

#include "injection_panel.hpp"

#include "vk_layer_setup.hpp"

#include <cmath>

#include <windows.h>

#include "inject/dll_probe.hpp"
#include "inject/injector.hpp"
#include "ui/app_chrome.hpp"

namespace coop
{

namespace
{

const ImVec4 kGreen(0.4f, 1.0f, 0.4f, 1.0f);
const ImVec4 kRed(1.0f, 0.45f, 0.4f, 1.0f);
const ImVec4 kGrey(0.7f, 0.7f, 0.7f, 1.0f);

std::string narrow(const std::wstring& w)
{
	if (w.empty())
	{
		return {};
	}
	const int len = WideCharToMultiByte(CP_UTF8, 0, w.c_str(), static_cast<int>(w.size()), nullptr, 0, nullptr, nullptr);
	std::string out(static_cast<std::size_t>(len), '\0');
	WideCharToMultiByte(CP_UTF8, 0, w.c_str(), static_cast<int>(w.size()), out.data(), len, nullptr, nullptr);
	return out;
}

// Case-insensitive equality of two image names (e.g. "snb.exe").
bool iequals_name(const std::wstring& a, const std::wstring& b)
{
	const std::string x = narrow(a), y = narrow(b);
	if (x.size() != y.size())
	{
		return false;
	}
	for (std::size_t i = 0; i < x.size(); ++i)
	{
		if (::tolower(static_cast<unsigned char>(x[i])) != ::tolower(static_cast<unsigned char>(y[i])))
		{
			return false;
		}
	}
	return true;
}

bool contains_ci(const std::wstring& haystack, const char* needle_utf8)
{
	if (needle_utf8 == nullptr || needle_utf8[0] == '\0')
	{
		return true;
	}
	const std::string hay = narrow(haystack);
	std::string h = hay, n = needle_utf8;
	for (char& c : h)
	{
		c = static_cast<char>(::tolower(static_cast<unsigned char>(c)));
	}
	for (char& c : n)
	{
		c = static_cast<char>(::tolower(static_cast<unsigned char>(c)));
	}
	return h.find(n) != std::string::npos;
}

// Absolute path to coop_hook.dll, assumed to sit next to the host executable.
std::wstring hook_dll_path()
{
	wchar_t buffer[MAX_PATH] = {};
	const DWORD len = GetModuleFileNameW(nullptr, buffer, MAX_PATH);
	std::wstring path(buffer, len);
	const std::size_t slash = path.find_last_of(L"\\/");
	if (slash != std::wstring::npos)
	{
		path.resize(slash + 1);
	}
	path += L"coop_hook.dll";
	return path;
}

} // namespace

InjectionPanel::InjectionPanel()
{
	refresh_targets();
}

void InjectionPanel::refresh_targets()
{
	windows_ = list_windows();
	processes_ = list_processes();
}

InjectionPanel::~InjectionPanel()
{
	// Graceful tool exit: unhook everything so a still-running game returns to normal (the DLL
	// stays injected, dormant). Short timeout -- the flags persist in the section the DLL keeps
	// alive, so the unhook completes even if the process exits before it confirms.
	disconnect_graceful(/*timeout_ms=*/300);
	if (vk_layer_enabled_)
	{
		unregister_vk_layer(); // don't leave the implicit layer registered after the tool closes
	}
	close_target_handle();
}

void InjectionPanel::disconnect_graceful(int timeout_ms)
{
	// Tell the DLL to remove every hook so the game runs as if it was never touched, then wait
	// (bounded) for it to confirm before we drop the channel. The flags persist in the section the
	// DLL keeps alive, so it unhooks even if we time out or exit first -- the wait just lets us
	// observe a clean game. The DLL is left injected (dormant) for a later reconnect; we never eject.
	if (server_.running())
	{
		server_.request_unhook_all();
		for (int waited = 0; waited < timeout_ms && !server_.all_hooks_removed(); waited += 10)
		{
			Sleep(10);
		}
	}
	server_.stop();
	injected_ = false;
	close_target_handle();
}

void InjectionPanel::close_target_handle()
{
	if (target_process_ != nullptr)
	{
		CloseHandle(target_process_);
		target_process_ = nullptr;
	}
	target_state_ = TargetState::NotInjected;
	dll_alive_ = false;
	last_heartbeat_ = 0;
	last_heartbeat_time_ = 0.0;
}

void InjectionPanel::tick()
{
	update_liveness();
	auto_reattach_tick();
}

void InjectionPanel::auto_reattach_tick()
{
	if (!auto_reattach_ || target_state_ != TargetState::Terminated || selected_name_.empty())
	{
		return;
	}
	// Poll the process list a couple of times a second (cheap, and we want to catch the
	// relaunch early to read the exact audio format before the game creates its client).
	const double now = ImGui::GetTime();
	if (now - last_auto_poll_ < 0.5)
	{
		return;
	}
	last_auto_poll_ = now;
	refresh_processes();
	for (const ProcessEntry& e : processes_)
	{
		if (iequals_name(e.exe_name, selected_name_))
		{
			// The same game relaunched -> tear down the stale channel and re-attach to it.
			server_.stop();
			close_target_handle();
			injected_ = false;
			selected_pid_ = e.pid;
			inject_selected(); // logs success/failure in status_; retried next poll on failure
			break;
		}
	}
}

void InjectionPanel::update_liveness()
{
	if (!injected_)
	{
		target_state_ = TargetState::NotInjected;
		return;
	}

	// Process gone? The handle was opened with SYNCHRONIZE at inject time, so a
	// signaled wait means it exited. This is authoritative even if the heartbeat
	// happened to look alive a moment ago.
	if (target_process_ != nullptr && WaitForSingleObject(target_process_, 0) == WAIT_OBJECT_0)
	{
		target_state_ = TargetState::Terminated;
		dll_alive_ = false;
		return;
	}

	// Still running: alive vs hung from the hook heartbeat (advances ~4x/s). A live
	// process whose heartbeat stalled for ~2 s is frozen, not gone -- a distinct state.
	const std::uint32_t hb = server_.hook_status().heartbeat;
	const double now = ImGui::GetTime();
	if (hb != last_heartbeat_)
	{
		last_heartbeat_ = hb;
		last_heartbeat_time_ = now;
		dll_alive_ = true;
	}
	else if (now - last_heartbeat_time_ > 2.0)
	{
		dll_alive_ = false;
	}
	target_state_ = dll_alive_ ? TargetState::Alive : TargetState::Hung;
}

void InjectionPanel::refresh_processes()
{
	processes_ = list_processes();
}

void InjectionPanel::publish_subsystem_state()
{
	// Publish the desired subsystem state before the hook's next reconcile tick.
	server_.set_subsystem_enabled(HookSubsys_Input, want_input_);
	server_.set_subsystem_enabled(HookSubsys_Focus, want_focus_);
	server_.set_subsystem_enabled(HookSubsys_Audio, want_audio_);
	server_.set_subsystem_enabled(HookSubsys_Video, want_video_);
	server_.set_subsystem_enabled(HookSubsys_Mkb, want_mkb_);
	server_.set_cursor_clip_allowed(!release_cursor_);
}

void InjectionPanel::begin_liveness_tracking()
{
	injected_ = true;
	// Track liveness: a SYNCHRONIZE|QUERY handle lets us notice the game exiting,
	// and seeding the heartbeat clock avoids a spurious "hung" before the first beat.
	close_target_handle(); // drop any handle from a previous target
	target_process_ = OpenProcess(SYNCHRONIZE | PROCESS_QUERY_LIMITED_INFORMATION, FALSE, selected_pid_);
	target_state_ = TargetState::Alive;
	last_heartbeat_ = 0;
	last_heartbeat_time_ = ImGui::GetTime();
	dll_alive_ = true;
}

void InjectionPanel::reconnect_selected()
{
	// Re-attach to a DLL that's already injected and alive (a prior session left it dormant after a
	// graceful disconnect, or the tool restarted): bring the channel back up on the SAME per-pid
	// section the DLL still holds and re-publish the desired subsystem state -- no re-injection.
	if (!server_.start(selected_pid_))
	{
		status_ = "Failed to re-attach shared memory.";
		status_color_ = kRed;
		return;
	}
	publish_subsystem_state();
	begin_liveness_tracking();
	status_ = "Reconnected to " + narrow(selected_name_) + " (pid " + std::to_string(selected_pid_) +
	          ") -- reused the injected DLL.";
	status_color_ = kGreen;
}

void InjectionPanel::inject_selected()
{
	if (selected_pid_ == 0)
	{
		status_ = "Select a target process first.";
		status_color_ = kRed;
		return;
	}

	// If our DLL is already injected and alive in this target (left dormant by a graceful disconnect,
	// or surviving a tool restart -- it keeps the section alive), reconnect to it instead of
	// injecting a second time.
	if (hook_dll_alive(selected_pid_))
	{
		reconnect_selected();
		return;
	}

	// Bring up the shared-memory channel before injecting so the hook finds it
	// immediately on load.
	if (!server_.start(selected_pid_))
	{
		status_ = "Failed to create shared memory.";
		status_color_ = kRed;
		return;
	}

	publish_subsystem_state();

	const InjectResult result = inject_dll(selected_pid_, hook_dll_path());
	if (result.status == InjectStatus::Ok)
	{
		begin_liveness_tracking();
		status_ = "Injected into " + narrow(selected_name_) + " (pid " + std::to_string(selected_pid_) + ").";
		status_color_ = kGreen;
	}
	else
	{
		server_.stop();
		status_ = std::string("Injection failed: ") + to_string(result.status);
		if (result.os_error != 0)
		{
			status_ += " [err " + std::to_string(result.os_error) + "]";
		}
		status_color_ = kRed;
	}
}

void InjectionPanel::reattach()
{
	if (selected_name_.empty())
	{
		return;
	}

	// Find live processes that share the original target's image name.
	refresh_processes();
	std::vector<unsigned long> matches;
	for (const ProcessEntry& e : processes_)
	{
		if (iequals_name(e.exe_name, selected_name_))
		{
			matches.push_back(e.pid);
		}
	}

	const std::string name = narrow(selected_name_);
	if (matches.empty())
	{
		status_ = "No running \"" + name + "\" to re-attach to.";
		status_color_ = kRed;
		return;
	}
	if (matches.size() > 1)
	{
		// Don't guess which instance: filter the picker to the matches so the operator
		// chooses, then injects via the normal button.
		snprintf(filter_, sizeof(filter_), "%s", name.c_str());
		selected_pid_ = 0;
		status_ = "Multiple \"" + name + "\" running -- pick one below, then Inject & Connect.";
		status_color_ = kGrey;
		return;
	}

	// Exactly one: tear down the stale channel and inject into the new pid.
	server_.stop();
	close_target_handle();
	injected_ = false;
	selected_pid_ = matches.front();
	inject_selected();
}

#ifdef COOP_TEST_HARNESS
unsigned long InjectionPanel::dev_inject_by_name(const std::wstring& image_name)
{
	refresh_processes();
	for (const ProcessEntry& e : processes_)
	{
		if (iequals_name(e.exe_name, image_name))
		{
			selected_pid_ = e.pid;
			selected_name_ = e.exe_name;
			inject_selected();
			return injected_ ? selected_pid_ : 0;
		}
	}
	return 0;
}
#endif

void InjectionPanel::publish(const std::array<PadInfo, kMaxPads>& pads)
{
	if (!test_input_.load(std::memory_order_relaxed))
	{
		server_.publish(pads);
		return;
	}

	// Synthetic, unmistakably non-human pattern: left stick sweeps a circle and
	// A is pressed every other second. If the game moves on its own to this, the
	// forwarding pipeline is proven end-to-end.
	const unsigned long long ms = GetTickCount64();
	const double t = static_cast<double>(ms) / 1000.0;

	std::array<PadInfo, kMaxPads> synthetic{};
	PadInfo& pad = synthetic[0];
	pad.connected = true;
	pad.source = "synthetic test";
	pad.state.connected = 1;
	pad.state.packet = static_cast<std::uint32_t>(ms);
	pad.state.thumb_lx = static_cast<std::int16_t>(std::cos(t) * 30000.0);
	pad.state.thumb_ly = static_cast<std::int16_t>(std::sin(t) * 30000.0);
	if ((ms / 1000) % 2 == 0)
	{
		pad.state.buttons |= 0x1000; // XINPUT_GAMEPAD_A
	}
	server_.publish(synthetic);
}

void InjectionPanel::draw_hook_list(const HookStatusView& status)
{
	static const char* kSubsysName[] = {"Input", "Focus", "Audio", "Video", "MKB"};

	const std::uint32_t n = status.hook_entry_count < kMaxHookEntries ? status.hook_entry_count : kMaxHookEntries;
	if (n == 0)
	{
		return;
	}

	if (!ImGui::CollapsingHeader("Installed hooks", ImGuiTreeNodeFlags_DefaultOpen))
	{
		return;
	}
	if (ImGui::BeginTable("hooks", 3, ImGuiTableFlags_Borders | ImGuiTableFlags_SizingStretchProp))
	{
		ImGui::TableSetupColumn("Hook");
		ImGui::TableSetupColumn("On", ImGuiTableColumnFlags_WidthFixed);
		ImGui::TableSetupColumn("Calls", ImGuiTableColumnFlags_WidthFixed);
		ImGui::TableHeadersRow();

		// Group rows by subsystem so related hooks sit together.
		for (std::uint32_t sub = 0; sub < HookSubsys_Count; ++sub)
		{
			bool header_done = false;
			for (std::uint32_t i = 0; i < n; ++i)
			{
				const HookEntry& e = status.hook_entries[i];
				if (e.subsystem != sub)
				{
					continue;
				}
				if (!header_done)
				{
					ImGui::TableNextRow();
					ImGui::TableNextColumn();
					ImGui::TextDisabled("%s", kSubsysName[sub < HookSubsys_Count ? sub : 0]);
					ImGui::TableNextColumn();
					ImGui::TableNextColumn();
					header_done = true;
				}
				ImGui::TableNextRow();
				ImGui::TableNextColumn();
				ImGui::TextUnformatted(e.name);
				ImGui::TableNextColumn();
				if (e.installed)
				{
					ImGui::TextColored(kGreen, "yes");
				}
				else
				{
					ImGui::TextDisabled("no");
				}
				ImGui::TableNextColumn();
				ImGui::Text("%llu", static_cast<unsigned long long>(e.calls));
			}
		}
		ImGui::EndTable();
	}
}

// True if any hook of `subsystem` is currently installed in the game.
static bool subsystem_installed(const HookStatusView& status, std::uint32_t subsystem)
{
	const std::uint32_t n = status.hook_entry_count < kMaxHookEntries ? status.hook_entry_count : kMaxHookEntries;
	for (std::uint32_t i = 0; i < n; ++i)
	{
		if (status.hook_entries[i].subsystem == subsystem && status.hook_entries[i].installed)
		{
			return true;
		}
	}
	return false;
}

void InjectionPanel::draw_subsystem_controls(const HookStatusView& status)
{
	ImGui::SeparatorText("Subsystems (hook / unhook)");

	struct Row
	{
		const char* label;
		std::uint32_t subsystem;
		bool* want;
		const char* depends;
	};
	const Row rows[] = {
	    {"Input forwarding (XInput)", HookSubsys_Input, &want_input_, "controller input reaches the game"},
	    {"Focus spoof", HookSubsys_Focus, &want_focus_, "the game keeps running unfocused"},
	    {"Audio render-hook", HookSubsys_Audio, &want_audio_, "audio mirror without echo"},
	    {"Video Present-hook", HookSubsys_Video, &want_video_, "Video mirror can use the hooked source"},
	    {"Mouse + keyboard forwarding", HookSubsys_Mkb, &want_mkb_, "clicks/keys reach the game"},
	};

	for (const Row& r : rows)
	{
		ImGui::PushID(r.label);
		if (ImGui::Checkbox(r.label, r.want))
		{
			server_.set_subsystem_enabled(r.subsystem, *r.want);
		}
		ImGui::SameLine();
		const bool on = subsystem_installed(status, r.subsystem);
		if (*r.want != on)
		{
			ImGui::TextColored(kGrey, "(%s...)", *r.want ? "installing" : "removing");
		}
		else
		{
			ImGui::TextColored(on ? kGreen : kGrey, on ? "installed" : "off");
		}
		if (!*r.want)
		{
			ImGui::TextDisabled("   off: %s won't work", r.depends);
		}
		ImGui::PopID();
	}

	// Cursor release is a Focus sub-option for games that clip/recenter the mouse
	// (e.g. Trails through Daybreak), which would otherwise trap the operator.
	if (ImGui::Checkbox("Release operator cursor (free the game's clip)  [F2]", &release_cursor_))
	{
		server_.set_cursor_clip_allowed(!release_cursor_);
	}
}

void InjectionPanel::draw_hook_status(bool debug_details)
{
	if (!server_.running())
	{
		return;
	}

	const HookStatusView status = server_.hook_status();

	ImGui::SeparatorText("Hook status");
	if (!injected_)
	{
		ImGui::TextColored(kGrey, "Not injected.");
		return;
	}

	switch (target_state_)
	{
	case TargetState::Alive:
		ImGui::TextColored(kGreen, "Hook DLL loaded in pid %lu (heartbeat %u)", server_.target_pid(),
		                   status.heartbeat);
		break;
	case TargetState::Hung:
		ImGui::TextColored(kRed, "Target not responding -- heartbeat stalled (frozen?).");
		break;
	case TargetState::Terminated:
		ImGui::TextColored(kRed, "Target process has exited.");
		break;
	case TargetState::NotInjected:
		break;
	}

	// The game is gone or frozen -> its hooks can't act on toggles, so lock them.
	ImGui::BeginDisabled(target_state_ != TargetState::Alive);
	draw_subsystem_controls(status);
	ImGui::EndDisabled();
	ImGui::TextDisabled("Controller poll rates are in the Controllers panel.");

	draw_hook_list(status);

	if (!debug_details)
	{
		return; // everything below is diagnostic detail
	}

	// Focus-API usage: tells us whether the game even consults these (which we
	// spoof) when deciding it lost focus.
	ImGui::Text("Focus API calls  FG:%llu  Active:%llu  Focus:%llu",
	            static_cast<unsigned long long>(status.focus_calls[coop::FocusApi_Foreground]),
	            static_cast<unsigned long long>(status.focus_calls[coop::FocusApi_Active]),
	            static_cast<unsigned long long>(status.focus_calls[coop::FocusApi_Focus]));

	// Input-path diagnostics: a focus-gated detection path would explain a game
	// that only accepts the controller when it has true focus.
	ImGui::SeparatorText("Input path");
	if (status.raw_input_gamepad)
	{
		ImGui::TextColored(status.raw_input_gamepad_sink ? kGreen : kRed, "Raw Input gamepad: yes (INPUTSINK %s)",
		                   status.raw_input_gamepad_sink ? "set -> bg ok" : "MISSING -> focus-gated!");
	}
	else if (status.raw_input_registered)
	{
		ImGui::TextColored(kGrey, "Raw Input: registered, but not for a gamepad usage");
	}
	else
	{
		ImGui::TextColored(kGrey, "Raw Input: not registered");
	}
	ImGui::TextColored(status.dinput_loaded ? kRed : kGrey, "DirectInput dll loaded: %s",
	                   status.dinput_loaded ? "yes (could be foreground-gated)" : "no");
}

void InjectionPanel::draw(bool debug_details)
{
	apply_panel_layout(Panel::Injection);
	ImGui::Begin("Injection");

	if (server_.running())
	{
		if (target_state_ == TargetState::Terminated)
		{
			ImGui::TextColored(kRed, "Target (pid %lu) has terminated.", server_.target_pid());
		}
		else if (target_state_ == TargetState::Hung)
		{
			ImGui::TextColored(kRed, "Target (pid %lu) is not responding.", server_.target_pid());
		}
		else
		{
			ImGui::TextColored(kGreen, "Connected to pid %lu", server_.target_pid());
		}
		if (ImGui::Button("Disconnect"))
		{
			// Leave the game vanilla: unhook everything before dropping the channel. The DLL stays
			// injected (dormant), so it can be reconnected later without re-injecting.
			disconnect_graceful(/*timeout_ms=*/700);
			status_ = "Disconnected (game unhooked; DLL left injected).";
			status_color_ = kGrey;
		}
		// A relaunched game has a new pid; re-attach by image name without hunting for
		// it in the list. Only offered once the old target is gone.
		if (target_state_ == TargetState::Terminated && !selected_name_.empty())
		{
			ImGui::SameLine();
			if (ImGui::Button("Re-attach"))
			{
				reattach();
			}
			ImGui::SameLine();
			ImGui::TextDisabled("(relaunched %s)", narrow(selected_name_).c_str());
		}
		// Session-only auto re-attach: tick it (while attached or terminated), then kill +
		// relaunch the game and it re-injects itself early -- the kill+relaunch fix for a
		// wrong audio format, without picking a target again.
		if (!selected_name_.empty())
		{
			ImGui::Checkbox("Auto re-attach this game on relaunch", &auto_reattach_);
			if (auto_reattach_ && target_state_ == TargetState::Terminated)
			{
				ImGui::SameLine();
				ImGui::TextColored(kGrey, "(watching for %s...)", narrow(selected_name_).c_str());
			}
			// Opt-in Vulkan capture layer: for Vulkan games that initialize Vulkan immediately
			// (where even auto-attach injects too late -- see the red banner), register a per-user
			// implicit layer scoped to this game so the next launch is captured from the first frame.
			// Removed when unticked or the host exits.
			if (ImGui::Checkbox("Set up Vulkan layer (for immediate-init Vulkan games)", &vk_layer_enabled_))
			{
				if (vk_layer_enabled_)
				{
					vk_layer_enabled_ = register_vk_layer(selected_name_);
				}
				else
				{
					unregister_vk_layer();
				}
			}
			if (ImGui::IsItemHovered())
			{
				ImGui::SetTooltip("Registers a per-user (HKCU, no admin) implicit Vulkan layer scoped to\n"
				                  "this game, so a relaunch is captured before Vulkan init. Pair with\n"
				                  "Auto re-attach. Removed when you untick it or close the tool.");
			}
		}
		ImGui::Separator();
	}

	ImGui::TextUnformatted("Target window");
	if (ImGui::Button("Refresh"))
	{
		refresh_targets();
	}
	ImGui::SameLine();
	ImGui::SetNextItemWidth(-1.0f);
	ImGui::InputTextWithHint("##wfilter", "filter by title or process...", window_filter_, sizeof(window_filter_));

	if (ImGui::BeginListBox("##windows", ImVec2(-1.0f, 180.0f)))
	{
		for (const WindowEntry& w : windows_)
		{
			if (!contains_ci(w.title, window_filter_) && !contains_ci(w.exe_name, window_filter_))
			{
				continue;
			}
			const bool selected = w.pid == selected_pid_;
			char label[400];
			snprintf(label, sizeof(label), "%-32s  [%s  %lu]", narrow(w.title).c_str(),
			         narrow(w.exe_name).c_str(), w.pid);
			if (ImGui::Selectable(label, selected))
			{
				selected_pid_ = w.pid;
				selected_name_ = w.exe_name;
			}
		}
		ImGui::EndListBox();
	}

	// The full process list is the advanced fallback (e.g. a windowless game host),
	// kept out of the way unless the operator wants it.
	if (debug_details)
	{
		ImGui::SeparatorText("All processes (advanced)");
		ImGui::SetNextItemWidth(-1.0f);
		ImGui::InputTextWithHint("##filter", "filter by name...", filter_, sizeof(filter_));
		if (ImGui::BeginListBox("##processes", ImVec2(-1.0f, 160.0f)))
		{
			for (const ProcessEntry& entry : processes_)
			{
				if (!contains_ci(entry.exe_name, filter_))
				{
					continue;
				}
				const bool selected = entry.pid == selected_pid_;
				char label[300];
				snprintf(label, sizeof(label), "%-40s  %lu", narrow(entry.exe_name).c_str(), entry.pid);
				if (ImGui::Selectable(label, selected))
				{
					selected_pid_ = entry.pid;
					selected_name_ = entry.exe_name;
				}
			}
			ImGui::EndListBox();
		}
	}

	const bool can_inject = selected_pid_ != 0;
	ImGui::BeginDisabled(!can_inject);
	if (ImGui::Button("Inject & Connect", ImVec2(-1.0f, 0.0f)))
	{
		inject_selected();
	}
	ImGui::EndDisabled();

	if (!status_.empty())
	{
		ImGui::TextColored(status_color_, "%s", status_.c_str());
	}

	draw_hook_status(debug_details);

	record_panel_fit("Injection");
	ImGui::End();
}

} // namespace coop