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() {
NvfbcDestroyInstance();
nvfbcLib.close();
nvfbcLib.Close();
}
bool Initialize() override {
if(!nvfbcLib.load()) {
if(!nvfbcLib.Load()) {
return false;
}

83
agent/src/nvfbc_library.cpp
View file

@ -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<NvFBC_CreateFunctionExType>(GetProcAddress(m_handle, "NvFBC_CreateEx"));
pfn_set_global_flags = reinterpret_cast<NvFBC_SetGlobalFlagsType>(GetProcAddress(m_handle, "NvFBC_SetGlobalFlags"));
pfn_get_status = reinterpret_cast<NvFBC_GetStatusExFunctionType>(GetProcAddress(m_handle, "NvFBC_GetStatusEx"));
pfn_enable = reinterpret_cast<NvFBC_EnableFunctionType>(GetProcAddress(m_handle, "NvFBC_Enable"));
pNvFBC_CreateEx = reinterpret_cast<NvFBC_CreateFunctionExType>(GetProcAddress(m_handle, "NvFBC_CreateEx"));
pNvFBC_SetGlobalFlags = reinterpret_cast<NvFBC_SetGlobalFlagsType>(GetProcAddress(m_handle, "NvFBC_SetGlobalFlags"));
pNvFBC_GetStatusEx = reinterpret_cast<NvFBC_GetStatusExFunctionType>(GetProcAddress(m_handle, "NvFBC_GetStatusEx"));
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)) {
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<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) {
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);
}

27
agent/src/nvfbc_library.hpp
View file

@ -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 };
};