CoopAllTheThings/vk_layer/coop_vk_layer.cpp

// CoopAllTheThings Vulkan capture layer.
//
// A real Vulkan *implicit layer* the loader inserts at vkCreateInstance -- guaranteed to be in
// the chain before the game resolves vkQueuePresentKHR. This is the reliable early-presence path
// for games that initialize Vulkan immediately (which inject-after-launch + the inline-hook
// vk_hook can't catch). The actual read-back (copy the presented image into the shared keyed-mutex
// texture, off the present thread) is the shared coop::hook::VkCapture; this file only gets the
// layer into the dispatch chain and feeds VkCapture each present.
//
// Scoping: an implicit layer loads into *every* Vulkan app, so the layer only *captures* when it
// recognises the process as the host's target -- env COOP_VK_LAYER_FORCE=1 (tests), or this
// process's image basename matches %TEMP%\coop_vk_target.txt (written by the host's "Set up Vulkan
// layer" checkbox). Otherwise it's a pure pass-through.
//
// The loader/layer interface structs (VkLayer*CreateInfo, VkNegotiateLayerInterface) live in
// vk_layer.h, which Vulkan-Headers doesn't ship, so they're declared here to the stable
// loader-interface-version-2 ABI.
#include <atomic>
#include <cstdarg>
#include <cstdio>
#include <cstring>
#include <string>
#include <vector>

#include <windows.h>

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

#define VK_NO_PROTOTYPES
#define VK_USE_PLATFORM_WIN32_KHR
#include <vulkan/vulkan.h>

#include "coop/shared_memory.hpp"
#include "ipc_client.hpp"
#include "vk_capture.hpp"

// --- Loader/layer interface (interface version 2) ---------------------------
extern "C"
{
	typedef enum VkLayerFunction_
	{
		COOP_VK_LAYER_LINK_INFO = 0,
		COOP_VK_LOADER_DATA_CALLBACK = 1,
		COOP_VK_LOADER_LAYER_CREATE_DEVICE_CALLBACK = 2,
		COOP_VK_LOADER_FEATURES = 3,
	} CoopVkLayerFunction;

	typedef PFN_vkVoidFunction(VKAPI_PTR* PFN_GetPhysicalDeviceProcAddr)(VkInstance, const char*);

	typedef struct VkLayerInstanceLink_
	{
		struct VkLayerInstanceLink_* pNext;
		PFN_vkGetInstanceProcAddr pfnNextGetInstanceProcAddr;
		PFN_GetPhysicalDeviceProcAddr pfnNextGetPhysicalDeviceProcAddr;
	} VkLayerInstanceLink;

	typedef struct VkLayerInstanceCreateInfo
	{
		VkStructureType sType; // 1000000000 = VK_STRUCTURE_TYPE_LOADER_INSTANCE_CREATE_INFO
		const void* pNext;
		CoopVkLayerFunction function;
		union {
			VkLayerInstanceLink* pLayerInfo;
			void* pfnCallback; // other callbacks (unused here); keeps the union pointer-sized
		} u;
	} VkLayerInstanceCreateInfo;

	typedef struct VkLayerDeviceLink_
	{
		struct VkLayerDeviceLink_* pNext;
		PFN_vkGetInstanceProcAddr pfnNextGetInstanceProcAddr;
		PFN_vkGetDeviceProcAddr pfnNextGetDeviceProcAddr;
	} VkLayerDeviceLink;

	typedef struct VkLayerDeviceCreateInfo
	{
		VkStructureType sType; // 1000000001 = VK_STRUCTURE_TYPE_LOADER_DEVICE_CREATE_INFO
		const void* pNext;
		CoopVkLayerFunction function;
		union {
			VkLayerDeviceLink* pLayerInfo;
			void* pfnCallback;
		} u;
	} VkLayerDeviceCreateInfo;

	typedef struct VkNegotiateLayerInterface
	{
		uint32_t sType; // 1 = LAYER_NEGOTIATE_INTERFACE_STRUCT
		void* pNext;
		uint32_t loaderLayerInterfaceVersion;
		PFN_vkGetInstanceProcAddr pfnGetInstanceProcAddr;
		PFN_vkGetDeviceProcAddr pfnGetDeviceProcAddr;
		PFN_GetPhysicalDeviceProcAddr pfnGetPhysicalDeviceProcAddr;
	} VkNegotiateLayerInterface;
}

