CoopAllTheThings/hook/src/opengl_hook.cpp

#include "opengl_hook.hpp"

#include <atomic>
#include <vector>

#include <windows.h>

#include <d3d11.h>
#include <dxgi1_2.h>

#include <safetyhook.hpp>

#include "coop/protocol.hpp"
#include "coop/shared_memory.hpp"
#include "debug_log.hpp"
#include "hook_guard.hpp"
#include "hook_registry.hpp"

namespace coop::hook
{

namespace
{

DetourGate g_gate; // drains in-flight swap detours before remove frees the shared D3D state

// OpenGL enums (avoid pulling in <GL/gl.h> / linking opengl32 at load time).
constexpr unsigned GL_RGBA = 0x1908;
constexpr unsigned GL_UNSIGNED_BYTE = 0x1401;
constexpr unsigned GL_PACK_ALIGNMENT = 0x0D05;

using PFN_glReadPixels = void(WINAPI*)(int, int, int, int, unsigned, unsigned, void*);
using PFN_glPixelStorei = void(WINAPI*)(unsigned, int);
using PFN_wglGetCurrentContext = HGLRC(WINAPI*)();

IpcClient* g_ipc = nullptr;
unsigned long g_pid = 0;

safetyhook::InlineHook g_hk_swapbuffers;  // gdi32!SwapBuffers
safetyhook::InlineHook g_hk_wglswap;      // opengl32!wglSwapBuffers
int g_id_swapbuffers = -1;
int g_id_wglswap = -1;

std::atomic<std::uint64_t> g_swaps{0};
std::atomic<std::uint64_t> g_frames_shared{0};

PFN_glReadPixels g_glReadPixels = nullptr;
PFN_glPixelStorei g_glPixelStorei = nullptr;
PFN_wglGetCurrentContext g_wglGetCurrentContext = nullptr;
bool g_gl_resolved = false;
bool g_unsupported_logged = false;

// Our own D3D11 device hosting the shared texture (the GL path has no D3D device).
ID3D11Device* g_device = nullptr;
ID3D11DeviceContext* g_ctx = nullptr;
ID3D11Texture2D* g_shared_tex = nullptr;
IDXGIKeyedMutex* g_shared_mutex = nullptr;
HANDLE g_shared_handle = nullptr;
UINT g_share_w = 0;
UINT g_share_h = 0;

std::vector<unsigned char> g_read_buf;  // glReadPixels target (bottom-up)
std::vector<unsigned char> g_flip_buf;  // vertically flipped, uploaded to D3D

// GetBuffer/ReleaseBuffer-style re-entrancy guard: wglSwapBuffers may call
// gdi32!SwapBuffers (or vice versa); capture only on the outermost call.
thread_local bool t_in_swap = false;

void resolve_gl()
{
	if (g_gl_resolved)
	{
		return;
	}
	HMODULE gl = GetModuleHandleW(L"opengl32.dll");
	if (gl == nullptr)
	{
		return; // not an OpenGL process (yet)
	}
	g_glReadPixels = reinterpret_cast<PFN_glReadPixels>(GetProcAddress(gl, "glReadPixels"));
	g_glPixelStorei = reinterpret_cast<PFN_glPixelStorei>(GetProcAddress(gl, "glPixelStorei"));
	g_wglGetCurrentContext = reinterpret_cast<PFN_wglGetCurrentContext>(GetProcAddress(gl, "wglGetCurrentContext"));
	g_gl_resolved = (g_glReadPixels != nullptr && g_wglGetCurrentContext != nullptr);
}

bool ensure_device()
{
	if (g_device != nullptr)
	{
		return true;
	}
	const HRESULT hr = D3D11CreateDevice(nullptr, D3D_DRIVER_TYPE_HARDWARE, nullptr, 0, nullptr, 0,
	                                     D3D11_SDK_VERSION, &g_device, nullptr, &g_ctx);
	if (FAILED(hr) || g_device == nullptr)
	{
		logf("opengl: D3D11CreateDevice failed hr=0x%08lX", static_cast<unsigned long>(hr));
		return false;
	}
	return true;
}

void release_shared()
{
	if (g_shared_mutex != nullptr)
	{
		g_shared_mutex->Release();
		g_shared_mutex = nullptr;
	}
	if (g_shared_tex != nullptr)
	{
		g_shared_tex->Release();
		g_shared_tex = nullptr;
	}
	if (g_shared_handle != nullptr)
	{
		CloseHandle(g_shared_handle);
		g_shared_handle = nullptr;
	}
	g_share_w = g_share_h = 0;
}

bool ensure_shared_texture(UINT w, UINT h)
{
	if (g_shared_tex != nullptr && g_share_w == w && g_share_h == h)
	{
		return true;
	}
	release_shared();

	D3D11_TEXTURE2D_DESC desc{};
	desc.Width = w;
	desc.Height = h;
	desc.MipLevels = 1;
	desc.ArraySize = 1;
	desc.Format = DXGI_FORMAT_R8G8B8A8_UNORM; // glReadPixels(GL_RGBA) byte order
	desc.SampleDesc.Count = 1;
	desc.Usage = D3D11_USAGE_DEFAULT;
	desc.BindFlags = D3D11_BIND_SHADER_RESOURCE;
	desc.MiscFlags = D3D11_RESOURCE_MISC_SHARED_NTHANDLE | D3D11_RESOURCE_MISC_SHARED_KEYEDMUTEX;

	if (FAILED(g_device->CreateTexture2D(&desc, nullptr, &g_shared_tex)) || g_shared_tex == nullptr)
	{
		return false;
	}
	IDXGIResource1* res = nullptr;
	if (FAILED(g_shared_tex->QueryInterface(__uuidof(IDXGIResource1), reinterpret_cast<void**>(&res))) ||
	    res == nullptr)
	{
		release_shared();
		return false;
	}
	const std::wstring name = video_share_name(g_pid);
	const HRESULT hr = res->CreateSharedHandle(
	    nullptr, DXGI_SHARED_RESOURCE_READ | DXGI_SHARED_RESOURCE_WRITE, name.c_str(), &g_shared_handle);
	res->Release();
	if (FAILED(hr) || g_shared_handle == nullptr)
	{
		release_shared();
		return false;
	}
	if (FAILED(g_shared_tex->QueryInterface(__uuidof(IDXGIKeyedMutex), reinterpret_cast<void**>(&g_shared_mutex))))
	{
		release_shared();
		return false;
	}
	g_share_w = w;
	g_share_h = h;
	logf("opengl: shared texture ready %ux%u name=%ls", w, h, name.c_str());
	return true;
}

// Read the current GL backbuffer for the window behind `hdc` and upload it.
void capture_gl(HDC hdc)
{
	resolve_gl();
	if (!g_gl_resolved || g_wglGetCurrentContext() == nullptr)
	{
		if (!g_unsupported_logged)
		{
			logf("opengl: no current GL context / glReadPixels; capture idle");
			g_unsupported_logged = true;
		}
		return;
	}

	HWND hwnd = WindowFromDC(hdc);
	RECT rc{};
	if (hwnd == nullptr || !GetClientRect(hwnd, &rc))
	{
		return;
	}
	const UINT w = static_cast<UINT>(rc.right - rc.left);
	const UINT h = static_cast<UINT>(rc.bottom - rc.top);
	if (w == 0 || h == 0)
	{
		return;
	}

	if (!ensure_device() || !ensure_shared_texture(w, h))
	{
		return;
	}

	const size_t bytes = static_cast<size_t>(w) * h * 4;
	if (g_read_buf.size() != bytes)
	{
		g_read_buf.resize(bytes);
		g_flip_buf.resize(bytes);
	}

	if (g_glPixelStorei != nullptr)
	{
		g_glPixelStorei(GL_PACK_ALIGNMENT, 1);
	}
	// Reads the back buffer of the current context (bottom-up, origin lower-left).
	g_glReadPixels(0, 0, static_cast<int>(w), static_cast<int>(h), GL_RGBA, GL_UNSIGNED_BYTE, g_read_buf.data());

	// Flip vertically so the image is top-down like a D3D backbuffer.
	const size_t row = static_cast<size_t>(w) * 4;
	for (UINT y = 0; y < h; ++y)
	{
		memcpy(g_flip_buf.data() + y * row, g_read_buf.data() + (h - 1 - y) * row, row);
	}

	if (g_shared_mutex->AcquireSync(kVideoMutexKey, 8) == S_OK)
	{
		g_ctx->UpdateSubresource(g_shared_tex, 0, nullptr, g_flip_buf.data(), static_cast<UINT>(row), 0);
		g_ctx->Flush();
		g_shared_mutex->ReleaseSync(kVideoMutexKey);
		g_frames_shared.fetch_add(1, std::memory_order_relaxed);
		if (g_ipc != nullptr)
		{
			g_ipc->publish_video_frame(w, h, static_cast<std::uint32_t>(DXGI_FORMAT_R8G8B8A8_UNORM));
		}
	}
}

BOOL WINAPI hk_SwapBuffers(HDC hdc)
{
	DetourGate::Guard guard(g_gate); // keep the shared D3D state alive for this whole detour
	hook_note_call(g_id_swapbuffers);
	g_swaps.fetch_add(1, std::memory_order_relaxed);
	const bool outer = !t_in_swap;
	if (outer)
	{
		t_in_swap = true;
		if (g_ipc != nullptr)
		{
			g_ipc->note_present(); // SwapBuffers is the GL present (counts toward video.present_calls)
		}
		capture_gl(hdc);
	}
	const BOOL r = g_hk_swapbuffers.stdcall<BOOL>(hdc); // __stdcall: call() is __cdecl on x86 -> crash
	if (outer)
	{
		t_in_swap = false;
	}
	return r;
}

BOOL WINAPI hk_wglSwapBuffers(HDC hdc)
{
	DetourGate::Guard guard(g_gate); // keep the shared D3D state alive for this whole detour
	hook_note_call(g_id_wglswap);
	g_swaps.fetch_add(1, std::memory_order_relaxed);
	const bool outer = !t_in_swap;
	if (outer)
	{
		t_in_swap = true;
		if (g_ipc != nullptr)
		{
			g_ipc->note_present(); // wglSwapBuffers is the GL present
		}
		capture_gl(hdc);
	}
	const BOOL r = g_hk_wglswap.stdcall<BOOL>(hdc); // __stdcall: call() is __cdecl on x86 -> crash
	if (outer)
	{
		t_in_swap = false;
	}
	return r;
}

} // namespace

bool install_opengl_hooks(IpcClient& ipc)
{
	g_ipc = &ipc;
	g_pid = GetCurrentProcessId();
	if (g_hk_swapbuffers || g_hk_wglswap)
	{
		return true; // already installed
	}

	g_id_swapbuffers = hook_register("SwapBuffers", HookSubsys_Video);
	g_id_wglswap = hook_register("wglSwapBuffers", HookSubsys_Video);
	g_unsupported_logged = false;

	// gdi32!SwapBuffers is always available (the common GL present call).
	if (HMODULE gdi = GetModuleHandleW(L"gdi32.dll"))
	{
		if (void* fn = reinterpret_cast<void*>(GetProcAddress(gdi, "SwapBuffers")))
		{
			g_hk_swapbuffers = safetyhook::create_inline(fn, reinterpret_cast<void*>(&hk_SwapBuffers));
		}
	}
	// opengl32!wglSwapBuffers if OpenGL is already loaded.
	if (HMODULE gl = GetModuleHandleW(L"opengl32.dll"))
	{
		if (void* fn = reinterpret_cast<void*>(GetProcAddress(gl, "wglSwapBuffers")))
		{
			g_hk_wglswap = safetyhook::create_inline(fn, reinterpret_cast<void*>(&hk_wglSwapBuffers));
		}
	}

	hook_set_installed(g_id_swapbuffers, static_cast<bool>(g_hk_swapbuffers));
	hook_set_installed(g_id_wglswap, static_cast<bool>(g_hk_wglswap));
	logf("install_opengl_hooks: SwapBuffers=%d wglSwapBuffers=%d", static_cast<bool>(g_hk_swapbuffers) ? 1 : 0,
	     static_cast<bool>(g_hk_wglswap) ? 1 : 0);
	return static_cast<bool>(g_hk_swapbuffers) || static_cast<bool>(g_hk_wglswap);
}

void remove_opengl_hooks()
{
	// DISABLE (not destroy) first: restores the SwapBuffers bytes under thread suspension (no new
	// detour) but keeps the trampolines alive, so an in-flight detour about to call .stdcall() (the
	// trampoline) doesn't have it freed under it. Destroying (= {}) before the drain frees the
	// trampoline immediately -- a UAF the uncapped mock-game storm (thousands of swaps/s) can hit.
	// Disable -> drain -> only then destroy.
	(void)g_hk_swapbuffers.disable();
	(void)g_hk_wglswap.disable();
	hook_set_installed(g_id_swapbuffers, false);
	hook_set_installed(g_id_wglswap, false);
	g_gate.drain();
	g_hk_swapbuffers = {}; // no detour in-flight or able to start now -> safe to free the trampolines
	g_hk_wglswap = {};
	release_shared();
	if (g_ctx != nullptr)
	{
		g_ctx->Release();
		g_ctx = nullptr;
	}
	if (g_device != nullptr)
	{
		g_device->Release();
		g_device = nullptr;
	}
	g_read_buf.clear();
	g_flip_buf.clear();
	g_swaps.store(0, std::memory_order_relaxed);
	g_frames_shared.store(0, std::memory_order_relaxed);
	g_gl_resolved = false;
	g_glReadPixels = nullptr;
	g_glPixelStorei = nullptr;
	g_wglGetCurrentContext = nullptr;
	g_ipc = nullptr;
}

std::uint64_t opengl_swaps()
{
	return g_swaps.load(std::memory_order_relaxed);
}

std::uint64_t opengl_frames_shared()
{
	return g_frames_shared.load(std::memory_order_relaxed);
}

} // namespace coop::hook