--- /dev/null
+/*
+ PowerDNS Versatile Database Driven Nameserver
+ Copyright (C) 2002 PowerDNS.COM BV
+
+ This program is free software; you can redistribute it and/or modify
+ it under the terms of the GNU General Public License as published by
+ the Free Software Foundation; either version 2 of the License, or
+ (at your option) any later version.
+
+ This program is distributed in the hope that it will be useful,
+ but WITHOUT ANY WARRANTY; without even the implied warranty of
+ MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ GNU General Public License for more details.
+
+ You should have received a copy of the GNU General Public License
+ along with this program; if not, write to the Free Software
+ Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
+*/
+\r
+/*!\r
+\file ntservice.cpp\r
+\brief This file contains the NTService class implementation.\r
+*/\r
+\r
+#include "utility.hh"\r
+#include <sstream>\r
+#include <iostream>\r
+#include "logger.hh"\r
+#include "ntservice.hh"\r
+\r
+#define L theL("pdns")\r
+\r
+\r
+// Default constructor.\r
+NTService::NTService( void )\r
+{\r
+ m_runningAsService = false;\r
+ m_errorCode = 0;\r
+ m_statusCode = 0;\r
+ m_serviceStatusHandle = NULL;\r
+}\r
+\r
+\r
+// Destructor.\r
+NTService::~NTService( void )\r
+{\r
+}\r
+\r
+\r
+// Returns whether the program is running as a service.\r
+bool NTService::isRunningAsService( void )\r
+{\r
+ return m_runningAsService;\r
+}\r
+\r
+\r
+// Registers the service.\r
+bool NTService::registerService( const std::string & description, bool registerLog )\r
+{\r
+ std::stringstream str;\r
+ HKEY key, pkey;\r
+ SC_HANDLE sc;\r
+ char temp[ 512 ];\r
+ DWORD flags;\r
+ \r
+ sc = OpenSCManager( NULL, NULL, SC_MANAGER_ALL_ACCESS );\r
+ if ( sc == NULL )\r
+ return false; // Could not open the Service Control Manager.\r
+\r
+ GetModuleFileName( NULL, temp, sizeof( temp ));\r
+\r
+ str << temp << " --ntservice";\r
+ if ( CreateService( \r
+ sc, \r
+ getServiceName().c_str(), \r
+ getServiceName().c_str(),\r
+ SERVICE_ALL_ACCESS,\r
+ SERVICE_WIN32_OWN_PROCESS,\r
+ SERVICE_AUTO_START,\r
+ SERVICE_ERROR_NORMAL,\r
+ str.str().c_str(),\r
+ NULL,\r
+ NULL,\r
+ NULL,\r
+ NULL,\r
+ NULL ) == NULL && GetLastError() != ERROR_SERVICE_EXISTS )\r
+ {\r
+ return false; // Don't we all like functions with 43 billion parameters?\r
+ }\r
+\r
+ CloseServiceHandle( sc );\r
+\r
+ str.str( "" );\r
+ \r
+ // Set description.\r
+ if ( !description.empty())\r
+ {\r
+ str << "SYSTEM\\CurrentControlSet\\Services\\" << getServiceName();\r
+\r
+ if ( RegCreateKey( HKEY_LOCAL_MACHINE, str.str().c_str(), &key ) != ERROR_SUCCESS )\r
+ return false;\r
+\r
+ if ( RegSetValueEx( key, "Description", 0, REG_SZ, reinterpret_cast< const unsigned char * >( description.c_str()), description.length()) != ERROR_SUCCESS )\r
+ {\r
+ RegCloseKey( key );\r
+ return false;\r
+ }\r
+\r
+ RegCloseKey( key );\r
+ }\r
+\r
+ // Register event log.\r
+ if ( registerLog )\r
+ {\r
+ str.str( "" );\r
+\r
+ str << "SYSTEM\\CurrentControlSet\\Services\\Eventlog\\Application\\" << getServiceName();\r
+ if ( RegCreateKey( HKEY_LOCAL_MACHINE, str.str().c_str(), &pkey ) != ERROR_SUCCESS )\r
+ return false;\r
+\r
+ flags = EVENTLOG_ERROR_TYPE | EVENTLOG_WARNING_TYPE | EVENTLOG_INFORMATION_TYPE;\r
+ \r
+ if ( RegSetValueEx( pkey, "TypesSupported", 0, REG_DWORD, reinterpret_cast< const unsigned char * >( &flags ), sizeof( flags )) != ERROR_SUCCESS )\r
+ {\r
+ RegCloseKey( pkey );\r
+ return false;\r
+ }\r
+\r
+ // For the message file this function assumes %SystemRoot%\\System32\\<servicename>msg.dll\r
+ str.str( "" );\r
+\r
+ char path[ MAX_PATH ];\r
+ GetCurrentDirectory( sizeof( path ), path );\r
+\r
+ str << path << "\\" << getServiceName() << "msg.dll";\r
+ if ( RegSetValueEx( pkey, "EventMessageFile", 0, REG_SZ, reinterpret_cast< const unsigned char * >( str.str().c_str()), str.str().length()) != ERROR_SUCCESS )\r
+ {\r
+ RegCloseKey( pkey );\r
+ return false;\r
+ }\r
+\r
+ RegCloseKey( pkey );\r
+ }\r
+ \r
+ return true;\r
+}\r
+\r
+\r
+// Calls the control handler.\r
+void WINAPI NTService::s_ctrlHandler( DWORD controlCode )\r
+{\r
+ NTService::instance()->ctrlHandler( controlCode );\r
+}\r
+\r
+\r
+// Calls the service's main function.\r
+void WINAPI NTService::s_serviceMain( DWORD argc, LPTSTR *argv )\r
+{\r
+ // IEEEEUUWWWW!!\r
+ \r
+ NTService::instance()->m_serviceStatusHandle = RegisterServiceCtrlHandler( NTService::instance()->getServiceName().c_str(), s_ctrlHandler );\r
+ if ( NTService::instance()->m_serviceStatusHandle == 0 )\r
+ {\r
+ // Could not register service ctrl handler.\r
+ return;\r
+ }\r
+\r
+ NTService::instance()->setStatus( SERVICE_START_PENDING );\r
+ // Initialize.\r
+ if ( !NTService::instance()->init())\r
+ {\r
+ NTService::instance()->setStatus( SERVICE_STOPPED, -1 );\r
+ return;\r
+ }\r
+\r
+ NTService::instance()->setStatus( SERVICE_RUNNING );\r
+ \r
+ // Run.\r
+ NTService::instance()->main( argc, argv );\r
+\r
+ NTService::instance()->setStatus( SERVICE_STOP_PENDING );\r
+\r
+ // Shut down.\r
+ NTService::instance()->shutdown();\r
+ \r
+ NTService::instance()->setStatus( SERVICE_STOPPED );\r
+}\r
+\r
+\r
+// Sets the service's status.\r
+void NTService::setStatus( DWORD status, DWORD error )\r
+{\r
+ SERVICE_STATUS stat;\r
+\r
+ if ( !m_serviceStatusHandle )\r
+ return;\r
+\r
+ stat.dwServiceType = SERVICE_WIN32_OWN_PROCESS;\r
+ stat.dwCurrentState = status;\r
+ stat.dwControlsAccepted = SERVICE_ACCEPT_STOP | SERVICE_ACCEPT_SHUTDOWN;\r
+ stat.dwWin32ExitCode = ( error ? ERROR_SERVICE_SPECIFIC_ERROR : NO_ERROR );\r
+ stat.dwServiceSpecificExitCode = error;\r
+ stat.dwCheckPoint = 0;\r
+ stat.dwWaitHint = 0;\r
+\r
+ SetServiceStatus( m_serviceStatusHandle, &stat ); \r
+}\r
+\r
+\r
+// Starts the service.\r
+int NTService::start( int argc, char *argv[], bool asService )\r
+{\r
+ int res = 0;\r
+ char name[ 128 ];\r
+\r
+ strncpy( name, getServiceName().c_str(), sizeof( name ));\r
+ \r
+ SERVICE_TABLE_ENTRY entries[] =\r
+ {\r
+ { name, s_serviceMain },\r
+ { NULL, NULL }\r
+ };\r
+ \r
+ if ( asService )\r
+ {\r
+ // Run as service.\r
+ m_runningAsService = true;\r
+\r
+ if ( StartServiceCtrlDispatcher( entries ))\r
+ return 0; // Success!\r
+\r
+ // StartServiceCtrlDispatcher() failed, check if we should run as a normal\r
+ // console program.\r
+ if ( GetLastError() != ERROR_FAILED_SERVICE_CONTROLLER_CONNECT )\r
+ return -1;\r
+\r
+ }\r
+\r
+ DLOG( L << "Running as a normal (console) program." << endl );\r
+\r
+ // Run as normal (console) program.\r
+ m_runningAsService = false;\r
+\r
+ if ( !init())\r
+ return -1; // Could not initialize.\r
+\r
+ // Run.\r
+ res = main( argc, argv );\r
+\r
+ shutdown();\r
+\r
+ return res;\r
+}\r
+\r
+\r
+// Stops the service.\r
+bool NTService::stop( void )\r
+{\r
+ if ( !isRunningAsService())\r
+ exit( 0 );\r
+\r
+ setStatus( SERVICE_STOPPED, 0 );\r
+\r
+ return true;\r
+}\r
+\r
+\r
+// Unregister service.\r
+bool NTService::unregisterService( void )\r
+{\r
+ HKEY key;\r
+ SC_HANDLE sc, svc;\r
+\r
+ sc = OpenSCManager( NULL, NULL, SC_MANAGER_ALL_ACCESS );\r
+ if ( sc == NULL )\r
+ return false;\r
+ \r
+ svc = OpenService( sc, getServiceName().c_str(), DELETE );\r
+ if ( GetLastError() != ERROR_SERVICE_DOES_NOT_EXIST )\r
+ {\r
+ if ( svc == NULL )\r
+ {\r
+ CloseServiceHandle( sc );\r
+ return false;\r
+ }\r
+\r
+ DeleteService( svc );\r
+ CloseServiceHandle( svc );\r
+ CloseServiceHandle( sc );\r
+ }\r
+\r
+ if ( RegOpenKeyEx( HKEY_LOCAL_MACHINE, "SYSTEM\\CurrentControlSet\\Services\\EventLog\\Application", 0, KEY_WRITE, &key ) != ERROR_SUCCESS )\r
+ return false;\r
+\r
+ RegDeleteKey( key, getServiceName().c_str());\r
+\r
+ RegCloseKey( key );\r
+ \r
+ return true;\r
+}\r