hazelnut/agent/CollabVMFbcAgent/capture_nvfbc.cpp

#include "capture.hpp"

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

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

namespace hazelnut {

	class FramebufferCapture_NVFBC final : public IFramebufferCapture {
		NVFBCRESULT nvfbcStatus;
		NvFBCLibrary nvfbcLib;

		NVFBC_TOSYS_GRAB_FRAME_PARAMS fbcSysGrabParams;
		NvFBCFrameGrabInfo grabInfo;
		NvFBCToSys* nvfbc;

		u32 width;
		u32 height;

		u32* pRawFramebuffer;
		unique_buffer<u32> convertedFramebuffer;

		u8* pDiffMap;
		u32 diffmapWidth;
		u32 diffmapHeight;

	   public:
		virtual ~FramebufferCapture_NVFBC() { 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 session to avoid memory leak
			if (nvfbc) {
				NvfbcDestroyInstance();
			}

			DWORD maxDisplayWidth = -1;
			DWORD maxDisplayHeight = -1;
			nvfbc = (NvFBCToSys*)nvfbcLib.create(NVFBC_TO_SYS, &maxDisplayWidth, &maxDisplayHeight);

			if (!nvfbc) {
				return false;
			}

			return true;
		}

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

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

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

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

			memset(&fbcSysGrabParams, 0, sizeof(NVFBC_TOSYS_GRAB_FRAME_PARAMS));

			// 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 = 16;
			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;

				diffmapWidth = (u32)ceil((f32)width / 32);
				diffmapHeight = (u32)ceil((f32)height / 32);

				convertedFramebuffer.resize(width * height);

				return DisplayCaptureResult::OkButResized;
			}

			// Splat to the converted framebuffer. It doesn't have padding on it.
			auto* pBufferData = convertedFramebuffer.data();
			for (u32 y = 0; y < grabInfo.dwHeight; ++y) {
				memcpy(&pBufferData[y * width], &pRawFramebuffer[(y * grabInfo.dwBufferWidth)], grabInfo.dwWidth * 4);
			}

			return DisplayCaptureResult::Ok;
		}

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

		DiffInformation GetDiffInformation() override {
			return DiffInformation{
				pDiffMap,
				diffmapWidth,
				diffmapHeight
			};
		}
        
	};

	IFramebufferCapture* CreateFramebufferCapture_NVFBC() {
		return new FramebufferCapture_NVFBC();
	}
} // namespace hazelnut