hazelnut/agent/src/capture_nvfbc.cpp

#include "capture.hpp"

// clang-format off
#include "Utils.hpp"

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

namespace hazelnut {

	void DiffMapWidthHeight(NVFBC_TOSYS_DIFFMAP_BLOCKSIZE blockSize, u32& width, u32& height) {
		switch(blockSize) {
			case NVFBC_TOSYS_DIFFMAP_BLOCKSIZE_16X16:
				width = 16;
				height = 16;
				break;

			case NVFBC_TOSYS_DIFFMAP_BLOCKSIZE_32X32:
				width = 32;
				height = 32;
				break;

			case NVFBC_TOSYS_DIFFMAP_BLOCKSIZE_64X64:
				width = 64;
				height = 64;
				break;

			default:
			case NVFBC_TOSYS_DIFFMAP_BLOCKSIZE_128X128:
				width = 128;
				height = 128;
				break;
		}
	}

	// From Looking Glass, cleaned up to not be macro slop
	u32 DiffMapDimension(u32 srcDim, NVFBC_TOSYS_DIFFMAP_BLOCKSIZE blockSize) {
		u32 blockShift = 0;

		// Calculate shift
		switch(blockSize) {
			case NVFBC_TOSYS_DIFFMAP_BLOCKSIZE_16X16: blockShift = 4; break;

			case NVFBC_TOSYS_DIFFMAP_BLOCKSIZE_32X32: blockShift = 5; break;

			case NVFBC_TOSYS_DIFFMAP_BLOCKSIZE_64X64: blockShift = 6; break;

			default:
			case NVFBC_TOSYS_DIFFMAP_BLOCKSIZE_128X128: blockShift = 7; break;
		}

		return (((srcDim) + (1 << (blockShift)) - 1) >> (blockShift));
	}

	struct NVFBCDisplayCapture final : public DisplayCaptureBase {
		NVFBCRESULT nvfbcStatus;
		NvFBCLibrary nvfbcLib;

		NVFBC_TOSYS_GRAB_FRAME_PARAMS fbcSysGrabParams;
		NVFBC_TOSYS_DIFFMAP_BLOCKSIZE blockSize;
		NvFBCFrameGrabInfo grabInfo;
		NvFBCToSys* nvfbc;

		u32 width;
		u32 height;

		u32* pRawFramebuffer;
		unique_buffer<u32> convertedFramebuffer;

		u8* pDiffMap;
		u32 diffmapWidth;
		u32 diffmapHeight;

	   public:
		virtual ~NVFBCDisplayCapture() { Shutdown(); }

		void Shutdown() {
			NvfbcDestroyInstance();
			nvfbcLib.close();
		}

		bool Initialize() override {
			if(!nvfbcLib.load()) {
				return false;
			}

			if(!NvfbcInitSession())
				return false;

			return true;
		}

		void NvfbcDestroyInstance() {
			if(nvfbc) {
				nvfbc->NvFBCToSysRelease();
				nvfbc = nullptr;

				// reset state while we're here I guess
				pRawFramebuffer = nullptr;
				pDiffMap = nullptr;

				width = 0;
				height = 0;
				diffmapWidth = 0;
				diffmapHeight = 0;
			}
		}

		bool NvfbcCreateInstance() {
			// Destroy an existing NvFBC instance to avoid resource leaks.
			if(nvfbc) {
				NvfbcDestroyInstance();
			}

			DWORD maxDisplayWidth = -1;
			DWORD maxDisplayHeight = -1;
			nvfbc = nvfbcLib.CreateToSys(&maxDisplayWidth, &maxDisplayHeight);

			if(!nvfbc) {
				return false;
			}

			return true;
		}

