CoopAllTheThings/host/src/inject/mkb_forward.cpp

#include "inject/mkb_forward.hpp"

#include <imgui.h>

#include "coop/protocol.hpp"
#include "injection_panel.hpp"
#include "inject/mkb_map.hpp"

namespace coop
{

namespace
{

// Map an ImGui key to a Win32 virtual-key. Returns 0 for keys we don't forward.
int imgui_key_to_vk(ImGuiKey k)
{
	if (k >= ImGuiKey_A && k <= ImGuiKey_Z)
	{
		return 'A' + (k - ImGuiKey_A);
	}
	if (k >= ImGuiKey_0 && k <= ImGuiKey_9)
	{
		return '0' + (k - ImGuiKey_0);
	}
	if (k >= ImGuiKey_Keypad0 && k <= ImGuiKey_Keypad9)
	{
		return VK_NUMPAD0 + (k - ImGuiKey_Keypad0);
	}
	if (k >= ImGuiKey_F1 && k <= ImGuiKey_F12)
	{
		return VK_F1 + (k - ImGuiKey_F1);
	}
	switch (k)
	{
	case ImGuiKey_Tab: return VK_TAB;
	case ImGuiKey_LeftArrow: return VK_LEFT;
	case ImGuiKey_RightArrow: return VK_RIGHT;
	case ImGuiKey_UpArrow: return VK_UP;
	case ImGuiKey_DownArrow: return VK_DOWN;
	case ImGuiKey_PageUp: return VK_PRIOR;
	case ImGuiKey_PageDown: return VK_NEXT;
	case ImGuiKey_Home: return VK_HOME;
	case ImGuiKey_End: return VK_END;
	case ImGuiKey_Insert: return VK_INSERT;
	case ImGuiKey_Delete: return VK_DELETE;
	case ImGuiKey_Backspace: return VK_BACK;
	case ImGuiKey_Space: return VK_SPACE;
	case ImGuiKey_Enter: return VK_RETURN;
	case ImGuiKey_Escape: return VK_ESCAPE;
	case ImGuiKey_LeftCtrl: return VK_LCONTROL;
	case ImGuiKey_LeftShift: return VK_LSHIFT;
	case ImGuiKey_LeftAlt: return VK_LMENU;
	case ImGuiKey_LeftSuper: return VK_LWIN;
	case ImGuiKey_RightCtrl: return VK_RCONTROL;
	case ImGuiKey_RightShift: return VK_RSHIFT;
	case ImGuiKey_RightAlt: return VK_RMENU;
	case ImGuiKey_RightSuper: return VK_RWIN;
	case ImGuiKey_Menu: return VK_APPS;
	case ImGuiKey_Apostrophe: return VK_OEM_7;
	case ImGuiKey_Comma: return VK_OEM_COMMA;
	case ImGuiKey_Minus: return VK_OEM_MINUS;
	case ImGuiKey_Period: return VK_OEM_PERIOD;
	case ImGuiKey_Slash: return VK_OEM_2;
	case ImGuiKey_Semicolon: return VK_OEM_1;
	case ImGuiKey_Equal: return VK_OEM_PLUS;
	case ImGuiKey_LeftBracket: return VK_OEM_4;
	case ImGuiKey_Backslash: return VK_OEM_5;
	case ImGuiKey_RightBracket: return VK_OEM_6;
	case ImGuiKey_GraveAccent: return VK_OEM_3;
	case ImGuiKey_CapsLock: return VK_CAPITAL;
	case ImGuiKey_ScrollLock: return VK_SCROLL;
	case ImGuiKey_NumLock: return VK_NUMLOCK;
	case ImGuiKey_PrintScreen: return VK_SNAPSHOT;
	case ImGuiKey_Pause: return VK_PAUSE;
	case ImGuiKey_KeypadDecimal: return VK_DECIMAL;
	case ImGuiKey_KeypadDivide: return VK_DIVIDE;
	case ImGuiKey_KeypadMultiply: return VK_MULTIPLY;
	case ImGuiKey_KeypadSubtract: return VK_SUBTRACT;
	case ImGuiKey_KeypadAdd: return VK_ADD;
	case ImGuiKey_KeypadEnter: return VK_RETURN;
	default: return 0;
	}
}

// Last valid game-client position, so a button release that lands on a letterbox
// bar still goes to the game (no stuck buttons).
int g_last_gx = 0;
int g_last_gy = 0;

// What we've forwarded as currently-held, so we can always release it -- even when forwarding stops
// (mirror off, focus lost, ImGui grabbed the input, the game died) -- so nothing sticks in the guest.
bool g_mouse_down[3] = {};
bool g_key_down[256] = {}; // indexed by VK

void release_held_keys(InjectionPanel& injection)
{
	for (int vk = 0; vk < 256; ++vk)
	{
		if (g_key_down[vk])
		{
			injection.push_mkb(MkbEvent{Mkb_KeyUp, static_cast<std::uint32_t>(vk), 0, 0});
			g_key_down[vk] = false;
		}
	}
}

void release_held_mouse(InjectionPanel& injection)
{
	for (int b = 0; b < 3; ++b)
	{
		if (g_mouse_down[b])
		{
			injection.push_mkb(MkbEvent{Mkb_MouseUp, static_cast<std::uint32_t>(b), g_last_gx, g_last_gy});
			g_mouse_down[b] = false;
		}
	}
}

} // namespace

void forward_mkb_frame(InjectionPanel& injection, HWND host_hwnd, bool mirroring, bool source_hooked)
{
	const HWND game = injection.game_hwnd();
	// Under RPT the guest's MKB lands on our (focused) window; only forward then, so the operator's
	// own desktop use isn't injected. If we can't forward for ANY reason -- subsystem off, we lost
	// focus, or the game is gone -- release everything we're still holding first, so a key/button
	// held at that moment doesn't stick down in the guest.
	if (!injection.mkb_enabled() || GetForegroundWindow() != host_hwnd || game == nullptr || !IsWindow(game))
	{
		release_held_keys(injection);
		release_held_mouse(injection);
		return;
	}

	ImGuiIO& io = ImGui::GetIO();

	// --- Keyboard (unless ImGui is using it for e.g. a text field -- then release what we hold) ---
	if (io.WantCaptureKeyboard)
	{
		release_held_keys(injection);
	}
	else
	{
		for (ImGuiKey k = ImGuiKey_NamedKey_BEGIN; k < ImGuiKey_NamedKey_END; k = static_cast<ImGuiKey>(k + 1))
		{
			const int vk = imgui_key_to_vk(k);
			if (vk == 0)
			{
				continue;
			}
			if (ImGui::IsKeyPressed(k, false))
			{
				injection.push_mkb(MkbEvent{Mkb_KeyDown, static_cast<std::uint32_t>(vk), 0, 0});
				g_key_down[vk & 0xFF] = true;
			}
			if (ImGui::IsKeyReleased(k))
			{
				injection.push_mkb(MkbEvent{Mkb_KeyUp, static_cast<std::uint32_t>(vk), 0, 0});
				g_key_down[vk & 0xFF] = false;
			}
		}
		for (int i = 0; i < io.InputQueueCharacters.Size; ++i)
		{
			const ImWchar c = io.InputQueueCharacters[i];
			if (c != 0)
			{
				injection.push_mkb(MkbEvent{Mkb_Char, static_cast<std::uint32_t>(c), 0, 0});
			}
		}
	}

	// --- Mouse (clicks + wheel only, and only while mirroring and ImGui isn't using the mouse) ---
	// When we're not forwarding the mouse, still release any button we hold (below), so it can't stick.
	const bool forwarding_mouse = mirroring && !io.WantCaptureMouse;
	if (!forwarding_mouse)
	{
		release_held_mouse(injection);
		return;
	}

	RECT host_client{};
	GetClientRect(host_hwnd, &host_client);

	MkbMapInput m;
	m.host_x = static_cast<int>(io.MousePos.x);
	m.host_y = static_cast<int>(io.MousePos.y);
	m.dst_w = host_client.right;
	m.dst_h = host_client.bottom;
	if (source_hooked)
	{
		// Hooked capture mirrors the backbuffer (client area), no decorations.
		const VideoShareView v = injection.video_share();
		m.src_w = m.client_w = static_cast<int>(v.width);
		m.src_h = m.client_h = static_cast<int>(v.height);
	}
	else
	{
		// WGC captures the whole window; the client area sits at a decoration offset.
		RECT wr{}, cr{};
		POINT client_origin{0, 0};
		GetWindowRect(game, &wr);
		GetClientRect(game, &cr);
		ClientToScreen(game, &client_origin);
		m.src_w = wr.right - wr.left;
		m.src_h = wr.bottom - wr.top;
		m.client_w = cr.right;
		m.client_h = cr.bottom;
		m.client_off_x = client_origin.x - wr.left;
		m.client_off_y = client_origin.y - wr.top;
	}

	int gx = 0, gy = 0;
	const bool on_game = map_host_to_game_client(m, gx, gy);
	if (on_game)
	{
		g_last_gx = gx;
		g_last_gy = gy;
	}
	const int mx = on_game ? gx : g_last_gx;
	const int my = on_game ? gy : g_last_gy;

	for (int button = 0; button < 3; ++button) // 0=left, 1=right, 2=middle
	{
		if (on_game && ImGui::IsMouseClicked(button))
		{
			injection.push_mkb(MkbEvent{Mkb_MouseDown, static_cast<std::uint32_t>(button), mx, my});
			g_mouse_down[button] = true;
		}
		if (ImGui::IsMouseReleased(button)) // send the up even off-game, so nothing sticks
		{
			injection.push_mkb(MkbEvent{Mkb_MouseUp, static_cast<std::uint32_t>(button), mx, my});
			g_mouse_down[button] = false;
		}
	}
	if (on_game && io.MouseWheel != 0.0f)
	{
		const int delta = static_cast<int>(io.MouseWheel * WHEEL_DELTA);
		injection.push_mkb(MkbEvent{Mkb_Wheel, static_cast<std::uint32_t>(delta), mx, my});
	}
}

} // namespace coop