CoopAllTheThings/hook/src/xinput_hook.cpp

#include "xinput_hook.hpp"

#include <array>
#include <atomic>
#include <vector>

#include <windows.h>
#include <xinput.h>

#include <safetyhook.hpp>

#include "hook_registry.hpp"

namespace coop::hook
{

namespace
{

// XInput guide-button bit, reported only by the undocumented ordinal-100
// XInputGetStateEx that many games use. Mirrors how Steam/x360ce expose it.
constexpr std::uint16_t kGuideButton = 0x0400;

IpcClient* g_ipc = nullptr;
std::vector<safetyhook::InlineHook> g_hooks;

// Registry ids for the hook list (one per logical export; shared across the
// xinput*.dll variants that may each export it).
int g_id_getstate = -1;
int g_id_getstateex = -1;
int g_id_getcaps = -1;
int g_id_setstate = -1;

// Last good snapshot, so a momentary failed IPC read (host mid-write) doesn't
// flicker the controller as disconnected inside the game.
std::array<CoopPadState, kMaxPads> g_cache;

void refresh_cache()
{
	if (g_ipc == nullptr)
	{
		return;
	}
	CoopPadState pads[kMaxPads];
	std::uint32_t count = 0;
	if (g_ipc->snapshot(pads, count))
	{
		for (std::uint32_t i = 0; i < kMaxPads; ++i)
		{
			g_cache[i] = pads[i];
		}
	}
}

void fill_gamepad(const CoopPadState& pad, XINPUT_GAMEPAD& out)
{
	out.wButtons = pad.buttons;
	out.bLeftTrigger = pad.left_trigger;
	out.bRightTrigger = pad.right_trigger;
	out.sThumbLX = pad.thumb_lx;
	out.sThumbLY = pad.thumb_ly;
	out.sThumbRX = pad.thumb_rx;
	out.sThumbRY = pad.thumb_ry;
}

// Core of every state query. `keep_guide` drops the guide bit for the plain
// (documented) XInputGetState, which must not report it.
DWORD query_state(DWORD user_index, XINPUT_STATE* state, bool keep_guide)
{
	if (state == nullptr || user_index >= kMaxPads)
	{
		return ERROR_DEVICE_NOT_CONNECTED;
	}
	if (g_ipc != nullptr)
	{
		g_ipc->note_state_query(user_index); // proves to the host the game is polling us
	}
	refresh_cache();
	const CoopPadState& pad = g_cache[user_index];
	if (!pad.connected)
	{
		return ERROR_DEVICE_NOT_CONNECTED;
	}

	XINPUT_STATE result = {};
	result.dwPacketNumber = pad.packet;
	fill_gamepad(pad, result.Gamepad);
	if (!keep_guide)
	{
		result.Gamepad.wButtons &= ~kGuideButton;
	}
	*state = result;
	return ERROR_SUCCESS;
}

DWORD WINAPI hk_XInputGetState(DWORD user_index, XINPUT_STATE* state)
{
	hook_note_call(g_id_getstate);
	return query_state(user_index, state, /*keep_guide=*/false);
}

DWORD WINAPI hk_XInputGetStateEx(DWORD user_index, XINPUT_STATE* state)
{
	hook_note_call(g_id_getstateex);
	return query_state(user_index, state, /*keep_guide=*/true);
}

DWORD WINAPI hk_XInputGetCapabilities(DWORD user_index, DWORD /*flags*/, XINPUT_CAPABILITIES* caps)
{
	hook_note_call(g_id_getcaps);
	if (caps == nullptr || user_index >= kMaxPads)
	{
		return ERROR_DEVICE_NOT_CONNECTED;
	}
	if (g_ipc != nullptr)
	{
		g_ipc->note_caps_query(user_index);
	}
	refresh_cache();
	if (!g_cache[user_index].connected)
	{
		return ERROR_DEVICE_NOT_CONNECTED;
	}

	// Advertise a standard wired Xbox-style gamepad with all controls present.
	XINPUT_CAPABILITIES result = {};
	result.Type = XINPUT_DEVTYPE_GAMEPAD;
	result.SubType = XINPUT_DEVSUBTYPE_GAMEPAD;
	result.Flags = 0;
	result.Gamepad.wButtons = 0xF3FF; // all standard buttons reachable
	result.Gamepad.bLeftTrigger = 0xFF;
	result.Gamepad.bRightTrigger = 0xFF;
	result.Gamepad.sThumbLX = static_cast<SHORT>(0x7FFF);
	result.Gamepad.sThumbLY = static_cast<SHORT>(0x7FFF);
	result.Gamepad.sThumbRX = static_cast<SHORT>(0x7FFF);
	result.Gamepad.sThumbRY = static_cast<SHORT>(0x7FFF);
	*caps = result;
	return ERROR_SUCCESS;
}

// Swallow rumble: it would otherwise be sent to whatever physical device sits at
// this index on the host machine. Forwarding it back to the guest is a later
// phase; for now report success so the game's logic is happy.
DWORD WINAPI hk_XInputSetState(DWORD user_index, XINPUT_VIBRATION* /*vibration*/)
{
	hook_note_call(g_id_setstate);
	if (user_index >= kMaxPads || !g_cache[user_index].connected)
	{
		return ERROR_DEVICE_NOT_CONNECTED;
	}
	return ERROR_SUCCESS;
}

void hook_export(HMODULE module, const char* name, void* detour, int registry_id)
{
	if (module == nullptr)
	{
		return;
	}
	if (void* target = reinterpret_cast<void*>(GetProcAddress(module, name)))
	{
		g_hooks.emplace_back(safetyhook::create_inline(target, detour));
		hook_set_installed(registry_id, true);
	}
}

void hook_ordinal(HMODULE module, WORD ordinal, void* detour, int registry_id)
{
	if (module == nullptr)
	{
		return;
	}
	if (void* target = reinterpret_cast<void*>(GetProcAddress(module, MAKEINTRESOURCEA(ordinal))))
	{
		g_hooks.emplace_back(safetyhook::create_inline(target, detour));
		hook_set_installed(registry_id, true);
	}
}

} // namespace

bool install_xinput_hooks(IpcClient& ipc)
{
	if (!g_hooks.empty())
	{
		return true; // already installed
	}
	g_ipc = &ipc;
	refresh_cache();

	g_id_getstate = hook_register("XInputGetState", HookSubsys_Input);
	g_id_getstateex = hook_register("XInputGetStateEx (ord 100)", HookSubsys_Input);
	g_id_getcaps = hook_register("XInputGetCapabilities", HookSubsys_Input);
	g_id_setstate = hook_register("XInputSetState", HookSubsys_Input);

	// A process generally loads exactly one of these, but hook every one that is
	// present so we don't miss the one the game actually calls.
	const wchar_t* modules[] = {L"xinput1_4.dll", L"xinput1_3.dll", L"xinput9_1_0.dll", L"xinputuap.dll"};
	for (const wchar_t* name : modules)
	{
		HMODULE module = GetModuleHandleW(name);
		if (module == nullptr)
		{
			continue;
		}
		hook_export(module, "XInputGetState", reinterpret_cast<void*>(&hk_XInputGetState), g_id_getstate);
		hook_ordinal(module, 100, reinterpret_cast<void*>(&hk_XInputGetStateEx), g_id_getstateex);
		hook_export(module, "XInputGetCapabilities", reinterpret_cast<void*>(&hk_XInputGetCapabilities),
		            g_id_getcaps);
		hook_export(module, "XInputSetState", reinterpret_cast<void*>(&hk_XInputSetState), g_id_setstate);
	}
	if (!g_hooks.empty())
	{
		g_ipc->mark_attached();
		return true;
	}
	return false;
}

void remove_xinput_hooks()
{
	g_hooks.clear(); // InlineHook destructor restores the original bytes
	hook_set_installed(g_id_getstate, false);
	hook_set_installed(g_id_getstateex, false);
	hook_set_installed(g_id_getcaps, false);
	hook_set_installed(g_id_setstate, false);
	g_ipc = nullptr;
}

} // namespace coop::hook