		bool NvfbcInitSession() {
			if(!NvfbcCreateInstance())
				return false;

			blockSize = NVFBC_TOSYS_DIFFMAP_BLOCKSIZE_32X32;

			// set up a session
			NVFBC_TOSYS_SETUP_PARAMS fbcSysSetupParams{};
			fbcSysSetupParams.dwVersion = NVFBC_TOSYS_SETUP_PARAMS_VER;
			fbcSysSetupParams.eMode = NVFBC_TOSYS_ARGB;
			fbcSysSetupParams.bWithHWCursor = true;
			fbcSysSetupParams.ppBuffer = reinterpret_cast<void**>(&pRawFramebuffer);

			// enable the difference map with configured blocksize
			// (currently hardcoded to 32x32).
			fbcSysSetupParams.bDiffMap = TRUE;
			fbcSysSetupParams.ppDiffMap = reinterpret_cast<void**>(&pDiffMap);
			fbcSysSetupParams.eDiffMapBlockSize = blockSize;

			auto status = nvfbc->NvFBCToSysSetUp(&fbcSysSetupParams);
			if(status != NVFBC_SUCCESS)
				return false;

			fbcSysGrabParams = {};

			// set up grab params
			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 = 1000;
			return true;
		}

		DisplayCaptureResult CaptureFrame() override {
			auto nvStatus = nvfbc->NvFBCToSysGrabFrame(&this->fbcSysGrabParams);

			switch(nvStatus) {
				case NVFBC_SUCCESS: break;

				// Need to recreate the session. If it fails then we fail too.
				case NVFBC_ERROR_INVALIDATED_SESSION: {
					if(!NvfbcInitSession())
						return DisplayCaptureResult::Fail;

					// Recurse. This looks naughty, but will allow us to directly retry
					// the capture. (Plus if this causes issues whatever has happened is probably beyond saving)
					return CaptureFrame();
				} break;

				default: return DisplayCaptureResult::Fail;
			}

			if(width != grabInfo.dwWidth || height != grabInfo.dwHeight) {
				width = grabInfo.dwWidth;
				height = grabInfo.dwHeight;

				convertedFramebuffer.resize(width * height);

				return DisplayCaptureResult::OkButResized;
			}

			// Splat to the converted framebuffer. It doesn't have padding on it.
			auto* pBufferData = (u8*)convertedFramebuffer.data();
			auto* pSrcData = (u8*)&pRawFramebuffer[0];

			for(u32 y = 0; y < grabInfo.dwHeight; ++y) {


				// Convert to BGRA
				// FIXME: Make this SIMD. I can't into this very well
#if 0
				usize srcStart = (y * grabInfo.dwBufferWidth) * 4;
				usize dstStart = (y * width) * 4;
				for(u32 x = 0; x < grabInfo.dwWidth * 4; x += 4) {
					pBufferData[(dstStart + x) + 0] = pSrcData[(srcStart + x) + 2]; // B
					pBufferData[(dstStart + x) + 1] = pSrcData[(srcStart + x) + 1]; // G
					pBufferData[(dstStart + x) + 2] = pSrcData[(srcStart + x) + 0]; // R
					pBufferData[(dstStart + x) + 3] = 0xff;							// A
				}
#endif

				memcpy(&pBufferData[(y * width) * 4], &pRawFramebuffer[(y * grabInfo.dwBufferWidth)], grabInfo.dwWidth * 4);
			}

			// sleep
			// nvfbc->NvFBCToSysGPUBasedCPUSleep(16000);

			return DisplayCaptureResult::Ok;
		}

		FramebufferInformation GetFramebufferInformation() override { return FramebufferInformation { convertedFramebuffer.data(), width, height }; }

		DiffInformation GetDiffInformation() override {
			//diffmapWidth = (u32)ceil((f32)width / 32);
			//diffmapHeight = (u32)ceil((f32)height / 32);

			diffmapWidth = DiffMapDimension(width, blockSize);
			diffmapHeight = DiffMapDimension(height, blockSize);
			return DiffInformation { pDiffMap, diffmapWidth, diffmapHeight };
		}
	};

	DisplayCaptureBase* CreateNVFBCDisplayCapture() {
		return new NVFBCDisplayCapture();
	}
} // namespace hazelnut