pdns/ntservice.cc

   1 /*
   2     PowerDNS Versatile Database Driven Nameserver
   3     Copyright (C) 2002  PowerDNS.COM BV
   4
   5     This program is free software; you can redistribute it and/or modify
   6     it under the terms of the GNU General Public License version 2
   7     as published by the Free Software Foundation
   8
   9
  10     This program is distributed in the hope that it will be useful,
  11     but WITHOUT ANY WARRANTY; without even the implied warranty of
  12     MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  13     GNU General Public License for more details.
  14
  15     You should have received a copy of the GNU General Public License
  16     along with this program; if not, write to the Free Software
  17     Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
  18 */
  19
  20 /*!
  21 \file ntservice.cpp
  22 \brief This file contains the NTService class implementation.
  23 */
  24
  25 #include "utility.hh"
  26 #include <sstream>
  27 #include <iostream>
  28 #include "logger.hh"
  29 #include "ntservice.hh"
  30
  31 #define L theL("pdns")
  32
  33
  34 // Default constructor.
  35 NTService::NTService( void )
  36 {
  37   m_runningAsService    = false;
  38   m_errorCode           = 0;
  39   m_statusCode          = 0;
  40   m_serviceStatusHandle = NULL;
  41 }
  42
  43
  44 // Destructor.
  45 NTService::~NTService( void )
  46 {
  47 }
  48
  49
  50 // Returns whether the program is running as a service.
  51 bool NTService::isRunningAsService( void )
  52 {
  53   return m_runningAsService;
  54 }
  55
  56
  57 // Registers the service.
  58 bool NTService::registerService( const std::string & description, bool registerLog )
  59 {
  60   std::stringstream str;
  61   HKEY              key, pkey;
  62   SC_HANDLE         sc;
  63   char              temp[ 512 ];
  64   DWORD             flags;
  65
  66   sc = OpenSCManager( NULL, NULL, SC_MANAGER_ALL_ACCESS );
  67   if ( sc == NULL )
  68     return false; // Could not open the Service Control Manager.
  69
  70   GetModuleFileName( NULL, temp, sizeof( temp ));
  71
  72   str << temp << " --ntservice";
  73
  74   if ( CreateService(
  75     sc,
  76     getServiceName().c_str(),
  77     getServiceName().c_str(),
  78     SERVICE_ALL_ACCESS,
  79     SERVICE_WIN32_OWN_PROCESS,
  80     SERVICE_AUTO_START,
  81     SERVICE_ERROR_NORMAL,
  82     str.str().c_str(),
  83     NULL,
  84     NULL,
  85     NULL,
  86     NULL,
  87     NULL ) == NULL  )
  88   {
  89         if(GetLastError() != ERROR_SERVICE_EXISTS)
  90       return false; // Don't we all like functions with 43 billion parameters?
  91   }
  92   CloseServiceHandle( sc );
  93
  94   str.str( "" );
  95
  96   // Set description.
  97   if ( !description.empty())
  98   {
  99     str << "SYSTEM\\CurrentControlSet\\Services\\" << getServiceName();
 100
 101     if ( RegCreateKey( HKEY_LOCAL_MACHINE, str.str().c_str(), &key ) != ERROR_SUCCESS )
 102       return false;
 103
 104     if ( RegSetValueEx( key, "Description", 0, REG_SZ, reinterpret_cast< const unsigned char * >( description.c_str()), description.length()) != ERROR_SUCCESS )
 105     {
 106       RegCloseKey( key );
 107       return false;
 108     }
 109
 110     RegCloseKey( key );
 111   }
 112
 113   // Register event log.
 114   if ( registerLog )
 115   {
 116     str.str( "" );
 117
 118     str << "SYSTEM\\CurrentControlSet\\Services\\Eventlog\\Application\\" << getServiceName();
 119     if ( RegCreateKey( HKEY_LOCAL_MACHINE, str.str().c_str(), &pkey ) != ERROR_SUCCESS )
 120       return false;
 121
 122     flags = EVENTLOG_ERROR_TYPE | EVENTLOG_WARNING_TYPE | EVENTLOG_INFORMATION_TYPE;
 123
 124     if ( RegSetValueEx( pkey, "TypesSupported", 0, REG_DWORD, reinterpret_cast< const unsigned char * >( &flags ), sizeof( flags )) != ERROR_SUCCESS )
 125     {
 126       RegCloseKey( pkey );
 127       return false;
 128     }
 129
 130     // For the message file this function assumes %SystemRoot%\\System32\\<servicename>msg.dll
 131     str.str( "" );
 132
 133     char path[ MAX_PATH ];
 134     GetCurrentDirectory( sizeof( path ), path );
 135
 136     // FIXME: This really should be: str << path << "\\" << getServiceName() << "msg.dll";
 137     str << path << "\\Eventlogger.dll";
 138     if ( RegSetValueEx( pkey, "EventMessageFile", 0, REG_SZ, reinterpret_cast< const unsigned char * >( str.str().c_str()), str.str().length()) != ERROR_SUCCESS )
 139     {
 140       RegCloseKey( pkey );
 141       return false;
 142     }
 143
 144     RegCloseKey( pkey );
 145   }
 146
 147   return true;
 148 }
 149
 150
 151 // Calls the control handler.
 152 void WINAPI NTService::s_ctrlHandler( DWORD controlCode )
 153 {
 154   NTService::instance()->ctrlHandler( controlCode );
 155 }
 156
 157
 158 // Calls the service's main function.
 159 void WINAPI NTService::s_serviceMain( DWORD argc, LPTSTR *argv )
 160 {
 161   // IEEEEUUWWWW!!
 162
 163   NTService::instance()->m_serviceStatusHandle = RegisterServiceCtrlHandler( NTService::instance()->getServiceName().c_str(), s_ctrlHandler );
 164   if ( NTService::instance()->m_serviceStatusHandle == 0 )
 165   {
 166     // Could not register service ctrl handler.
 167     return;
 168   }
 169
 170   NTService::instance()->setStatus( SERVICE_START_PENDING );
 171   // Initialize.
 172   if ( !NTService::instance()->init())
 173   {
 174     NTService::instance()->setStatus( SERVICE_STOPPED, -1 );
 175     return;
 176   }
 177
 178   NTService::instance()->setStatus( SERVICE_RUNNING );
 179
 180   // Run.
 181   NTService::instance()->main( argc, argv );
 182
 183   NTService::instance()->setStatus( SERVICE_STOP_PENDING );
 184
 185   // Shut down.
 186   NTService::instance()->shutdown();
 187
 188   NTService::instance()->setStatus( SERVICE_STOPPED );
 189 }
 190
 191
 192 // Sets the service's status.
 193 void NTService::setStatus( DWORD status, DWORD error )
 194 {
 195   SERVICE_STATUS stat;
 196
 197   if ( !m_serviceStatusHandle )
 198     return;
 199
 200   stat.dwServiceType              = SERVICE_WIN32_OWN_PROCESS;
 201   stat.dwCurrentState             = status;
 202   stat.dwControlsAccepted         = SERVICE_ACCEPT_STOP | SERVICE_ACCEPT_SHUTDOWN;
 203   stat.dwWin32ExitCode            = ( error ? ERROR_SERVICE_SPECIFIC_ERROR : NO_ERROR );
 204   stat.dwServiceSpecificExitCode  = error;
 205   stat.dwCheckPoint               = 0;
 206   stat.dwWaitHint                 = 0;
 207
 208   SetServiceStatus( m_serviceStatusHandle, &stat );
 209 }
 210
 211
 212 // Starts the service.
 213 int NTService::start( int argc, char *argv[], bool asService )
 214 {
 215   int res = 0;
 216   char name[ 128 ];
 217
 218   strncpy( name, getServiceName().c_str(), sizeof( name ));
 219
 220   SERVICE_TABLE_ENTRY entries[] =
 221   {
 222     { name, s_serviceMain },
 223     { NULL, NULL }
 224   };
 225
 226   if ( asService )
 227   {
 228     // Run as service.
 229     m_runningAsService = true;
 230
 231     if ( StartServiceCtrlDispatcher( entries ))
 232       return 0; // Success!
 233
 234     // StartServiceCtrlDispatcher() failed, check if we should run as a normal
 235     // console program.
 236     if ( GetLastError() != ERROR_FAILED_SERVICE_CONTROLLER_CONNECT )
 237       return -1;
 238
 239   }
 240
 241   DLOG( L << "Running as a normal (console) program." << endl );
 242
 243   // Run as normal (console) program.
 244   m_runningAsService = false;
 245
 246   if ( !init())
 247     return -1; // Could not initialize.
 248
 249   // Run.
 250   res = main( argc, argv );
 251
 252   shutdown();
 253
 254   return res;
 255 }
 256
 257
 258 // Stops the service.
 259 bool NTService::stop( void )
 260 {
 261   if ( !isRunningAsService())
 262     exit( 0 );
 263
 264   setStatus( SERVICE_STOPPED, 0 );
 265
 266   return true;
 267 }
 268
 269
 270 // Unregister service.
 271 bool NTService::unregisterService( void )
 272 {
 273   HKEY      key;
 274   SC_HANDLE sc, svc;
 275
 276   sc = OpenSCManager( NULL, NULL, SC_MANAGER_ALL_ACCESS );
 277   if ( sc == NULL )
 278     return false;
 279
 280   svc = OpenService( sc, getServiceName().c_str(), DELETE );
 281   if ( GetLastError() != ERROR_SERVICE_DOES_NOT_EXIST )
 282   {
 283     if ( svc == NULL )
 284     {
 285       CloseServiceHandle( sc );
 286       return false;
 287     }
 288
 289     DeleteService( svc );
 290     CloseServiceHandle( svc );
 291     CloseServiceHandle( sc );
 292   }
 293
 294   if ( RegOpenKeyEx( HKEY_LOCAL_MACHINE, "SYSTEM\\CurrentControlSet\\Services\\EventLog\\Application", 0, KEY_WRITE, &key ) != ERROR_SUCCESS )
 295     return false;
 296
 297   RegDeleteKey( key, getServiceName().c_str());
 298
 299   RegCloseKey( key );
 300
 301   return true;
 302 }