From 01b10b3ee135f5d0486b059aaf546118799dc583 Mon Sep 17 00:00:00 2001 From: Geoffrey McRae Date: Wed, 1 Nov 2017 03:23:46 +1100 Subject: [PATCH] [host] converted program to a windows hidden/background user app. NvFBC will not start when running as a service, likely intentional security of the API to prevent it from being abused. --- host/Capture/NvFBC.cpp | 20 ++++- host/Service.cpp | 7 +- host/Service.h | 2 +- host/kvm-ivshmem-host.vcxproj | 8 +- host/main.cpp | 157 ++-------------------------------- 5 files changed, 32 insertions(+), 162 deletions(-) diff --git a/host/Capture/NvFBC.cpp b/host/Capture/NvFBC.cpp index 97e71207..dadf11a7 100644 --- a/host/Capture/NvFBC.cpp +++ b/host/Capture/NvFBC.cpp @@ -55,13 +55,25 @@ namespace Capture status.dwVersion = NVFBC_STATUS_VER; status.dwAdapterIdx = 0; - if (m_fnGetStatusEx(&status) != NVFBC_SUCCESS) + NVFBCRESULT ret = m_fnGetStatusEx(&status); + if (ret != NVFBC_SUCCESS) { - DEBUG_ERROR("Failed to get NvFBC status"); - DeInitialize(); - return false; + DEBUG_INFO("Attempting to enable NvFBC"); + if (m_fnEnable(NVFBC_STATE_ENABLE) == NVFBC_SUCCESS) + { + DEBUG_INFO("Success, attempting to get status again"); + ret = m_fnGetStatusEx(&status); + } + + if (ret != NVFBC_SUCCESS) + { + DEBUG_ERROR("Failed to get NvFBC status"); + DeInitialize(); + return false; + } } + if (!status.bIsCapturePossible) { DEBUG_ERROR("Capture is not possible, unsupported device or driver"); diff --git a/host/Service.cpp b/host/Service.cpp index 28ad834e..8c87da65 100644 --- a/host/Service.cpp +++ b/host/Service.cpp @@ -97,7 +97,7 @@ void Service::DeInitialize() m_initialized = false; } -bool Service::Process(HANDLE stopEvent) +bool Service::Process() { if (!m_initialized) return false; @@ -142,10 +142,7 @@ bool Service::Process(HANDLE stopEvent) // wait for the host to notify that is it is ready to proceed ResetEvent(m_readyEvent); - while( - stopEvent == INVALID_HANDLE_VALUE || - (WaitForSingleObject(stopEvent, 0) == WAIT_OBJECT_0) - ) + while(true) { if (!m_ivshmem->RingDoorbell(header->hostID, 0)) { diff --git a/host/Service.h b/host/Service.h index 9bc35c7b..c89fef7e 100644 --- a/host/Service.h +++ b/host/Service.h @@ -19,7 +19,7 @@ public: bool Initialize(); void DeInitialize(); - bool Process(HANDLE stopEvent); + bool Process(); private: static Service * m_instance; diff --git a/host/kvm-ivshmem-host.vcxproj b/host/kvm-ivshmem-host.vcxproj index 7fca4ddf..98c3f669 100644 --- a/host/kvm-ivshmem-host.vcxproj +++ b/host/kvm-ivshmem-host.vcxproj @@ -94,7 +94,7 @@ ..\;.\;%(AdditionalIncludeDirectories) - Console + Windows true kernel32.lib;user32.lib;gdi32.lib;winspool.lib;comdlg32.lib;advapi32.lib;shell32.lib;ole32.lib;oleaut32.lib;uuid.lib;odbc32.lib;odbccp32.lib;setupapi.lib;%(AdditionalDependencies) %(AdditionalLibraryDirectories) @@ -111,7 +111,7 @@ ..\;.\;%(AdditionalIncludeDirectories) - Console + Windows true kernel32.lib;user32.lib;gdi32.lib;winspool.lib;comdlg32.lib;advapi32.lib;shell32.lib;ole32.lib;oleaut32.lib;uuid.lib;odbc32.lib;odbccp32.lib;setupapi.lib;%(AdditionalDependencies) %(AdditionalLibraryDirectories) @@ -130,7 +130,7 @@ ..\;.\;%(AdditionalIncludeDirectories) - Console + Windows true true true @@ -151,7 +151,7 @@ ..\;.\;%(AdditionalIncludeDirectories) - Console + Windows true true true diff --git a/host/main.cpp b/host/main.cpp index 9f69b60d..51fb35f8 100644 --- a/host/main.cpp +++ b/host/main.cpp @@ -4,161 +4,22 @@ #include "Service.h" -#define SERVICE_NAME "kvm-ivshmem-host" - -//============================================================================= - -struct App +int WINAPI WinMain(HINSTANCE hInstance, HINSTANCE hPrevInstance, PSTR szCmdParam, int iCmdShow) { - SERVICE_STATUS serviceStatus; - SERVICE_STATUS_HANDLE statusHandle; - HANDLE serviceStopEvent; -}; +#ifdef DEBUG + AllocConsole(); +#endif -struct App app = -{ - {0}, - NULL, - INVALID_HANDLE_VALUE -}; - -//============================================================================= - -VOID WINAPI ServiceMain(DWORD argc, LPTSTR *argv); -VOID WINAPI ServiceCtrlHandler(DWORD CtrlCode); -DWORD WINAPI ServiceWorkerThread(LPVOID lpParam); - -//============================================================================= - -int main(int argc, TCHAR *argv[]) -{ - SERVICE_TABLE_ENTRY ServiceTable[] = - { - {_T(SERVICE_NAME), (LPSERVICE_MAIN_FUNCTION)ServiceMain}, - {NULL, NULL} - }; - - if (StartServiceCtrlDispatcher(ServiceTable) == FALSE) - { - DWORD lastError = GetLastError(); - if (lastError == ERROR_FAILED_SERVICE_CONTROLLER_CONNECT) - { - ServiceWorkerThread(INVALID_HANDLE_VALUE); - getc(stdin); - return 0; - } - - return lastError; - } - - return 0; -} - -//============================================================================= - -VOID WINAPI ServiceMain(DWORD argc, LPTSTR *argv) -{ - app.statusHandle = RegisterServiceCtrlHandler(_T(SERVICE_NAME), ServiceCtrlHandler); - if (!app.statusHandle) - return; - - ZeroMemory(&app.serviceStatus, sizeof(app.serviceStatus)); - app.serviceStatus.dwServiceType = SERVICE_WIN32_OWN_PROCESS; - app.serviceStatus.dwControlsAccepted = 0; - app.serviceStatus.dwCurrentState = SERVICE_START_PENDING; - app.serviceStatus.dwWin32ExitCode = 0; - app.serviceStatus.dwServiceSpecificExitCode = 0; - app.serviceStatus.dwCheckPoint = 0; - - if (SetServiceStatus(app.statusHandle, &app.serviceStatus) == FALSE) - { - DEBUG_ERROR("SetServiceStatus failed"); - return; - } - - app.serviceStopEvent = CreateEvent(NULL, TRUE, FALSE, NULL); - if (app.serviceStopEvent == INVALID_HANDLE_VALUE) - { - app.serviceStatus.dwControlsAccepted = 0; - app.serviceStatus.dwCurrentState = SERVICE_STOPPED; - app.serviceStatus.dwWin32ExitCode = GetLastError(); - app.serviceStatus.dwCheckPoint = 1; - if (SetServiceStatus(app.statusHandle, &app.serviceStatus) == FALSE) - { - DEBUG_ERROR("SetServiceStatus failed"); - return; - } - } - - app.serviceStatus.dwControlsAccepted = SERVICE_ACCEPT_STOP; - app.serviceStatus.dwCurrentState = SERVICE_RUNNING; - app.serviceStatus.dwWin32ExitCode = 0; - app.serviceStatus.dwCheckPoint = 0; - if (SetServiceStatus(app.statusHandle, &app.serviceStatus) == FALSE) - { - DEBUG_ERROR("SetServiceStatus failed"); - return; - } - - HANDLE hThread = CreateThread(NULL, 0, ServiceWorkerThread, NULL, 0, NULL); - WaitForSingleObject(hThread, INFINITE); - - CloseHandle(app.serviceStopEvent); - app.serviceStatus.dwControlsAccepted = 0; - app.serviceStatus.dwCurrentState = SERVICE_STOPPED; - app.serviceStatus.dwWin32ExitCode = 0; - app.serviceStatus.dwCheckPoint = 3; - - if (SetServiceStatus(app.statusHandle, &app.serviceStatus) == FALSE) - { - DEBUG_ERROR("SetServiceStatus failed"); - return; - } -} - -//============================================================================= - -VOID WINAPI ServiceCtrlHandler(DWORD CtrlCode) -{ - switch (CtrlCode) - { - case SERVICE_CONTROL_STOP: - if (app.serviceStatus.dwCurrentState != SERVICE_RUNNING) - break; - - app.serviceStatus.dwControlsAccepted = 0; - app.serviceStatus.dwCurrentState = SERVICE_STOP_PENDING; - app.serviceStatus.dwWin32ExitCode = 0; - app.serviceStatus.dwCheckPoint = 4; - - if (SetServiceStatus(app.statusHandle, &app.serviceStatus) == FALSE) - DEBUG_ERROR("SetServiceStatus failed"); - - SetEvent(app.serviceStopEvent); - break; - - default: - break; - } -} - -//============================================================================= - -DWORD WINAPI ServiceWorkerThread(LPVOID lpParam) -{ Service *svc = svc->Get(); if (!svc->Initialize()) { DEBUG_ERROR("Failed to initialize service"); - return ERROR; + return -1; } - while (WaitForSingleObject(app.serviceStopEvent, 0) != WAIT_OBJECT_0) - if (!svc->Process(app.serviceStopEvent)) - break; + while (true) + svc->Process(); svc->DeInitialize(); - return ERROR_SUCCESS; -} - -//============================================================================= \ No newline at end of file + return 0; +} \ No newline at end of file