modeco80/hazelnut
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

clean up nvfbclibrary more

Browse source
This commit is contained in:
Lily Tsuru 2024-12-01 01:39:48 -05:00
parent 599472a176
commit 11766f6385
3 changed files with 62 additions and 52 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
agent/src/capture_nvfbc.cpp
View file

@ -79,11 +79,11 @@ namespace hazelnut {
void Shutdown() { void Shutdown() {
NvfbcDestroyInstance(); NvfbcDestroyInstance();
nvfbcLib.close(); nvfbcLib.Close();
} }
bool Initialize() override { bool Initialize() override {
if(!nvfbcLib.load()) { if(!nvfbcLib.Load()) {
return false; return false;
} }

83
agent/src/nvfbc_library.cpp
View file

@ -1,4 +1,5 @@
#include "nvfbc_library.hpp" #include "nvfbc_library.hpp"
#include "NvFBC/nvFBCToSys.h" #include "NvFBC/nvFBCToSys.h"
namespace { namespace {
@ -29,7 +30,7 @@ namespace {
} }
// pBuf must be valid from [0..MAX_PATH] // pBuf must be valid from [0..MAX_PATH]
void GetNvFBCDefaultPath(char* pBuf) { bool GetNvFBCLibraryPath(char* pBuf) {
// clang-format off // clang-format off
constexpr const char kPathFormatString[] = constexpr const char kPathFormatString[] =
#ifdef _WIN64 #ifdef _WIN64
@ -41,83 +42,88 @@ namespace {
#endif #endif
// clang-format on // clang-format on
if(pBuf == nullptr) [[unlikely]]
return false;
// Get the Windows system directory (there's probably a "nicer" way using // 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) // the SH* APIs but, this was previously nastier and using %SystemRoot%. I think this is OK)
char szWinDir[MAX_PATH] {}; char szWinDir[MAX_PATH] {};
if(GetWindowsDirectoryA(&szWinDir[0], MAX_PATH) == 0) { if(GetWindowsDirectoryA(&szWinDir[0], MAX_PATH) == 0) {
strcpy(&pBuf[0], "its fucked"); return false;
return;
} }
sprintf(&pBuf[0], kPathFormatString, szWinDir); sprintf(&pBuf[0], kPathFormatString, szWinDir);
return true;
} }
} // namespace } // namespace
NvFBCLibrary::~NvFBCLibrary() { NvFBCLibrary::~NvFBCLibrary() {
if(m_handle != nullptr) if(m_handle != nullptr)
close(); Close();
} }
bool NvFBCLibrary::load() { bool NvFBCLibrary::Load() {
if(m_handle != nullptr) if(m_handle != nullptr)
return true; return true;
char szDefaultPath[MAX_PATH] {}; char szNvfbcLibPath[MAX_PATH] {};
GetNvFBCDefaultPath(&szDefaultPath[0]);
m_handle = LoadLibraryA(szDefaultPath); if(!GetNvFBCLibraryPath(&szNvfbcLibPath[0]))
return false;
m_handle = LoadLibraryA(szNvfbcLibPath);
if(m_handle == nullptr) { if(m_handle == nullptr) {
return false; return false;
} }
// Load NvFBC exports // Load NvFBC exports
pfn_create = reinterpret_cast<NvFBC_CreateFunctionExType>(GetProcAddress(m_handle, "NvFBC_CreateEx")); pNvFBC_CreateEx = reinterpret_cast<NvFBC_CreateFunctionExType>(GetProcAddress(m_handle, "NvFBC_CreateEx"));
pfn_set_global_flags = reinterpret_cast<NvFBC_SetGlobalFlagsType>(GetProcAddress(m_handle, "NvFBC_SetGlobalFlags")); pNvFBC_SetGlobalFlags = reinterpret_cast<NvFBC_SetGlobalFlagsType>(GetProcAddress(m_handle, "NvFBC_SetGlobalFlags"));
pfn_get_status = reinterpret_cast<NvFBC_GetStatusExFunctionType>(GetProcAddress(m_handle, "NvFBC_GetStatusEx")); pNvFBC_GetStatusEx = reinterpret_cast<NvFBC_GetStatusExFunctionType>(GetProcAddress(m_handle, "NvFBC_GetStatusEx"));
pfn_enable = reinterpret_cast<NvFBC_EnableFunctionType>(GetProcAddress(m_handle, "NvFBC_Enable")); pNvFBC_Enable = reinterpret_cast<NvFBC_EnableFunctionType>(GetProcAddress(m_handle, "NvFBC_Enable"));
if((pfn_create == nullptr) || (pfn_set_global_flags == nullptr) || (pfn_get_status == nullptr) || (pfn_enable == nullptr)) { if((pNvFBC_CreateEx == nullptr) || (pNvFBC_SetGlobalFlags == nullptr) || (pNvFBC_GetStatusEx == nullptr) || (pNvFBC_Enable == nullptr)) {
close(); Close();
return false; return false;
} }
return true; return true;
} }
void NvFBCLibrary::close() { void NvFBCLibrary::Close() {
if(m_handle != nullptr) if(m_handle != nullptr)
FreeLibrary(m_handle); FreeLibrary(m_handle);
m_handle = NULL; m_handle = NULL;
pfn_create = NULL; pNvFBC_CreateEx = NULL;
pfn_get_status = NULL; pNvFBC_GetStatusEx = NULL;
pfn_enable = NULL; pNvFBC_Enable = NULL;
} }
NVFBCRESULT NvFBCLibrary::getStatus(NvFBCStatusEx* status) { NVFBCRESULT NvFBCLibrary::GetStatus(NvFBCStatusEx* status) {
return pfn_get_status((void*)status); return pNvFBC_GetStatusEx((void*)status);
} }
void NvFBCLibrary::setGlobalFlags(DWORD flags, int adapter) { void NvFBCLibrary::SetGlobalFlags(DWORD flags, int adapter) {
setTargetAdapter(adapter); SetTargetAdapter(adapter);
pfn_set_global_flags(flags); pNvFBC_SetGlobalFlags(flags);
} }
NVFBCRESULT NvFBCLibrary::createEx(NvFBCCreateParams* pParams) { NVFBCRESULT NvFBCLibrary::CreateEx(NvFBCCreateParams* pParams) {
return pfn_create((void*)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) if(NULL == m_handle)
return NULL; return NULL;
NVFBCRESULT res = NVFBC_SUCCESS; NVFBCRESULT res = NVFBC_SUCCESS;
NvFBCStatusEx status{}; NvFBCStatusEx status {};
status.dwVersion = NVFBC_STATUS_VER; status.dwVersion = NVFBC_STATUS_VER;
status.dwAdapterIdx = adapter; status.dwAdapterIdx = adapter;
res = getStatus(&status); res = GetStatus(&status);
if(res != NVFBC_SUCCESS) { if(res != NVFBC_SUCCESS) {
return nullptr; return nullptr;
@ -131,31 +137,34 @@ void* NvFBCLibrary::create(DWORD type, DWORD* maxWidth, DWORD* maxHeight, int ad
return nullptr; return nullptr;
} }
NvFBCCreateParams createParams{}; NvFBCCreateParams createParams {};
createParams.dwVersion = NVFBC_CREATE_PARAMS_VER; createParams.dwVersion = NVFBC_CREATE_PARAMS_VER;
createParams.dwInterfaceType = type; createParams.dwInterfaceType = type;
createParams.pDevice = devicePtr; createParams.pDevice = devicePtr;
createParams.dwAdapterIdx = adapter; createParams.dwAdapterIdx = adapter;
res = pfn_create(&createParams); res = pNvFBC_CreateEx(&createParams);
*maxWidth = createParams.dwMaxDisplayWidth; if(maxWidth)
*maxHeight = createParams.dwMaxDisplayHeight; *maxWidth = createParams.dwMaxDisplayWidth;
if(maxHeight)
*maxHeight = createParams.dwMaxDisplayHeight;
return createParams.pNvFBC; return createParams.pNvFBC;
} }
INvFBCToSys_v4* NvFBCLibrary::CreateToSys(DWORD* maxWidth, DWORD* maxHeight, int adapter) { INvFBCToSys_v4* NvFBCLibrary::CreateToSys(DWORD* maxWidth, DWORD* maxHeight, int adapter) {
return static_cast<NvFBCToSys*>(create(NVFBC_TO_SYS, maxWidth, maxHeight, adapter, nullptr)); return static_cast<NvFBCToSys*>(CreateRaw(NVFBC_TO_SYS, maxWidth, maxHeight, adapter, nullptr));
} }
NVFBCRESULT NvFBCLibrary::enable(NVFBC_STATE nvFBCState) { NVFBCRESULT NvFBCLibrary::Enable(NVFBC_STATE nvFBCState) {
return pfn_enable(nvFBCState); return pNvFBC_Enable(nvFBCState);
} }
void NvFBCLibrary::setTargetAdapter(int adapter) { void NvFBCLibrary::SetTargetAdapter(int adapter) {
char targetAdapter[10]{}; char targetAdapter[10] {};
_snprintf_s(targetAdapter, 10, 9, "%d", adapter); _snprintf_s(targetAdapter, 10, 9, "%d", adapter);
SetEnvironmentVariableA("NVFBC_TARGET_ADAPTER", targetAdapter); SetEnvironmentVariableA("NVFBC_TARGET_ADAPTER", targetAdapter);
} }

27
agent/src/nvfbc_library.hpp
View file

@ -18,37 +18,38 @@ class NvFBCLibrary {
~NvFBCLibrary(); ~NvFBCLibrary();
/// Tries to load the NvFBC dll /// Tries to load the NvFBC dll
bool load(); bool Load();
/// Close the NvFBC DLL. /// Close the NvFBC DLL.
void close(); void Close();
/// Get the status for the provided adapter, /// Get the status for the provided adapter,
/// if no adapter is provided the default adapter is used. /// 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. /// Sets the global flags for the provided adapter.
/// If no adapter is provided, the default adapter is used /// 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 /// 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); INvFBCToSys_v4* CreateToSys(DWORD* maxWidth, DWORD* maxHeight, int adapter = 0);
/// enable/disable NVFBC /// enable/disable NVFBC
NVFBCRESULT enable(NVFBC_STATE nvFBCState); NVFBCRESULT Enable(NVFBC_STATE nvFBCState);
protected: 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: private:
HMODULE m_handle { nullptr }; HMODULE m_handle { nullptr };
NvFBC_GetStatusExFunctionType pfn_get_status { nullptr }; NvFBC_GetStatusExFunctionType pNvFBC_GetStatusEx { nullptr };
NvFBC_SetGlobalFlagsType pfn_set_global_flags { nullptr }; NvFBC_SetGlobalFlagsType pNvFBC_SetGlobalFlags { nullptr };
NvFBC_CreateFunctionExType pfn_create { nullptr }; NvFBC_CreateFunctionExType pNvFBC_CreateEx { nullptr };
NvFBC_EnableFunctionType pfn_enable { nullptr }; NvFBC_EnableFunctionType pNvFBC_Enable { nullptr };
}; };