hazelnut/agent/CollabVMFbcAgent/CollabVMFbcAgent.cpp

// clang-format off
#pragma comment(lib, "ws2_32.lib")
#include <WinSock2.h>
#include <ws2tcpip.h>

#include <windows.h>
#include <stdio.h>
#include <string>
#include <vector>

#include "Utils.hpp"

#include "NvFBCLibrary.h"
#include <NvFBC/nvFBCToSys.h>
// clang-format on

#pragma pack(push, 1)
enum class MessageType : u32 {
	Resize, // tResizeMessage
	Data,	// tDataMessage
};

struct tMessageHeader {
	MessageType type;
	u32 datalen;
	// data
};

struct tResizeMessage {
	u32 width;
	u32 height;
};

struct tTile {
	u32 tile_x;
	u32 tile_y;

	u32 tile_width;
	u32 tile_height;

	// u32 tile_rgba[tile_width*tile_height]
};

struct tDataMessage {
	u32 tileCount;
};

#pragma pack(pop)

struct tileRect {
	u32 x, y, width, height;
};

// client for streamserver
class cStreamClient {
	SOCKET tcpSocket { -1 };

   public:
	cStreamClient() = default;

	~cStreamClient() {
		if(tcpSocket != -1) {
			closesocket(tcpSocket);
			tcpSocket = -1;
		}
	}

	bool Connect(const char* address, int port) {
		tcpSocket = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
		if(tcpSocket == -1) {
			return false;
		}

		sockaddr_in clientSvc {};
		clientSvc.sin_family = AF_INET;
		inet_pton(AF_INET, address, &clientSvc.sin_addr.s_addr);
		//		clientSvc.sin_addr.s_addr = inet_addr(address);
		clientSvc.sin_port = htons(port);

		if(connect(tcpSocket, (SOCKADDR*)&clientSvc, sizeof(clientSvc)) == SOCKET_ERROR) {
			printf("No connection socket. Fuck you\n");
			return false;
		}

		// disable Nagle's alrogithm
		int nodelay = 1;
		setsockopt(tcpSocket, IPPROTO_TCP, TCP_NODELAY, (const char*)&nodelay, sizeof(nodelay));
		return true;
	}

	void SendResize(const tResizeMessage& resize) {
		tMessageHeader header;
		header.type = MessageType::Resize;
		header.datalen = sizeof(tResizeMessage);

		send(tcpSocket, (const char*)&header, sizeof(header), 0);
		send(tcpSocket, (const char*)&resize, sizeof(resize), 0);
	}

	void SendData(const unique_buffer<u32>& data, u32 width, u32 height, const std::vector<tileRect>& tiles) {
		tMessageHeader header;
		header.type = MessageType::Data;
		// header.datalen = data.get_size() * sizeof(UINT32);
		header.datalen = sizeof(tDataMessage);

		auto* pData = data.data();

		// send tile header
		bool sendFull = false;
		tDataMessage dm;
		if(tiles.empty()) {
			dm.tileCount = 1;
			sendFull = true;
		} else {
			dm.tileCount = (u32)tiles.size();
		}

		// we have writev() at home
		// writev() at home:
		send(tcpSocket, (const char*)&header, sizeof(header), 0);
		send(tcpSocket, (const char*)&dm, sizeof(dm), 0);

		// send each tile
		if(!sendFull) {
			for(auto& tile : tiles) {
				tTile tile_wire { tile.x, tile.y, tile.width, tile.height };

				std::vector<u32> data;
				data.resize(tile.width * tile.height * 4);

				for(u32 y = 0; y < tile.height; ++y) {
					auto* pTileLineStart = &pData[(tile.y + y) * width + tile.x];
					memcpy(&data[y * tile.width], pTileLineStart, tile.width * sizeof(u32));
				}

				// send header
				send(tcpSocket, (const char*)&tile_wire, (int)sizeof(tile_wire), 0);
				send(tcpSocket, (const char*)&data[0], (int)data.size(), 0);

				// send data now
				/*for (auto y = tile.y; y < tile.y + tile.height; ++y) {
					auto* pTileLineStart = &pData[y * width + tile.x];
					send(tcpSocket, (const char*)pTileLineStart, tile.width * sizeof(UINT32), 0);
				}*/
			}
		} else {
			tTile tDummyTile { 0, 0, width, height };

			send(tcpSocket, (const char*)&tDummyTile, sizeof(tDummyTile), 0);
			send(tcpSocket, (const char*)&data.data()[0], (i32)(data.get_size() * sizeof(u32)), 0);
		}

		// send(tcpSocket, (const char*)&data.data()[0], data.get_size() * sizeof(UINT32), 0);
	}
};

// FIXME: make this a class, I've clearly got the ability to pull it into one

// stuff
NvFBCLibrary* nvfbcLibrary = nullptr;
NvFBCToSys* nvfbcToSys = nullptr;

// filled in by NVFBC
unsigned char* frameBuffer = nullptr;
unsigned char* diffMap = nullptr;

bool nvfbcSetup() {
	//! Setup the frame grab
	NVFBC_TOSYS_SETUP_PARAMS fbcSysSetupParams = { 0 };
	fbcSysSetupParams.dwVersion = NVFBC_TOSYS_SETUP_PARAMS_VER;
	fbcSysSetupParams.eMode = NVFBC_TOSYS_ARGB;
	fbcSysSetupParams.bWithHWCursor = true;
	fbcSysSetupParams.ppBuffer = (void**)&frameBuffer;

	// enable a 32x32 difference map
	fbcSysSetupParams.bDiffMap = TRUE;
	fbcSysSetupParams.ppDiffMap = (void**)&diffMap;
	fbcSysSetupParams.eDiffMapBlockSize = NVFBC_TOSYS_DIFFMAP_BLOCKSIZE_32X32;

	auto status = nvfbcToSys->NvFBCToSysSetUp(&fbcSysSetupParams);
	return status == NVFBC_SUCCESS;
}

