diff --git a/src/service.cpp b/src/service.cpp index 879639e..9c146c2 100644 --- a/src/service.cpp +++ b/src/service.cpp @@ -17,11 +17,15 @@ * along with UDPT. If not, see . */ #include "service.hpp" +#include #ifdef WIN32 namespace UDPT { + SERVICE_STATUS_HANDLE Service::s_hServiceStatus = nullptr; + SERVICE_STATUS Service::s_serviceStatus = { 0 }; + Service::Service(const boost::program_options::variables_map& conf) : m_conf(conf) { @@ -32,10 +36,10 @@ namespace UDPT } - void Service::install() + void Service::install(const std::string& config_path) { std::string& binaryPath = getFilename(); - binaryPath = "\"" + binaryPath + "\""; + binaryPath = "\"" + binaryPath + "\" -c \"" + config_path + "\""; std::shared_ptr svcMgr = getServiceManager(SC_MANAGER_CREATE_SERVICE); { SC_HANDLE installedService = ::CreateService(reinterpret_cast(svcMgr.get()), @@ -87,15 +91,119 @@ namespace UDPT { const_cast(m_conf["service.name"].as().c_str()), reinterpret_cast(&Service::serviceMain) }, {0, 0} }; + if (FALSE == ::StartServiceCtrlDispatcher(service)) { throw OSError(); } } + DWORD Service::handler(DWORD controlCode, DWORD dwEventType, LPVOID eventData, LPVOID context) + { + switch (controlCode) + { + case SERVICE_CONTROL_INTERROGATE: + return NO_ERROR; + + case SERVICE_CONTROL_STOP: + { + reportServiceStatus(SERVICE_STOP_PENDING, 0, 3000); + Tracker::getInstance().stop(); + + return NO_ERROR; + } + + default: + return ERROR_CALL_NOT_IMPLEMENTED; + } + } + + void Service::reportServiceStatus(DWORD currentState, DWORD dwExitCode, DWORD dwWaitHint) + { + static DWORD checkpoint = 1; + + if (currentState == SERVICE_STOPPED || currentState == SERVICE_RUNNING) + { + checkpoint = 0; + } + else + { + ++checkpoint; + } + + switch (currentState) + { + case SERVICE_RUNNING: + s_serviceStatus.dwControlsAccepted = SERVICE_ACCEPT_STOP; + break; + + default: + s_serviceStatus.dwControlsAccepted = 0; + } + + s_serviceStatus.dwCheckPoint = checkpoint; + s_serviceStatus.dwCurrentState = currentState; + s_serviceStatus.dwWin32ExitCode = dwExitCode; + s_serviceStatus.dwWaitHint = dwWaitHint; + + ::SetServiceStatus(s_hServiceStatus, &s_serviceStatus); + } + VOID Service::serviceMain(DWORD argc, LPCSTR argv[]) { + boost::log::sources::severity_channel_logger_mt<> logger(boost::log::keywords::channel = "service"); + + wchar_t *commandLine = ::GetCommandLineW(); + int argCount = 0; + std::shared_ptr args(::CommandLineToArgvW(commandLine, &argCount), ::LocalFree); + if (nullptr == args) + { + BOOST_LOG_SEV(logger, boost::log::trivial::fatal) << "Failed parse command-line."; + ::exit(-1); + } + if (3 != argCount) + { + BOOST_LOG_SEV(logger, boost::log::trivial::fatal) << "Bad command-line length (must have exactly 2 arguments)."; + ::exit(-1); + } + + if (std::wstring(args.get()[1]) != L"-c") + { + BOOST_LOG_SEV(logger, boost::log::trivial::fatal) << "Argument 1 must be \"-c\"."; + ::exit(-1); + } + + std::wstring wFilename(args.get()[2]); + std::string cFilename(wFilename.begin(), wFilename.end()); + + boost::program_options::options_description& configOptions = UDPT::Tracker::getConfigOptions(); + boost::program_options::variables_map config; + boost::program_options::basic_parsed_options parsed_options = boost::program_options::parse_config_file(cFilename.c_str(), configOptions); + boost::program_options::store(parsed_options, config); + + s_hServiceStatus = ::RegisterServiceCtrlHandlerEx(config["service.name"].as().c_str(), Service::handler, NULL); + if (nullptr == s_hServiceStatus) + { + BOOST_LOG_SEV(logger, boost::log::trivial::fatal) << "Failed to register service control handler."; + ::exit(-1); + } + + s_serviceStatus.dwServiceType = SERVICE_WIN32_OWN_PROCESS; + s_serviceStatus.dwServiceSpecificExitCode = 0; + + reportServiceStatus(SERVICE_START_PENDING, 0, 0); + + { + UDPT::Tracker& tracker = UDPT::Tracker::getInstance(); + tracker.start(config); + + reportServiceStatus(SERVICE_RUNNING, 0, 0); + + tracker.wait(); + + reportServiceStatus(SERVICE_STOPPED, 0, 0); + } } std::shared_ptr Service::getService(DWORD access) diff --git a/src/service.hpp b/src/service.hpp index 4255756..2cefeb8 100644 --- a/src/service.hpp +++ b/src/service.hpp @@ -18,9 +18,12 @@ */ #pragma once +#include +#include #include #include "multiplatform.h" #include "exceptions.h" +#include "tracker.hpp" #ifdef WIN32 namespace UDPT @@ -33,7 +36,7 @@ namespace UDPT virtual ~Service(); - void install(); + void install(const std::string& config_path); void uninstall(); @@ -45,9 +48,15 @@ namespace UDPT private: const boost::program_options::variables_map& m_conf; + static SERVICE_STATUS_HANDLE s_hServiceStatus; + + static SERVICE_STATUS s_serviceStatus; + std::shared_ptr getService(DWORD access); - static VOID WINAPI handler(DWORD controlCode); + static DWORD WINAPI handler(DWORD controlCode, DWORD dwEventType, LPVOID eventData, LPVOID context); + + static void reportServiceStatus(DWORD currentState, DWORD dwExitCode, DWORD dwWaitHint); static VOID WINAPI serviceMain(DWORD argc, LPCSTR argv[]); diff --git a/vs/UDPT/UDPT.vcxproj b/vs/UDPT/UDPT.vcxproj index 99855c3..22afae6 100644 --- a/vs/UDPT/UDPT.vcxproj +++ b/vs/UDPT/UDPT.vcxproj @@ -59,7 +59,7 @@ true - ws2_32.lib;sqlite3.lib;advapi32.lib + ws2_32.lib;sqlite3.lib;advapi32.lib;shell32.lib Console @@ -78,7 +78,7 @@ - ws2_32.lib;sqlite3.lib;advapi32.lib + ws2_32.lib;sqlite3.lib;advapi32.lib;shell32.lib Console