diff --git a/agent/src/capture_nvfbc.cpp b/agent/src/capture_nvfbc.cpp index 3676c98..a2393f2 100644 --- a/agent/src/capture_nvfbc.cpp +++ b/agent/src/capture_nvfbc.cpp @@ -79,11 +79,11 @@ namespace hazelnut { void Shutdown() { NvfbcDestroyInstance(); - nvfbcLib.close(); + nvfbcLib.Close(); } bool Initialize() override { - if(!nvfbcLib.load()) { + if(!nvfbcLib.Load()) { return false; } diff --git a/agent/src/nvfbc_library.cpp b/agent/src/nvfbc_library.cpp index 3eb17d6..b792395 100644 --- a/agent/src/nvfbc_library.cpp +++ b/agent/src/nvfbc_library.cpp @@ -1,4 +1,5 @@ #include "nvfbc_library.hpp" + #include "NvFBC/nvFBCToSys.h" namespace { @@ -29,7 +30,7 @@ namespace { } // pBuf must be valid from [0..MAX_PATH] - void GetNvFBCDefaultPath(char* pBuf) { + bool GetNvFBCLibraryPath(char* pBuf) { // clang-format off constexpr const char kPathFormatString[] = #ifdef _WIN64 @@ -41,83 +42,88 @@ namespace { #endif // clang-format on + if(pBuf == nullptr) [[unlikely]] + return false; + // Get the Windows system directory (there's probably a "nicer" way using // the SH* APIs but, this was previously nastier and using %SystemRoot%. I think this is OK) char szWinDir[MAX_PATH] {}; if(GetWindowsDirectoryA(&szWinDir[0], MAX_PATH) == 0) { - strcpy(&pBuf[0], "its fucked"); - return; + return false; } sprintf(&pBuf[0], kPathFormatString, szWinDir); + return true; } } // namespace NvFBCLibrary::~NvFBCLibrary() { if(m_handle != nullptr) - close(); + Close(); } -bool NvFBCLibrary::load() { +bool NvFBCLibrary::Load() { if(m_handle != nullptr) return true; - char szDefaultPath[MAX_PATH] {}; - GetNvFBCDefaultPath(&szDefaultPath[0]); + char szNvfbcLibPath[MAX_PATH] {}; - m_handle = LoadLibraryA(szDefaultPath); + if(!GetNvFBCLibraryPath(&szNvfbcLibPath[0])) + return false; + + m_handle = LoadLibraryA(szNvfbcLibPath); if(m_handle == nullptr) { return false; } // Load NvFBC exports - pfn_create = reinterpret_cast(GetProcAddress(m_handle, "NvFBC_CreateEx")); - pfn_set_global_flags = reinterpret_cast(GetProcAddress(m_handle, "NvFBC_SetGlobalFlags")); - pfn_get_status = reinterpret_cast(GetProcAddress(m_handle, "NvFBC_GetStatusEx")); - pfn_enable = reinterpret_cast(GetProcAddress(m_handle, "NvFBC_Enable")); + pNvFBC_CreateEx = reinterpret_cast(GetProcAddress(m_handle, "NvFBC_CreateEx")); + pNvFBC_SetGlobalFlags = reinterpret_cast(GetProcAddress(m_handle, "NvFBC_SetGlobalFlags")); + pNvFBC_GetStatusEx = reinterpret_cast(GetProcAddress(m_handle, "NvFBC_GetStatusEx")); + pNvFBC_Enable = reinterpret_cast(GetProcAddress(m_handle, "NvFBC_Enable")); - if((pfn_create == nullptr) || (pfn_set_global_flags == nullptr) || (pfn_get_status == nullptr) || (pfn_enable == nullptr)) { - close(); + if((pNvFBC_CreateEx == nullptr) || (pNvFBC_SetGlobalFlags == nullptr) || (pNvFBC_GetStatusEx == nullptr) || (pNvFBC_Enable == nullptr)) { + Close(); return false; } return true; } -void NvFBCLibrary::close() { +void NvFBCLibrary::Close() { if(m_handle != nullptr) FreeLibrary(m_handle); m_handle = NULL; - pfn_create = NULL; - pfn_get_status = NULL; - pfn_enable = NULL; + pNvFBC_CreateEx = NULL; + pNvFBC_GetStatusEx = NULL; + pNvFBC_Enable = NULL; } -NVFBCRESULT NvFBCLibrary::getStatus(NvFBCStatusEx* status) { - return pfn_get_status((void*)status); +NVFBCRESULT NvFBCLibrary::GetStatus(NvFBCStatusEx* status) { + return pNvFBC_GetStatusEx((void*)status); } -void NvFBCLibrary::setGlobalFlags(DWORD flags, int adapter) { - setTargetAdapter(adapter); - pfn_set_global_flags(flags); +void NvFBCLibrary::SetGlobalFlags(DWORD flags, int adapter) { + SetTargetAdapter(adapter); + pNvFBC_SetGlobalFlags(flags); } -NVFBCRESULT NvFBCLibrary::createEx(NvFBCCreateParams* pParams) { - return pfn_create((void*)pParams); +NVFBCRESULT NvFBCLibrary::CreateEx(NvFBCCreateParams* pParams) { + return pNvFBC_CreateEx((void*)pParams); } -void* NvFBCLibrary::create(DWORD type, DWORD* maxWidth, DWORD* maxHeight, int adapter, void* devicePtr) { +void* NvFBCLibrary::CreateRaw(DWORD type, DWORD* maxWidth, DWORD* maxHeight, int adapter, void* devicePtr) { if(NULL == m_handle) return NULL; NVFBCRESULT res = NVFBC_SUCCESS; - NvFBCStatusEx status{}; + NvFBCStatusEx status {}; status.dwVersion = NVFBC_STATUS_VER; status.dwAdapterIdx = adapter; - res = getStatus(&status); + res = GetStatus(&status); if(res != NVFBC_SUCCESS) { return nullptr; @@ -131,31 +137,34 @@ void* NvFBCLibrary::create(DWORD type, DWORD* maxWidth, DWORD* maxHeight, int ad return nullptr; } - NvFBCCreateParams createParams{}; + NvFBCCreateParams createParams {}; createParams.dwVersion = NVFBC_CREATE_PARAMS_VER; createParams.dwInterfaceType = type; createParams.pDevice = devicePtr; createParams.dwAdapterIdx = adapter; - res = pfn_create(&createParams); + res = pNvFBC_CreateEx(&createParams); - *maxWidth = createParams.dwMaxDisplayWidth; - *maxHeight = createParams.dwMaxDisplayHeight; + if(maxWidth) + *maxWidth = createParams.dwMaxDisplayWidth; + + if(maxHeight) + *maxHeight = createParams.dwMaxDisplayHeight; return createParams.pNvFBC; } INvFBCToSys_v4* NvFBCLibrary::CreateToSys(DWORD* maxWidth, DWORD* maxHeight, int adapter) { - return static_cast(create(NVFBC_TO_SYS, maxWidth, maxHeight, adapter, nullptr)); + return static_cast(CreateRaw(NVFBC_TO_SYS, maxWidth, maxHeight, adapter, nullptr)); } -NVFBCRESULT NvFBCLibrary::enable(NVFBC_STATE nvFBCState) { - return pfn_enable(nvFBCState); +NVFBCRESULT NvFBCLibrary::Enable(NVFBC_STATE nvFBCState) { + return pNvFBC_Enable(nvFBCState); } -void NvFBCLibrary::setTargetAdapter(int adapter) { - char targetAdapter[10]{}; +void NvFBCLibrary::SetTargetAdapter(int adapter) { + char targetAdapter[10] {}; _snprintf_s(targetAdapter, 10, 9, "%d", adapter); SetEnvironmentVariableA("NVFBC_TARGET_ADAPTER", targetAdapter); } \ No newline at end of file diff --git a/agent/src/nvfbc_library.hpp b/agent/src/nvfbc_library.hpp index 6c290d1..776b5ca 100755 --- a/agent/src/nvfbc_library.hpp +++ b/agent/src/nvfbc_library.hpp @@ -18,37 +18,38 @@ class NvFBCLibrary { ~NvFBCLibrary(); /// Tries to load the NvFBC dll - bool load(); + bool Load(); /// Close the NvFBC DLL. - void close(); + void Close(); /// Get the status for the provided adapter, /// if no adapter is provided the default adapter is used. - NVFBCRESULT getStatus(NvFBCStatusEx* status); + NVFBCRESULT GetStatus(NvFBCStatusEx* status); /// Sets the global flags for the provided adapter. /// If no adapter is provided, the default adapter is used - void setGlobalFlags(DWORD flags, int adapter = 0); + void SetGlobalFlags(DWORD flags, int adapter = 0); /// Creates an instance of the provided NvFBC type if possible - NVFBCRESULT createEx(NvFBCCreateParams* pParams); + NVFBCRESULT CreateEx(NvFBCCreateParams* pParams); - /// Creates an instance of the provided NvFBC type if possible. - void* create(DWORD type, DWORD* maxWidth, DWORD* maxHeight, int adapter = 0, void* devicePtr = NULL); INvFBCToSys_v4* CreateToSys(DWORD* maxWidth, DWORD* maxHeight, int adapter = 0); /// enable/disable NVFBC - NVFBCRESULT enable(NVFBC_STATE nvFBCState); + NVFBCRESULT Enable(NVFBC_STATE nvFBCState); protected: - void setTargetAdapter(int adapter = 0); + void SetTargetAdapter(int adapter = 0); + + /// Creates an instance of the provided NvFBC type if possible. + void* CreateRaw(DWORD type, DWORD* maxWidth, DWORD* maxHeight, int adapter = 0, void* devicePtr = NULL); private: HMODULE m_handle { nullptr }; - NvFBC_GetStatusExFunctionType pfn_get_status { nullptr }; - NvFBC_SetGlobalFlagsType pfn_set_global_flags { nullptr }; - NvFBC_CreateFunctionExType pfn_create { nullptr }; - NvFBC_EnableFunctionType pfn_enable { nullptr }; + NvFBC_GetStatusExFunctionType pNvFBC_GetStatusEx { nullptr }; + NvFBC_SetGlobalFlagsType pNvFBC_SetGlobalFlags { nullptr }; + NvFBC_CreateFunctionExType pNvFBC_CreateEx { nullptr }; + NvFBC_EnableFunctionType pNvFBC_Enable { nullptr }; };