namespace
{
// The loader tags its chain-link structs with small, loader-internal sType values (not the
// 1000000000-range): VK_STRUCTURE_TYPE_LOADER_INSTANCE_CREATE_INFO = 47, _DEVICE = 48. These are
// from the (unvendored) vk_layer.h and are stable across loader versions.
constexpr VkStructureType kLoaderInstanceCreateInfo = static_cast<VkStructureType>(47);
constexpr VkStructureType kLoaderDeviceCreateInfo = static_cast<VkStructureType>(48);

using coop::hook::IpcClient;
using coop::hook::VkCapture;

IpcClient g_ipc;
std::atomic<bool> g_ipc_tried{false}; // reaper-thread side: connect once on the first published frame
bool g_active = false;                // do we capture in this process? (scoping)

// Chain dispatch.
PFN_vkGetInstanceProcAddr g_next_gipa = nullptr;
PFN_vkGetDeviceProcAddr g_next_gdpa = nullptr;
PFN_vkQueuePresentKHR g_real_present = nullptr;
PFN_vkCreateSwapchainKHR g_real_create_swapchain = nullptr;
PFN_vkGetSwapchainImagesKHR g_get_swapchain_images = nullptr; // for our CreateSwapchain image tracking

VkInstance g_instance = VK_NULL_HANDLE;
VkPhysicalDevice g_phys = VK_NULL_HANDLE;
VkDevice g_device = VK_NULL_HANDLE;
std::uint32_t g_qfam = 0;

VkCapture g_cap; // the shared, off-present-thread read-back

struct SwapInfo
{
	VkSwapchainKHR sc;
	VkFormat fmt;
	std::uint32_t w, h;
	std::vector<VkImage> images;
};
std::vector<SwapInfo> g_swaps;

bool eq(const char* a, const char* b)
{
	return std::strcmp(a, b) == 0;
}

// Optional file trace for debugging the chain dispatch (enable with COOP_VK_LAYER_LOG).
void logvk(const char* fmt, ...)
{
	static int enabled = -1;
	if (enabled < 0)
	{
		enabled = GetEnvironmentVariableW(L"COOP_VK_LAYER_LOG", nullptr, 0) != 0 ? 1 : 0;
	}
	if (enabled == 0)
	{
		return;
	}
	wchar_t dir[MAX_PATH] = {};
	if (GetTempPathW(MAX_PATH, dir) == 0)
	{
		return;
	}
	FILE* f = _wfopen((std::wstring(dir) + L"coop_vk_layer.log").c_str(), L"a");
	if (f == nullptr)
	{
		return;
	}
	va_list ap;
	va_start(ap, fmt);
	std::vfprintf(f, fmt, ap);
	va_end(ap);
	std::fputc('\n', f);
	std::fclose(f);
}

// Decide whether this process is the host's capture target (see file header).
bool decide_active()
{
	if (GetEnvironmentVariableW(L"COOP_VK_LAYER_FORCE", nullptr, 0) != 0)
	{
		return true;
	}
	wchar_t dir[MAX_PATH] = {};
	const DWORD n = GetTempPathW(MAX_PATH, dir);
	if (n == 0 || n >= MAX_PATH)
	{
		return false;
	}
	HANDLE f = CreateFileW((std::wstring(dir) + L"coop_vk_target.txt").c_str(), GENERIC_READ, FILE_SHARE_READ,
	                       nullptr, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, nullptr);
	if (f == INVALID_HANDLE_VALUE)
	{
		return false;
	}
	char want[MAX_PATH] = {};
	DWORD got = 0;
	ReadFile(f, want, sizeof(want) - 1, &got, nullptr);
	CloseHandle(f);
	// Trim trailing whitespace/newline.
	while (got > 0 && (want[got - 1] == '\n' || want[got - 1] == '\r' || want[got - 1] == ' '))
	{
		want[--got] = '\0';
	}
	if (got == 0)
	{
		return false;
	}
	wchar_t self[MAX_PATH] = {};
	GetModuleFileNameW(nullptr, self, MAX_PATH);
	const wchar_t* base = wcsrchr(self, L'\\');
	base = base != nullptr ? base + 1 : self;
	char self8[MAX_PATH] = {};
	WideCharToMultiByte(CP_UTF8, 0, base, -1, self8, sizeof(self8), nullptr, nullptr);
	return _stricmp(self8, want) == 0;
}

const SwapInfo* find_swap(VkSwapchainKHR sc)
{
	for (const SwapInfo& s : g_swaps)
	{
		if (s.sc == sc)
		{
			return &s;
		}
	}
	return nullptr;
}

// Per-second present-rate trace (enable with COOP_VK_LAYER_LOG): confirms the game keeps its frame
// rate while capturing (the whole point of the off-present-thread read-back) and how many of those
// presents the reaper actually mirrored.
void trace_present_rate()
{
	static int enabled = -1;
	if (enabled < 0)
	{
		enabled = GetEnvironmentVariableW(L"COOP_VK_LAYER_LOG", nullptr, 0) != 0 ? 1 : 0;
	}
	if (enabled == 0)
	{
		return;
	}
	static LARGE_INTEGER freq{};
	static LARGE_INTEGER last{};
	static int count = 0;
	static std::uint64_t last_published = 0;
	if (freq.QuadPart == 0)
	{
		QueryPerformanceFrequency(&freq);
		QueryPerformanceCounter(&last);
	}
	++count;
	LARGE_INTEGER now;
	QueryPerformanceCounter(&now);
	const double sec = static_cast<double>(now.QuadPart - last.QuadPart) / static_cast<double>(freq.QuadPart);
	if (sec >= 1.0)
	{
		const std::uint64_t pub = g_cap.frames_published();
		logvk("present rate %.1f/s  captured %.1f/s (game keeps its rate; capture is off the present thread)",
		      count / sec, (pub - last_published) / sec);
		last = now;
		count = 0;
		last_published = pub;
	}
}

VKAPI_ATTR VkResult VKAPI_CALL layer_QueuePresentKHR(VkQueue queue, const VkPresentInfoKHR* pi)
{
	trace_present_rate(); // unconditional so an inactive (pass-through) run gives a baseline to compare
	if (g_active && g_ipc.connected())
	{
		g_ipc.note_present();
	}
	if (g_active && pi != nullptr && pi->swapchainCount == 1)
	{
		const SwapInfo* s = find_swap(pi->pSwapchains[0]);
		if (s != nullptr && pi->pImageIndices[0] < s->images.size())
		{
			VkSemaphore chained = VK_NULL_HANDLE;
			if (g_cap.present(s->images[pi->pImageIndices[0]], s->fmt, s->w, s->h, pi->pWaitSemaphores,
			                  pi->waitSemaphoreCount, chained))
			{
				VkPresentInfoKHR p = *pi;
				p.waitSemaphoreCount = 1;
				p.pWaitSemaphores = &chained;
				return g_real_present(queue, &p);
			}
		}
	}
	return g_real_present(queue, pi);
}

VKAPI_ATTR VkResult VKAPI_CALL layer_CreateSwapchainKHR(VkDevice device, const VkSwapchainCreateInfoKHR* ci,
                                                       const VkAllocationCallbacks* a, VkSwapchainKHR* out)
{
	const VkResult r = g_real_create_swapchain(device, ci, a, out);
	if (g_active && r == VK_SUCCESS && out && g_get_swapchain_images)
	{
		SwapInfo info{};
		info.sc = *out;
		info.fmt = ci->imageFormat;
		info.w = ci->imageExtent.width;
		info.h = ci->imageExtent.height;
		std::uint32_t n = 0;
		g_get_swapchain_images(device, *out, &n, nullptr);
		info.images.resize(n);
		g_get_swapchain_images(device, *out, &n, info.images.data());
		g_swaps.push_back(std::move(info));
	}
	return r;
}

// Resolve the device functions the read-back needs and start the capture component.
void start_capture(VkDevice dev)
{
	VkCapture::Fns f{};
#define LOAD(field, vkname) f.field = reinterpret_cast<PFN_##vkname>(g_next_gdpa(dev, #vkname))
	LOAD(GetDeviceQueue, vkGetDeviceQueue);
	LOAD(CreateCommandPool, vkCreateCommandPool);
	LOAD(DestroyCommandPool, vkDestroyCommandPool);
	LOAD(AllocateCommandBuffers, vkAllocateCommandBuffers);
	LOAD(BeginCommandBuffer, vkBeginCommandBuffer);
	LOAD(EndCommandBuffer, vkEndCommandBuffer);
	LOAD(ResetCommandBuffer, vkResetCommandBuffer);
	LOAD(CmdPipelineBarrier, vkCmdPipelineBarrier);
	LOAD(CmdCopyImageToBuffer, vkCmdCopyImageToBuffer);
	LOAD(QueueSubmit, vkQueueSubmit);
	LOAD(CreateFence, vkCreateFence);
	LOAD(DestroyFence, vkDestroyFence);
	LOAD(WaitForFences, vkWaitForFences);
	LOAD(ResetFences, vkResetFences);
	LOAD(GetFenceStatus, vkGetFenceStatus);
	LOAD(CreateSemaphore, vkCreateSemaphore);
	LOAD(DestroySemaphore, vkDestroySemaphore);
	LOAD(CreateBuffer, vkCreateBuffer);
	LOAD(DestroyBuffer, vkDestroyBuffer);
	LOAD(GetBufferMemoryRequirements, vkGetBufferMemoryRequirements);
	LOAD(AllocateMemory, vkAllocateMemory);
	LOAD(FreeMemory, vkFreeMemory);
	LOAD(BindBufferMemory, vkBindBufferMemory);
	LOAD(MapMemory, vkMapMemory);
	LOAD(UnmapMemory, vkUnmapMemory);
	LOAD(InvalidateMappedMemoryRanges, vkInvalidateMappedMemoryRanges);
	LOAD(DeviceWaitIdle, vkDeviceWaitIdle);
#undef LOAD
	f.GetPhysicalDeviceMemoryProperties = reinterpret_cast<PFN_vkGetPhysicalDeviceMemoryProperties>(
	    g_next_gipa(g_instance, "vkGetPhysicalDeviceMemoryProperties"));
	g_get_swapchain_images =
	    reinterpret_cast<PFN_vkGetSwapchainImagesKHR>(g_next_gdpa(dev, "vkGetSwapchainImagesKHR"));

	g_cap.init(g_phys, dev, g_qfam, f, GetCurrentProcessId(), [](std::uint32_t w, std::uint32_t h) {
		// Runs on the reaper thread after each frame is published to the shared texture. Connect the
		// IPC channel lazily here (the host may not have created it yet at device-create time), then
		// publish the frame so the host's generation counter advances.
		if (!g_ipc_tried.exchange(true))
		{
			g_ipc.connect(/*attempts=*/40, /*delay_ms=*/25);
		}
		if (g_ipc.connected())
		{
			g_ipc.publish_video_frame(w, h, static_cast<std::uint32_t>(DXGI_FORMAT_R8G8B8A8_UNORM));
		}
	});
}

VKAPI_ATTR PFN_vkVoidFunction VKAPI_CALL layer_gdpa(VkDevice device, const char* name);

VKAPI_ATTR VkResult VKAPI_CALL layer_CreateDevice(VkPhysicalDevice phys, const VkDeviceCreateInfo* ci,
                                                  const VkAllocationCallbacks* a, VkDevice* out)
{
	auto* link = reinterpret_cast<VkLayerDeviceCreateInfo*>(const_cast<void*>(ci->pNext));
	while (link != nullptr &&
	       !(link->sType == kLoaderDeviceCreateInfo && link->function == COOP_VK_LAYER_LINK_INFO))
	{
		link = reinterpret_cast<VkLayerDeviceCreateInfo*>(const_cast<void*>(link->pNext));
	}
	if (link == nullptr)
	{
		logvk("CreateDevice: LINK_INFO not found");
		return VK_ERROR_INITIALIZATION_FAILED;
	}
	PFN_vkGetInstanceProcAddr next_gipa = link->u.pLayerInfo->pfnNextGetInstanceProcAddr;
	PFN_vkGetDeviceProcAddr next_gdpa = link->u.pLayerInfo->pfnNextGetDeviceProcAddr;
	link->u.pLayerInfo = link->u.pLayerInfo->pNext; // advance the chain for the next layer
	auto create = reinterpret_cast<PFN_vkCreateDevice>(next_gipa(g_instance, "vkCreateDevice"));
	const VkResult r = create(phys, ci, a, out);
	logvk("CreateDevice: result=%d active=%d", (int)r, g_active ? 1 : 0);
	if (r == VK_SUCCESS && out != nullptr && g_device == VK_NULL_HANDLE)
	{
		g_phys = phys;
		g_device = *out;
		g_next_gdpa = next_gdpa;
		g_qfam = ci->queueCreateInfoCount > 0 ? ci->pQueueCreateInfos[0].queueFamilyIndex : 0;
		g_real_present = reinterpret_cast<PFN_vkQueuePresentKHR>(next_gdpa(*out, "vkQueuePresentKHR"));
		g_real_create_swapchain = reinterpret_cast<PFN_vkCreateSwapchainKHR>(next_gdpa(*out, "vkCreateSwapchainKHR"));
		if (g_active)
		{
			start_capture(*out);
		}
	}
	return r;
}

VKAPI_ATTR VkResult VKAPI_CALL layer_CreateInstance(const VkInstanceCreateInfo* ci,
                                                    const VkAllocationCallbacks* a, VkInstance* out)
{
	auto* link = reinterpret_cast<VkLayerInstanceCreateInfo*>(const_cast<void*>(ci->pNext));
	while (link != nullptr &&
	       !(link->sType == kLoaderInstanceCreateInfo && link->function == COOP_VK_LAYER_LINK_INFO))
	{
		link = reinterpret_cast<VkLayerInstanceCreateInfo*>(const_cast<void*>(link->pNext));
	}
	if (link == nullptr)
	{
		logvk("CreateInstance: LINK_INFO not found");
		return VK_ERROR_INITIALIZATION_FAILED;
	}
	PFN_vkGetInstanceProcAddr next_gipa = link->u.pLayerInfo->pfnNextGetInstanceProcAddr;
	link->u.pLayerInfo = link->u.pLayerInfo->pNext; // advance the chain
	auto create = reinterpret_cast<PFN_vkCreateInstance>(next_gipa(nullptr, "vkCreateInstance"));
	const VkResult r = create(ci, a, out);
	if (r == VK_SUCCESS && out != nullptr)
	{
		g_instance = *out;
		g_next_gipa = next_gipa;
		g_active = decide_active();
	}
	logvk("CreateInstance: result=%d active=%d", (int)r, g_active ? 1 : 0);
	return r;
}

VKAPI_ATTR void VKAPI_CALL layer_DestroyDevice(VkDevice device, const VkAllocationCallbacks* a)
{
	auto destroy = reinterpret_cast<PFN_vkDestroyDevice>(g_next_gdpa(device, "vkDestroyDevice"));
	if (g_active && device == g_device)
	{
		g_cap.shutdown(); // joins the reaper, drains the device, frees the read-back resources
		g_swaps.clear();
		g_device = VK_NULL_HANDLE;
	}
	destroy(device, a);
}

VKAPI_ATTR void VKAPI_CALL layer_DestroyInstance(VkInstance instance, const VkAllocationCallbacks* a)
{
	auto destroy = reinterpret_cast<PFN_vkDestroyInstance>(g_next_gipa(instance, "vkDestroyInstance"));
	g_instance = VK_NULL_HANDLE;
	destroy(instance, a);
}

VKAPI_ATTR PFN_vkVoidFunction VKAPI_CALL layer_gdpa(VkDevice device, const char* name)
{
	if (name == nullptr)
	{
		return nullptr;
	}
	if (eq(name, "vkGetDeviceProcAddr"))
		return reinterpret_cast<PFN_vkVoidFunction>(&layer_gdpa);
	if (eq(name, "vkQueuePresentKHR"))
		return reinterpret_cast<PFN_vkVoidFunction>(&layer_QueuePresentKHR);
	if (eq(name, "vkCreateSwapchainKHR"))
		return reinterpret_cast<PFN_vkVoidFunction>(&layer_CreateSwapchainKHR);
	if (eq(name, "vkDestroyDevice"))
		return reinterpret_cast<PFN_vkVoidFunction>(&layer_DestroyDevice);
	return g_next_gdpa != nullptr ? g_next_gdpa(device, name) : nullptr;
}

VKAPI_ATTR PFN_vkVoidFunction VKAPI_CALL layer_gipa(VkInstance instance, const char* name)
{
	if (name == nullptr)
	{
		return nullptr;
	}
	if (eq(name, "vkGetInstanceProcAddr"))
		return reinterpret_cast<PFN_vkVoidFunction>(&layer_gipa);
	if (eq(name, "vkCreateInstance"))
		return reinterpret_cast<PFN_vkVoidFunction>(&layer_CreateInstance);
	if (eq(name, "vkCreateDevice"))
		return reinterpret_cast<PFN_vkVoidFunction>(&layer_CreateDevice);
	if (eq(name, "vkDestroyInstance"))
		return reinterpret_cast<PFN_vkVoidFunction>(&layer_DestroyInstance);
	if (eq(name, "vkGetDeviceProcAddr"))
		return reinterpret_cast<PFN_vkVoidFunction>(&layer_gdpa);
	return g_next_gipa != nullptr ? g_next_gipa(instance, name) : nullptr;
}

} // namespace

extern "C" __declspec(dllexport) VkResult VKAPI_CALL
    vkNegotiateLoaderLayerInterfaceVersion(VkNegotiateLayerInterface* pVersionStruct)
{
	logvk("negotiate: requestedVersion=%u", pVersionStruct->loaderLayerInterfaceVersion);
	if (pVersionStruct->loaderLayerInterfaceVersion > 2)
	{
		pVersionStruct->loaderLayerInterfaceVersion = 2;
	}
	pVersionStruct->pfnGetInstanceProcAddr = layer_gipa;
	pVersionStruct->pfnGetDeviceProcAddr = layer_gdpa;
	pVersionStruct->pfnGetPhysicalDeviceProcAddr = nullptr;
	return VK_SUCCESS;
}

// Also export the entry points directly, for loaders that probe them by name.
extern "C" __declspec(dllexport) PFN_vkVoidFunction VKAPI_CALL coop_vkGetInstanceProcAddr(VkInstance i,
                                                                                          const char* n)
{
	return layer_gipa(i, n);
}
extern "C" __declspec(dllexport) PFN_vkVoidFunction VKAPI_CALL coop_vkGetDeviceProcAddr(VkDevice d, const char* n)
{
	return layer_gdpa(d, n);
}