/mordor/daemon.cpp

  1// Copyright (c) 2010 - Mozy, Inc.
  2
  3#include "daemon.h"
  4
  5#include "log.h"
  6#include "main.h"
  7
  8namespace Mordor {
  9namespace Daemon {
 10
 11static Logger::ptr g_log = Log::lookup("mordor:daemon");
 12
 13boost::signals2::signal<void ()> onTerminate;
 14boost::signals2::signal<void ()> onInterrupt;
 15boost::signals2::signal<void ()> onReload;
 16boost::signals2::signal<void ()> onPause;
 17boost::signals2::signal<void ()> onContinue;
 18
 19#define CTRL_CASE(ctrl)             \
 20    case ctrl:                      \
 21        return os << #ctrl;         \
 22
 23#ifdef WINDOWS
 24static boost::function<int (int, char **)> g_daemonMain;
 25
 26namespace {
 27struct ServiceStatus
 28{
 29    SERVICE_STATUS_HANDLE serviceHandle;
 30    SERVICE_STATUS status;
 31};
 32}
 33
 34static DWORD acceptedControls()
 35{
 36    DWORD accepted = 0;
 37    if (!onTerminate.empty())
 38        accepted |= SERVICE_ACCEPT_STOP;
 39    if (!onReload.empty())
 40        accepted |= SERVICE_ACCEPT_PARAMCHANGE;
 41    if (!onPause.empty())
 42        accepted |= SERVICE_ACCEPT_PAUSE_CONTINUE;
 43    return accepted;
 44}
 45
 46namespace {
 47
 48enum ServiceCtrl {};
 49
 50std::ostream &operator <<(std::ostream &os, ServiceCtrl ctrl)
 51{
 52    switch (ctrl) {
 53        CTRL_CASE(SERVICE_CONTROL_STOP);
 54        CTRL_CASE(SERVICE_CONTROL_PAUSE);
 55        CTRL_CASE(SERVICE_CONTROL_CONTINUE);
 56        CTRL_CASE(SERVICE_CONTROL_INTERROGATE);
 57        CTRL_CASE(SERVICE_CONTROL_SHUTDOWN);
 58        CTRL_CASE(SERVICE_CONTROL_PARAMCHANGE);
 59        CTRL_CASE(SERVICE_CONTROL_NETBINDADD);
 60        CTRL_CASE(SERVICE_CONTROL_NETBINDREMOVE);
 61        CTRL_CASE(SERVICE_CONTROL_NETBINDENABLE);
 62        CTRL_CASE(SERVICE_CONTROL_NETBINDDISABLE);
 63        CTRL_CASE(SERVICE_CONTROL_DEVICEEVENT);
 64        CTRL_CASE(SERVICE_CONTROL_HARDWAREPROFILECHANGE);
 65        CTRL_CASE(SERVICE_CONTROL_POWEREVENT);
 66        CTRL_CASE(SERVICE_CONTROL_SESSIONCHANGE);
 67        CTRL_CASE(SERVICE_CONTROL_PRESHUTDOWN);
 68#ifdef SERVICE_CONTROL_TIMECHANGE
 69        CTRL_CASE(SERVICE_CONTROL_TIMECHANGE);
 70#endif
 71#ifdef SERVICE_CONTROL_TRIGGEREVENT
 72        CTRL_CASE(SERVICE_CONTROL_TRIGGEREVENT);
 73#endif
 74        default:
 75            return os << (int)ctrl;
 76    }
 77}
 78
 79}
 80
 81static DWORD WINAPI HandlerEx(DWORD dwControl, DWORD dwEventType,
 82    LPVOID lpEventData, LPVOID lpContext)
 83{
 84    ServiceStatus *status = (ServiceStatus *)lpContext;
 85    MORDOR_LOG_DEBUG(g_log) << "Received service control "
 86        << (ServiceCtrl)dwControl;
 87
 88    switch (dwControl) {
 89        case SERVICE_CONTROL_INTERROGATE:
 90            break;
 91        case SERVICE_CONTROL_PARAMCHANGE:
 92            onReload();
 93            break;
 94        case SERVICE_CONTROL_STOP:
 95            if (onTerminate.empty())
 96                ExitProcess(0);
 97            onTerminate();
 98            break;
 99        case SERVICE_CONTROL_PAUSE:
100            if (!onPause.empty())
101                status->status.dwCurrentState = SERVICE_PAUSED;
102            onPause();
103            break;
104        case SERVICE_CONTROL_CONTINUE:
105            onContinue();
106            status->status.dwCurrentState = SERVICE_RUNNING;
107            break;
108        default:
109            return ERROR_CALL_NOT_IMPLEMENTED;
110    }
111    MORDOR_LOG_TRACE(g_log) << "Service control " << (ServiceCtrl)dwControl
112        << " handled";
113    // If it's stop, onTerminate may cause ServiceMain to exit immediately,
114    // before onTerminate returns, and our pointer to status on ServiceMain's
115    // stack would then be invalid; we update the service status in ServiceMain
116    // anyway, so it's okay to skip it here
117    if (dwControl != SERVICE_CONTROL_STOP) {
118        status->status.dwControlsAccepted = acceptedControls();
119        SetServiceStatus(status->serviceHandle, &status->status);
120    }
121    return NO_ERROR;
122}
123
124static void WINAPI ServiceMain(DWORD argc, LPWSTR *argvW)
125{
126    ServiceStatus status;
127    status.status.dwServiceType = SERVICE_WIN32_OWN_PROCESS;
128    status.status.dwCurrentState = SERVICE_RUNNING;
129    status.status.dwControlsAccepted = acceptedControls();
130    status.status.dwWin32ExitCode = NO_ERROR;
131    status.status.dwServiceSpecificExitCode = 0;
132    status.status.dwCheckPoint = 0;
133    status.status.dwWaitHint = 0;
134
135    status.serviceHandle = RegisterServiceCtrlHandlerExW(NULL,
136        &HandlerEx, &status);
137    if (!status.serviceHandle)
138        return;
139
140    if (!SetServiceStatus(status.serviceHandle, &status.status))
141        return;
142
143    MORDOR_LOG_INFO(g_log) << "Starting service";
144    char **argv = CommandLineToUtf8(argc, argvW);
145    if (argv) {
146        try {
147            status.status.dwServiceSpecificExitCode =
148                g_daemonMain(argc, argv);
149            LocalFree(argv);
150        } catch (...) {
151            LocalFree(argv);
152            throw;
153        }
154    } else {
155        status.status.dwServiceSpecificExitCode = GetLastError();
156    }
157    status.status.dwCurrentState = SERVICE_STOPPED;
158    if (status.status.dwServiceSpecificExitCode)
159        status.status.dwWin32ExitCode = ERROR_SERVICE_SPECIFIC_ERROR;
160    SetServiceStatus(status.serviceHandle, &status.status);
161    MORDOR_LOG_INFO(g_log) << "Service stopped";
162}
163
164namespace {
165
166enum ConsoleCtrl {};
167
168std::ostream &operator <<(std::ostream &os, ConsoleCtrl ctrl)
169{
170    switch (ctrl) {
171        CTRL_CASE(CTRL_C_EVENT);
172        CTRL_CASE(CTRL_BREAK_EVENT);
173        CTRL_CASE(CTRL_CLOSE_EVENT);
174        CTRL_CASE(CTRL_LOGOFF_EVENT);
175        CTRL_CASE(CTRL_SHUTDOWN_EVENT);
176        default:
177            return os << (int)ctrl;
178    }
179}
180
181}
182
183static BOOL WINAPI HandlerRoutine(DWORD dwCtrlType)
184{
185    MORDOR_LOG_DEBUG(g_log) << "Received console control "
186        << (ConsoleCtrl)dwCtrlType;
187    switch (dwCtrlType) {
188        case CTRL_C_EVENT:
189            if (onInterrupt.empty())
190                return FALSE;
191            onInterrupt();
192            break;
193        case CTRL_CLOSE_EVENT:
194        case CTRL_BREAK_EVENT:
195            if (onTerminate.empty())
196                return FALSE;
197            onTerminate();
198            break;
199        default:
200            return FALSE;
201    }
202    MORDOR_LOG_TRACE(g_log) << "Console control " << (ServiceCtrl)dwCtrlType
203        << " handled";
204    return TRUE;
205}
206
207int run(int argc, char **argv,
208    boost::function<int (int, char **)> daemonMain, bool)
209{
210    SERVICE_TABLE_ENTRYW ServiceStartTable[] = {
211        { L"", &ServiceMain },
212        { NULL,                NULL }
213    };
214    g_daemonMain = daemonMain;
215    if (!StartServiceCtrlDispatcherW(ServiceStartTable)) {
216        if (GetLastError() == ERROR_FAILED_SERVICE_CONTROLLER_CONNECT) {
217            MORDOR_LOG_INFO(g_log) << "Starting console process";
218            SetConsoleCtrlHandler(&HandlerRoutine, TRUE);
219            int result = daemonMain(argc, argv);
220            SetConsoleCtrlHandler(&HandlerRoutine, FALSE);
221            MORDOR_LOG_INFO(g_log) << "Console process stopped";
222            return result;
223        } else {
224            return GetLastError();
225        }
226    } else {
227        return 0;
228    }
229}
230#else
231
232#include <errno.h>
233#include <signal.h>
234#include <sys/wait.h>
235boost::function<bool (pid_t, int)> onChildProcessExit;
236
237static sigset_t blockedSignals()
238{
239    sigset_t result;
240    sigemptyset(&result);
241    sigaddset(&result, SIGTERM);
242    sigaddset(&result, SIGINT);
243    sigaddset(&result, SIGHUP);
244    sigaddset(&result, SIGTSTP);
245    sigaddset(&result, SIGCONT);
246    return result;
247}
248
249static void unblockAndRaise(int signal)
250{
251    sigset_t set;
252    sigemptyset(&set);
253    sigaddset(&set, signal);
254    pthread_sigmask(SIG_UNBLOCK, &set, NULL);
255    raise(signal);
256    pthread_sigmask(SIG_BLOCK, &set, NULL);
257}
258
259namespace {
260
261enum Signal {};
262
263std::ostream &operator <<(std::ostream &os, Signal signal)
264{
265    switch ((int)signal) {
266        CTRL_CASE(SIGTERM);
267        CTRL_CASE(SIGINT);
268        CTRL_CASE(SIGTSTP);
269        CTRL_CASE(SIGCONT);
270        CTRL_CASE(SIGHUP);
271        default:
272            return os << (int)signal;
273    }
274}
275
276}
277
278static void *signal_thread(void *arg)
279{
280    sigset_t mask = blockedSignals();
281    while (true) {
282        int caught;
283        int rc = sigwait(&mask, &caught);
284        if (rc != 0)
285            continue;
286        MORDOR_LOG_DEBUG(g_log) << "Received signal " << (Signal)caught;
287        switch (caught) {
288            case SIGTERM:
289                if (onTerminate.empty()) {
290                    unblockAndRaise(caught);
291                    continue;
292                }
293                onTerminate();
294                break;
295            case SIGINT:
296                if (onInterrupt.empty()) {
297                    unblockAndRaise(caught);
298                    continue;
299                }
300                onInterrupt();
301                break;
302            case SIGTSTP:
303                onPause();
304                unblockAndRaise(caught);
305                break;
306            case SIGCONT:
307                onContinue();
308                unblockAndRaise(caught);
309                break;
310            case SIGHUP:
311                if (onReload.empty()) {
312                    unblockAndRaise(caught);
313                    continue;
314                }
315                onReload();
316                break;
317            default:
318                continue;
319        }
320        MORDOR_LOG_TRACE(g_log) << "Signal " << (Signal)caught
321            << " handled";
322    }
323
324    return NULL;
325}
326
327#ifndef OSX
328static bool shouldDaemonize(char **enviro)
329{
330    if (!enviro)
331        return false;
332    std::string parent;
333    for (const char *env = *enviro; *env; env += strlen(env) + 1) {
334        const char *equals = strchr(env, '=');
335        if (equals != env + 1 || *env != '_')
336            continue;
337        parent = equals + 1;
338        break;
339    }
340    if (parent.size() >= 12 &&
341        strncmp(parent.c_str(), "/etc/init.d/", 12) == 0)
342        return true;
343    if (parent.size() >= 17 &&
344        strcmp(parent.c_str() + parent.size() - 17, "start-stop-daemon") == 0)
345        return true;
346    return false;
347}
348
349static bool shouldDaemonizeDueToParent()
350{
351    std::ostringstream os;
352    os << "/proc/" << getppid() << "/environ";
353    std::string parentEnviron;
354    parentEnviron.resize(65536);
355    int fd = open(os.str().c_str(), O_RDONLY);
356    if (fd < 0)
357        return false;
358    int size = read(fd, &parentEnviron[0], 65536);
359    close(fd);
360    if (size < 0)
361        return false;
362    parentEnviron.resize(size + 1);
363    const char *parentEnviro = parentEnviron.c_str();
364    return shouldDaemonize((char **)&parentEnviro);
365}
366#endif
367
368// running user specified main
369static int startMain(int argc, char** argv,
370        boost::function<int (int, char**)> userMain)
371{
372    // Mask signals from other threads so we can handle them ourselves
373    sigset_t mask = blockedSignals();
374    int rc = pthread_sigmask(SIG_BLOCK, &mask, NULL);
375    if (rc != 0)
376        return errno;
377
378    // Create the thread that will handle signals for us
379    pthread_t signal_thread_id;
380    rc = pthread_create(&signal_thread_id, NULL, &signal_thread, NULL);
381    if (rc != 0)
382        return errno;
383
384    // Run the daemon's main
385    MORDOR_LOG_INFO(g_log) << "Starting main in daemon";
386    rc = userMain(argc, argv);
387    MORDOR_LOG_INFO(g_log) << "Main stopped in daemon";
388    return rc;
389}
390
391// start main in a forked process and monitor it, restart it in case
392// it dies abnormally.
393static int watchdog(int argc, char** argv,
394        boost::function<int (int, char**)> userMain)
395{
396    for (;;) {
397        pid_t pid = fork();
398        switch (pid) {
399        case -1:
400            MORDOR_LOG_ERROR(g_log) << "Unable to fork(2): " << errno;
401            return errno;
402        case 0:
403            MORDOR_LOG_INFO(g_log) << "Child #" << getpid() << " started";
404            return startMain(argc, argv, userMain);
405        default:
406            MORDOR_LOG_INFO(g_log) << "Watchdog starts monitoring child #" << pid;
407            int status;
408            if (waitpid(pid, &status, 0) == -1) {
409                MORDOR_LOG_ERROR(g_log) << "Failed to waitpid(2): " << errno;
410                return errno;
411            }
412            MORDOR_LOG_INFO(g_log) << "Child #" << pid << " dies with status:"
413                                   << status;
414            bool done = false;
415            // check if program wish to quit
416            if (!onChildProcessExit.empty()) {
417                done = onChildProcessExit(pid, status);
418                MORDOR_LOG_INFO(g_log) << "onChildProcessExit returns " << done;
419            } else {
420                MORDOR_LOG_INFO(g_log) << "onChildProcessExit is not set, "
421                                       << "restart child by default";
422            }
423            if (done) {
424                MORDOR_LOG_INFO(g_log) << "Watchdog stopped";
425                return 0;
426            }
427        }
428    }
429}
430
431int run(int argc, char **argv,
432    boost::function<int (int, char **)> daemonMain,
433    bool enableWatchdog)
434{
435#ifndef OSX
436    // Check for being run from /etc/init.d or start-stop-daemon as a hint to
437    // daemonize
438    if (shouldDaemonize(environ) || shouldDaemonizeDueToParent()) {
439        MORDOR_LOG_VERBOSE(g_log) << "Daemonizing";
440        if (daemon(0, 0) == -1)
441            return errno;
442    }
443#endif
444
445    if (enableWatchdog) {
446        return watchdog(argc, argv, daemonMain);
447    } else {
448        return startMain(argc, argv, daemonMain);
449    }
450}
451#endif
452
453}}