michael@0: /* This Source Code Form is subject to the terms of the Mozilla Public
michael@0:  * License, v. 2.0. If a copy of the MPL was not distributed with this
michael@0:  * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
michael@0: 
michael@0: #include <windows.h>
michael@0: #include <shlwapi.h>
michael@0: #include <stdio.h>
michael@0: #include <wchar.h>
michael@0: #include <shlobj.h>
michael@0: 
michael@0: #include "serviceinstall.h"
michael@0: #include "maintenanceservice.h"
michael@0: #include "servicebase.h"
michael@0: #include "workmonitor.h"
michael@0: #include "uachelper.h"
michael@0: #include "updatehelper.h"
michael@0: 
michael@0: SERVICE_STATUS gSvcStatus = { 0 }; 
michael@0: SERVICE_STATUS_HANDLE gSvcStatusHandle = nullptr; 
michael@0: HANDLE gWorkDoneEvent = nullptr;
michael@0: HANDLE gThread = nullptr;
michael@0: bool gServiceControlStopping = false;
michael@0: 
michael@0: // logs are pretty small, about 20 lines, so 10 seems reasonable.
michael@0: #define LOGS_TO_KEEP 10
michael@0: 
michael@0: BOOL GetLogDirectoryPath(WCHAR *path);
michael@0: 
michael@0: int 
michael@0: wmain(int argc, WCHAR **argv)
michael@0: {
michael@0:   // If command-line parameter is "install", install the service
michael@0:   // or upgrade if already installed
michael@0:   // If command line parameter is "forceinstall", install the service
michael@0:   // even if it is older than what is already installed.
michael@0:   // If command-line parameter is "upgrade", upgrade the service
michael@0:   // but do not install it if it is not already installed.
michael@0:   // If command line parameter is "uninstall", uninstall the service.
michael@0:   // Otherwise, the service is probably being started by the SCM.
michael@0:   bool forceInstall = !lstrcmpi(argv[1], L"forceinstall");
michael@0:   if (!lstrcmpi(argv[1], L"install") || forceInstall) {
michael@0:     WCHAR updatePath[MAX_PATH + 1];
michael@0:     if (GetLogDirectoryPath(updatePath)) {
michael@0:       LogInit(updatePath, L"maintenanceservice-install.log");
michael@0:     }
michael@0: 
michael@0:     SvcInstallAction action = InstallSvc;
michael@0:     if (forceInstall) {
michael@0:       action = ForceInstallSvc;
michael@0:       LOG(("Installing service with force specified..."));
michael@0:     } else {
michael@0:       LOG(("Installing service..."));
michael@0:     }
michael@0: 
michael@0:     bool ret = SvcInstall(action);
michael@0:     if (!ret) {
michael@0:       LOG_WARN(("Could not install service.  (%d)", GetLastError()));
michael@0:       LogFinish();
michael@0:       return 1;
michael@0:     }
michael@0: 
michael@0:     LOG(("The service was installed successfully"));
michael@0:     LogFinish();
michael@0:     return 0;
michael@0:   } 
michael@0: 
michael@0:   if (!lstrcmpi(argv[1], L"upgrade")) {
michael@0:     WCHAR updatePath[MAX_PATH + 1];
michael@0:     if (GetLogDirectoryPath(updatePath)) {
michael@0:       LogInit(updatePath, L"maintenanceservice-install.log");
michael@0:     }
michael@0: 
michael@0:     LOG(("Upgrading service if installed..."));
michael@0:     if (!SvcInstall(UpgradeSvc)) {
michael@0:       LOG_WARN(("Could not upgrade service.  (%d)", GetLastError()));
michael@0:       LogFinish();
michael@0:       return 1;
michael@0:     }
michael@0: 
michael@0:     LOG(("The service was upgraded successfully"));
michael@0:     LogFinish();
michael@0:     return 0;
michael@0:   }
michael@0: 
michael@0:   if (!lstrcmpi(argv[1], L"uninstall")) {
michael@0:     WCHAR updatePath[MAX_PATH + 1];
michael@0:     if (GetLogDirectoryPath(updatePath)) {
michael@0:       LogInit(updatePath, L"maintenanceservice-uninstall.log");
michael@0:     }
michael@0:     LOG(("Uninstalling service..."));
michael@0:     if (!SvcUninstall()) {
michael@0:       LOG_WARN(("Could not uninstall service.  (%d)", GetLastError()));
michael@0:       LogFinish();
michael@0:       return 1;
michael@0:     }
michael@0:     LOG(("The service was uninstalled successfully"));
michael@0:     LogFinish();
michael@0:     return 0;
michael@0:   }
michael@0: 
michael@0:   SERVICE_TABLE_ENTRYW DispatchTable[] = { 
michael@0:     { SVC_NAME, (LPSERVICE_MAIN_FUNCTIONW) SvcMain }, 
michael@0:     { nullptr, nullptr } 
michael@0:   }; 
michael@0: 
michael@0:   // This call returns when the service has stopped. 
michael@0:   // The process should simply terminate when the call returns.
michael@0:   if (!StartServiceCtrlDispatcherW(DispatchTable)) {
michael@0:     LOG_WARN(("StartServiceCtrlDispatcher failed.  (%d)", GetLastError()));
michael@0:   }
michael@0: 
michael@0:   return 0;
michael@0: }
michael@0: 
michael@0: /**
michael@0:  * Obtains the base path where logs should be stored
michael@0:  *
michael@0:  * @param  path The out buffer for the backup log path of size MAX_PATH + 1
michael@0:  * @return TRUE if successful.
michael@0:  */
michael@0: BOOL
michael@0: GetLogDirectoryPath(WCHAR *path)
michael@0: {
michael@0:   HRESULT hr = SHGetFolderPathW(nullptr, CSIDL_COMMON_APPDATA, nullptr, 
michael@0:     SHGFP_TYPE_CURRENT, path);
michael@0:   if (FAILED(hr)) {
michael@0:     return FALSE;
michael@0:   }
michael@0: 
michael@0:   if (!PathAppendSafe(path, L"Mozilla")) {
michael@0:     return FALSE;
michael@0:   }
michael@0:   // The directory should already be created from the installer, but
michael@0:   // just to be safe in case someone deletes.
michael@0:   CreateDirectoryW(path, nullptr);
michael@0: 
michael@0:   if (!PathAppendSafe(path, L"logs")) {
michael@0:     return FALSE;
michael@0:   }
michael@0:   CreateDirectoryW(path, nullptr);
michael@0:   return TRUE;
michael@0: }
michael@0: 
michael@0: /**
michael@0:  * Calculated a backup path based on the log number.
michael@0:  *
michael@0:  * @param  path      The out buffer to store the log path of size MAX_PATH + 1
michael@0:  * @param  basePath  The base directory where the calculated path should go
michael@0:  * @param  logNumber The log number, 0 == updater.log
michael@0:  * @return TRUE if successful.
michael@0:  */
michael@0: BOOL
michael@0: GetBackupLogPath(LPWSTR path, LPCWSTR basePath, int logNumber)
michael@0: {
michael@0:   WCHAR logName[64] = { L'\0' };
michael@0:   wcsncpy(path, basePath, sizeof(logName) / sizeof(logName[0]) - 1);
michael@0:   if (logNumber <= 0) {
michael@0:     swprintf(logName, sizeof(logName) / sizeof(logName[0]),
michael@0:              L"maintenanceservice.log");
michael@0:   } else {
michael@0:     swprintf(logName, sizeof(logName) / sizeof(logName[0]),
michael@0:              L"maintenanceservice-%d.log", logNumber);
michael@0:   }
michael@0:   return PathAppendSafe(path, logName);
michael@0: }
michael@0: 
michael@0: /**
michael@0:  * Moves the old log files out of the way before a new one is written.
michael@0:  * If you for example keep 3 logs, then this function will do:
michael@0:  *   updater2.log -> updater3.log
michael@0:  *   updater1.log -> updater2.log
michael@0:  *   updater.log -> updater1.log
michael@0:  * Which clears room for a new updater.log in the basePath directory
michael@0:  * 
michael@0:  * @param basePath      The base directory path where log files are stored
michael@0:  * @param numLogsToKeep The number of logs to keep
michael@0:  */
michael@0: void
michael@0: BackupOldLogs(LPCWSTR basePath, int numLogsToKeep)
michael@0: {
michael@0:   WCHAR oldPath[MAX_PATH + 1];
michael@0:   WCHAR newPath[MAX_PATH + 1];
michael@0:   for (int i = numLogsToKeep; i >= 1; i--) {
michael@0:     if (!GetBackupLogPath(oldPath, basePath, i -1)) {
michael@0:       continue;
michael@0:     }
michael@0: 
michael@0:     if (!GetBackupLogPath(newPath, basePath, i)) {
michael@0:       continue;
michael@0:     }
michael@0: 
michael@0:     if (!MoveFileExW(oldPath, newPath, MOVEFILE_REPLACE_EXISTING)) {
michael@0:       continue;
michael@0:     }
michael@0:   }
michael@0: }
michael@0: 
michael@0: /**
michael@0:  * Ensures the service is shutdown once all work is complete.
michael@0:  * There is an issue on XP SP2 and below where the service can hang
michael@0:  * in a stop pending state even though the SCM is notified of a stopped
michael@0:  * state.  Control *should* be returned to StartServiceCtrlDispatcher from the 
michael@0:  * call to SetServiceStatus on a stopped state in the wmain thread.  
michael@0:  * Sometimes this is not the case though. This thread will terminate the process
michael@0:  * if it has been 5 seconds after all work is done and the process is still not
michael@0:  * terminated.  This thread is only started once a stopped state was sent to the
michael@0:  * SCM. The stop pending hang can be reproduced intermittently even if you set 
michael@0:  * a stopped state dirctly and never set a stop pending state.  It is safe to 
michael@0:  * forcefully terminate the process ourselves since all work is done once we 
michael@0:  * start this thread.
michael@0: */
michael@0: DWORD WINAPI
michael@0: EnsureProcessTerminatedThread(LPVOID)
michael@0: {
michael@0:   Sleep(5000);
michael@0:   exit(0);
michael@0:   return 0;
michael@0: }
michael@0: 
michael@0: void
michael@0: StartTerminationThread()
michael@0: {
michael@0:   // If the process does not self terminate like it should, this thread 
michael@0:   // will terminate the process after 5 seconds.
michael@0:   HANDLE thread = CreateThread(nullptr, 0, EnsureProcessTerminatedThread,
michael@0:                                nullptr, 0, nullptr);
michael@0:   if (thread) {
michael@0:     CloseHandle(thread);
michael@0:   }
michael@0: }
michael@0: 
michael@0: /**
michael@0:  * Main entry point when running as a service.
michael@0:  */
michael@0: void WINAPI
michael@0: SvcMain(DWORD argc, LPWSTR *argv)
michael@0: {
michael@0:   // Setup logging, and backup the old logs
michael@0:   WCHAR updatePath[MAX_PATH + 1];
michael@0:   if (GetLogDirectoryPath(updatePath)) {
michael@0:     BackupOldLogs(updatePath, LOGS_TO_KEEP);
michael@0:     LogInit(updatePath, L"maintenanceservice.log");
michael@0:   }
michael@0: 
michael@0:   // Disable every privilege we don't need. Processes started using
michael@0:   // CreateProcess will use the same token as this process.
michael@0:   UACHelper::DisablePrivileges(nullptr);
michael@0: 
michael@0:   // Register the handler function for the service
michael@0:   gSvcStatusHandle = RegisterServiceCtrlHandlerW(SVC_NAME, SvcCtrlHandler);
michael@0:   if (!gSvcStatusHandle) {
michael@0:     LOG_WARN(("RegisterServiceCtrlHandler failed.  (%d)", GetLastError()));
michael@0:     ExecuteServiceCommand(argc, argv);  
michael@0:     LogFinish();
michael@0:     exit(1);
michael@0:   } 
michael@0: 
michael@0:   // These values will be re-used later in calls involving gSvcStatus
michael@0:   gSvcStatus.dwServiceType = SERVICE_WIN32_OWN_PROCESS;
michael@0:   gSvcStatus.dwServiceSpecificExitCode = 0;
michael@0: 
michael@0:   // Report initial status to the SCM
michael@0:   ReportSvcStatus(SERVICE_START_PENDING, NO_ERROR, 3000);
michael@0: 
michael@0:   // This event will be used to tell the SvcCtrlHandler when the work is
michael@0:   // done for when a stop comamnd is manually issued.
michael@0:   gWorkDoneEvent = CreateEvent(nullptr, TRUE, FALSE, nullptr);
michael@0:   if (!gWorkDoneEvent) {
michael@0:     ReportSvcStatus(SERVICE_STOPPED, 1, 0);
michael@0:     StartTerminationThread();
michael@0:     return;
michael@0:   }
michael@0: 
michael@0:   // Initialization complete and we're about to start working on
michael@0:   // the actual command.  Report the service state as running to the SCM.
michael@0:   ReportSvcStatus(SERVICE_RUNNING, NO_ERROR, 0);
michael@0: 
michael@0:   // The service command was executed, stop logging and set an event
michael@0:   // to indicate the work is done in case someone is waiting on a
michael@0:   // service stop operation.
michael@0:   ExecuteServiceCommand(argc, argv);  
michael@0:   LogFinish();
michael@0: 
michael@0:   SetEvent(gWorkDoneEvent);
michael@0: 
michael@0:   // If we aren't already in a stopping state then tell the SCM we're stopped
michael@0:   // now.  If we are already in a stopping state then the SERVICE_STOPPED state
michael@0:   // will be set by the SvcCtrlHandler.
michael@0:   if (!gServiceControlStopping) {
michael@0:     ReportSvcStatus(SERVICE_STOPPED, NO_ERROR, 0);
michael@0:     StartTerminationThread();
michael@0:   }
michael@0: }
michael@0: 
michael@0: /**
michael@0:  * Sets the current service status and reports it to the SCM.
michael@0:  *  
michael@0:  * @param currentState  The current state (see SERVICE_STATUS)
michael@0:  * @param exitCode      The system error code
michael@0:  * @param waitHint      Estimated time for pending operation in milliseconds
michael@0:  */
michael@0: void
michael@0: ReportSvcStatus(DWORD currentState, 
michael@0:                 DWORD exitCode, 
michael@0:                 DWORD waitHint)
michael@0: {
michael@0:   static DWORD dwCheckPoint = 1;
michael@0: 
michael@0:   gSvcStatus.dwCurrentState = currentState;
michael@0:   gSvcStatus.dwWin32ExitCode = exitCode;
michael@0:   gSvcStatus.dwWaitHint = waitHint;
michael@0: 
michael@0:   if (SERVICE_START_PENDING == currentState || 
michael@0:       SERVICE_STOP_PENDING == currentState) {
michael@0:     gSvcStatus.dwControlsAccepted = 0;
michael@0:   } else {
michael@0:     gSvcStatus.dwControlsAccepted = SERVICE_ACCEPT_STOP | 
michael@0:                                     SERVICE_ACCEPT_SHUTDOWN;
michael@0:   }
michael@0: 
michael@0:   if ((SERVICE_RUNNING == currentState) ||
michael@0:       (SERVICE_STOPPED == currentState)) {
michael@0:     gSvcStatus.dwCheckPoint = 0;
michael@0:   } else {
michael@0:     gSvcStatus.dwCheckPoint = dwCheckPoint++;
michael@0:   }
michael@0: 
michael@0:   // Report the status of the service to the SCM.
michael@0:   SetServiceStatus(gSvcStatusHandle, &gSvcStatus);
michael@0: }
michael@0: 
michael@0: /**
michael@0:  * Since the SvcCtrlHandler should only spend at most 30 seconds before 
michael@0:  * returning, this function does the service stop work for the SvcCtrlHandler. 
michael@0: */
michael@0: DWORD WINAPI
michael@0: StopServiceAndWaitForCommandThread(LPVOID)
michael@0: {
michael@0:   do {
michael@0:     ReportSvcStatus(SERVICE_STOP_PENDING, NO_ERROR, 1000);
michael@0:   } while(WaitForSingleObject(gWorkDoneEvent, 100) == WAIT_TIMEOUT);
michael@0:   CloseHandle(gWorkDoneEvent);
michael@0:   gWorkDoneEvent = nullptr;
michael@0:   ReportSvcStatus(SERVICE_STOPPED, NO_ERROR, 0);
michael@0:   StartTerminationThread();
michael@0:   return 0;
michael@0: }
michael@0: 
michael@0: /**
michael@0:  * Called by SCM whenever a control code is sent to the service
michael@0:  * using the ControlService function.
michael@0:  */
michael@0: void WINAPI
michael@0: SvcCtrlHandler(DWORD dwCtrl)
michael@0: {
michael@0:   // After a SERVICE_CONTROL_STOP there should be no more commands sent to
michael@0:   // the SvcCtrlHandler.
michael@0:   if (gServiceControlStopping) {
michael@0:     return;
michael@0:   }
michael@0: 
michael@0:   // Handle the requested control code. 
michael@0:   switch(dwCtrl) {
michael@0:   case SERVICE_CONTROL_SHUTDOWN:
michael@0:   case SERVICE_CONTROL_STOP: {
michael@0:       gServiceControlStopping = true;
michael@0:       ReportSvcStatus(SERVICE_STOP_PENDING, NO_ERROR, 1000);
michael@0: 
michael@0:       // The SvcCtrlHandler thread should not spend more than 30 seconds in 
michael@0:       // shutdown so we spawn a new thread for stopping the service 
michael@0:       HANDLE thread = CreateThread(nullptr, 0,
michael@0:                                    StopServiceAndWaitForCommandThread,
michael@0:                                    nullptr, 0, nullptr);
michael@0:       if (thread) {
michael@0:         CloseHandle(thread);
michael@0:       } else {
michael@0:         // Couldn't start the thread so just call the stop ourselves.
michael@0:         // If it happens to take longer than 30 seconds the caller will
michael@0:         // get an error.
michael@0:         StopServiceAndWaitForCommandThread(nullptr);
michael@0:       }
michael@0:     }
michael@0:     break;
michael@0:   default:
michael@0:     break;
michael@0:   }
michael@0: }