implement daemon

shadowsocks/daemon.py

@@ -0,0 +1,173 @@
+#!/usr/bin/python
+# -*- coding: utf-8 -*-
+
+# Copyright (c) 2014 clowwindy
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+from __future__ import absolute_import, division, print_function, \
+    with_statement
+
+import os
+import sys
+import logging
+import signal
+import time
+from shadowsocks import common
+
+# this module is ported from ShadowVPN daemon.c
+
+
+def daemon_exec(config):
+    if 'daemon' in config:
+        if os.name != 'posix':
+            raise Exception('daemon mode is only supported in unix')
+        command = config['daemon']
+        if not command:
+            command = 'start'
+        pid_file = config['pid-file']
+        log_file = config['log-file']
+        command = common.to_str(command)
+        pid_file = common.to_str(pid_file)
+        log_file = common.to_str(log_file)
+        if command == 'start':
+            daemon_start(pid_file, log_file)
+        elif command == 'stop':
+            daemon_stop(pid_file)
+            # always exit after daemon_stop
+            sys.exit(0)
+        elif command == 'restart':
+            daemon_stop(pid_file)
+            daemon_start(pid_file, log_file)
+        else:
+            raise Exception('unsupported daemon command %s' % command)
+
+
+def write_pid_file(pid_file, pid):
+    import fcntl
+    import stat
+
+    try:
+        fd = os.open(pid_file, os.O_RDWR | os.O_CREAT,
+                     stat.S_IRUSR | stat.S_IWUSR)
+    except OSError as e:
+        logging.error(e)
+        return -1
+    flags = fcntl.fcntl(fd, fcntl.F_GETFD)
+    assert flags != -1
+    flags |= fcntl.FD_CLOEXEC
+    r = fcntl.fcntl(fd, fcntl.F_SETFD, flags)
+    assert r != -1
+    # There is no platform independent way to implement fcntl(fd, F_SETLK, &fl)
+    # via fcntl.fcntl. So use lockf instead
+    try:
+        fcntl.lockf(fd, fcntl.LOCK_EX | fcntl.LOCK_NB, 0, 0, os.SEEK_SET)
+    except IOError:
+        r = os.read(fd, 32)
+        if r:
+            logging.error('already started at pid %s' % common.to_str(r))
+        else:
+            logging.error('already started')
+        os.close(fd)
+        return -1
+    os.ftruncate(fd, 0)
+    os.write(fd, common.to_bytes(str(pid)))
+    return 0
+
+
+def freopen(f, mode, stream):
+    oldf = open(f, mode)
+    oldfd = oldf.fileno()
+    newfd = stream.fileno()
+    os.close(newfd)
+    os.dup2(oldfd, newfd)
+
+
+def daemon_start(pid_file, log_file):
+    # fork only once because we are sure parent will exit
+    pid = os.fork()
+    assert pid != -1
+
+    def handle_exit(signum, _):
+        sys.exit(0)
+
+    if pid > 0:
+        # parent waits for its child
+        signal.signal(signal.SIGINT, handle_exit)
+        time.sleep(5)
+        sys.exit(0)
+
+    # child signals its parent to exit
+    ppid = os.getppid()
+    pid = os.getpid()
+    if write_pid_file(pid_file, pid) != 0:
+        os.kill(ppid, signal.SIGINT)
+        sys.exit(1)
+
+    print('started')
+    os.kill(ppid, signal.SIGINT)
+
+    sys.stdin.close()
+    freopen(log_file, 'a', sys.stdout)
+    freopen(log_file, 'a', sys.stderr)
+
+
+def daemon_stop(pid_file):
+    import errno
+    try:
+        with open(pid_file) as f:
+            buf = f.read()
+            pid = common.to_str(buf)
+            if not buf:
+                logging.error('not running')
+    except IOError as e:
+        logging.error(e)
+        if e.errno == errno.ENOENT:
+            # always exit 0 if we are sure daemon is not running
+            logging.error('not running')
+            return
+        sys.exit(1)
+    pid = int(pid)
+    if pid > 0:
+        try:
+            os.kill(pid, signal.SIGTERM)
+        except OSError as e:
+            if e.errno == errno.ESRCH:
+                logging.error('not running')
+                # always exit 0 if we are sure daemon is not running
+                return
+            logging.error(e)
+            sys.exit(1)
+    else:
+        logging.error('pid is not positive: %d', pid)
+
+    # sleep for maximum 10s
+    for i in range(0, 200):
+        try:
+            # query for the pid
+            os.kill(pid, 0)
+        except OSError as e:
+            if e.errno == errno.ESRCH:
+                break
+        time.sleep(0.05)
+    else:
+        logging.error('timed out when stopping pid %d', pid)
+        sys.exit(1)
+    print('stopped')
+    os.unlink(pid_file)