void nvfbcDestroy() {
	if(nvfbcToSys) {
		nvfbcToSys->NvFBCToSysRelease();
		diffMap = nullptr;
		frameBuffer = nullptr;
	}
}

bool nvfbcCreate() {
	if(nvfbcToSys) {
		nvfbcDestroy();
	}

	DWORD maxDisplayWidth = -1, maxDisplayHeight = -1;

	//! Create an instance of NvFBCToSys
	nvfbcToSys = (NvFBCToSys*)nvfbcLibrary->create(NVFBC_TO_SYS, &maxDisplayWidth, &maxDisplayHeight);
	if(!nvfbcToSys) {
		return false;
	}

	return true;
}

int main(int argc, char** argv) {
	WSADATA data;
	if(WSAStartup(MAKEWORD(2, 2), &data) != NO_ERROR) {
		return 1;
	}

	cStreamClient* client = new cStreamClient();

	if(!client->Connect("192.168.1.149", 9438)) {
		printf("conn failed\n");
		return 1;
	}

	nvfbcLibrary = new NvFBCLibrary();

	//! Load NvFBC
	if(!nvfbcLibrary->load()) {
		fprintf(stderr, "Unable to load the NvFBC library\n");
		return -1;
	}

	if(!nvfbcCreate()) {
		fprintf(stderr, "Unable to create an instance of NvFBC\n");
		return -1;
	}

	u32 width = 0;
	u32 height = 0;
	bool firstFrame = true;
	NVFBCRESULT status = NVFBC_SUCCESS;

	BOOL bRecoveryDone = FALSE;

	NvFBCFrameGrabInfo grabInfo;
	unique_buffer<u32> buffer;

	if(nvfbcSetup()) {
		// Sleep so that ToSysSetUp forces a framebuffer update
		Sleep(100);

		NVFBC_TOSYS_GRAB_FRAME_PARAMS fbcSysGrabParams = { 0 };

		std::vector<tileRect> tiles {};

		// set up grab parameters
		fbcSysGrabParams.dwVersion = NVFBC_TOSYS_GRAB_FRAME_PARAMS_VER;
		fbcSysGrabParams.dwFlags = NVFBC_TOSYS_WAIT_WITH_TIMEOUT;
		fbcSysGrabParams.dwTargetWidth = 0;
		fbcSysGrabParams.dwTargetHeight = 0;
		fbcSysGrabParams.dwStartX = 0;
		fbcSysGrabParams.dwStartY = 0;
		fbcSysGrabParams.eGMode = NVFBC_TOSYS_SOURCEMODE_FULL;
		fbcSysGrabParams.pNvFBCFrameGrabInfo = &grabInfo;
		fbcSysGrabParams.dwWaitTime = 16;

		while(true) {
			// Grab the frame
			status = nvfbcToSys->NvFBCToSysGrabFrame(&fbcSysGrabParams);
			if(status == NVFBC_SUCCESS) {
				bRecoveryDone = FALSE;

				// handle resizing the buffer
				if(width != grabInfo.dwWidth || height != grabInfo.dwHeight) {
					width = grabInfo.dwWidth;
					height = grabInfo.dwHeight;

					buffer.resize(grabInfo.dwWidth * grabInfo.dwHeight);

					firstFrame = true;
					client->SendResize(tResizeMessage { width, height });
				}

				// splat the data into an unpadded buffer
				// REMOVE THIS
				for(u32 y = 0; y < grabInfo.dwHeight; ++y) {
					memcpy(&buffer.data()[y * grabInfo.dwWidth], &frameBuffer[(y * grabInfo.dwBufferWidth) * 4], grabInfo.dwWidth * 4);
				}

				tiles.clear();

				// diffmap
				if(firstFrame == false) {
					u32 dwDiffMapWidth = (u32)ceil((f32)width / 32);
					u32 dwDiffMapHeight = (u32)ceil((f32)height / 32);

					for(u32 y = 0; y < dwDiffMapHeight; ++y) {
						for(u32 x = 0; x < dwDiffMapWidth; ++x) {
							auto& bl = diffMap[y * dwDiffMapWidth + x];
							if(bl != 0) {
								tiles.push_back(tileRect {
								x * (width / dwDiffMapWidth),	// x
								y * (height / dwDiffMapHeight), // y
								width / dwDiffMapWidth,			// width
								height / dwDiffMapHeight		// height
								});
							}
						}
					}
				}

				// send that to the server
				client->SendData(buffer, width, height, tiles);

				if(firstFrame)
					firstFrame = false;
			} else {
				if(bRecoveryDone == TRUE) {
					fprintf(stderr, "Unable to recover from NvFBC Frame grab failure.\n");
					break;
				}

				if(status == NVFBC_ERROR_DYNAMIC_DISABLE) {
					fprintf(stderr, "NvFBC disabled. Quitting\n");
					break;
				}

				// Try to recover the session
				if(status == NVFBC_ERROR_INVALIDATED_SESSION) {
					fprintf(stderr, "Session Invalidated. Attempting recovery\n");

					if(!nvfbcCreate()) {
						fprintf(stderr, "Unable to re-create NvFBC\n");
						break;
					}

					if(nvfbcSetup()) {
						bRecoveryDone = TRUE;
					} else {
						fprintf(stderr, "Unable to recover from NvFBC Frame grab failure.\n");
						break;
					}
				}
			}
		}
	}

	if(status != NVFBC_SUCCESS) {
		fprintf(stderr, "Unable to setup frame grab.\n");
	}

	nvfbcDestroy();

	delete client;
	delete nvfbcLibrary;

	return 0;
}