Add the Unix module and enable daemonization.
authorMichael Orlitzky <michael@orlitzky.com>
Wed, 16 Jul 2014 21:10:27 +0000 (17:10 -0400)
committerMichael Orlitzky <michael@orlitzky.com>
Wed, 16 Jul 2014 21:10:27 +0000 (17:10 -0400)
Add some dependencies for the Unix module.

.ghci
halcyon.cabal
src/ExitCodes.hs
src/Main.hs
src/OptionalConfiguration.hs
src/Unix.hs [new file with mode: 0644]

diff --git a/.ghci b/.ghci
index 9d0bab9d4753fc692d6283c729422bba1fa9cdd9..8d30952b798b649ad74da1b26a9400047e8c2b63 100644 (file)
--- a/.ghci
+++ b/.ghci
@@ -14,6 +14,7 @@
   src/Twitter/Http.hs
   src/Twitter/User.hs
   src/Twitter/Status.hs
+  src/Unix.hs
   src/Usernames.hs
 :}
 
@@ -27,6 +28,7 @@ import StringUtils
 import Twitter.Http
 import Twitter.User
 import Twitter.Status
+import Unix
 import Usernames
 
 -- Use a calmer prompt.
index 877d6abdfd01ed98183c781ce315d9cac52654c3..da45e30a20c21e3f8e3d6ec2a110f3bd5eeb8dda 100644 (file)
@@ -23,15 +23,16 @@ executable halcyon
     configurator                >= 0.2,
     directory                   >= 1.2,
     filepath                    >= 1.3,
+    hdaemonize                  >= 0.4,
     http-client                 >= 0.3,
     http-client-tls             >= 0.2,
-    HUnit                       == 1.2.*,
     MissingH                    >= 1.2,
     process                     >= 1.1,
     old-locale                  >= 1,
     tagsoup                     >= 0.13,
     text                        >= 1.1,
     time                        >= 1.4,
+    unix                        >= 2.6,
     -- Test deps
     tasty                       >= 0.8,
     tasty-hunit                 >= 0.3
@@ -53,6 +54,8 @@ executable halcyon
     Twitter.Http
     Twitter.Status
     Twitter.User
+    Unix
+    Usernames
 
   ghc-options:
     -Wall
@@ -69,6 +72,7 @@ executable halcyon
     -optc-march=native
     -O2
 
+
 test-suite testsuite
   type: exitcode-stdio-1.0
   hs-source-dirs: src test
@@ -82,6 +86,7 @@ test-suite testsuite
     configurator                >= 0.2,
     directory                   >= 1.2,
     filepath                    >= 1.3,
+    hdaemonize                  >= 0.4,
     http-client                 >= 0.3,
     http-client-tls             >= 0.2,
     MissingH                    >= 1.2,
@@ -90,6 +95,7 @@ test-suite testsuite
     tagsoup                     >= 0.13,
     text                        >= 1.1,
     time                        >= 1.4,
+    unix                        >= 2.6,
     -- Test deps
     tasty                       >= 0.8,
     tasty-hunit                 >= 0.3
index 7597a54723e0fa28d466f195f5161e60698ea3c4..65a519869cf3f88ecfa86c64c0353b4d0c9c4553 100644 (file)
@@ -1,9 +1,20 @@
 -- | All exit codes that the program can return (excepting
 --   ExitSuccess).
-module ExitCodes ( exit_no_usernames )
+module ExitCodes (
+  exit_no_usernames,
+  exit_pidfile_exists )
 where
 
+
 -- | No usernames found on either the command-line or in a config
 --   file.
 exit_no_usernames :: Int
 exit_no_usernames = 1
+
+
+-- | When running as a daemon, the existence of a fixed PID file is
+--   used to determine whether or not the daemon is already
+--   running. If the PID file already exists, we shouldn't start.
+--
+exit_pidfile_exists :: Int
+exit_pidfile_exists = 2
index 3c33d1d8aadcc0274e933c8ce3d68bca1ee296b0..1d82ee4d5c6125653ac449fa10e3a65d84354318 100644 (file)
@@ -1,19 +1,24 @@
+{-# LANGUAGE DoAndIfThenElse #-}
+
 module Main
 where
 
 import Control.Concurrent ( forkIO, threadDelay )
+import Control.Exception ( throw )
 import Control.Monad ( forever, when )
 import Data.Aeson ( decode )
 import Data.Maybe ( fromMaybe, isNothing )
 import Data.Monoid ( (<>) )
 import Data.Time.LocalTime ( TimeZone, getCurrentTimeZone )
 import System.Console.CmdArgs.Default ( Default(..) )
+import System.Directory ( doesFileExist )
 import System.Exit ( ExitCode(..), exitWith )
 import System.IO ( hPutStrLn, stderr )
+import System.IO.Error ( catchIOError )
 
 import CommandLine ( get_args, show_help )
 import Configuration ( Cfg(..), merge_optional )
-import ExitCodes ( exit_no_usernames )
+import ExitCodes ( exit_no_usernames, exit_pidfile_exists )
 import qualified OptionalConfiguration as OC ( from_rc )
 import Mail (
   Message(..),
@@ -28,8 +33,10 @@ import Twitter.Status (
   pretty_print,
   utc_time_to_rfc822 )
 import Twitter.User ( User(..) )
+import Unix ( full_daemonize )
 import Usernames ( Usernames(..) )
 
+
 -- | A wrapper around threadDelay which takes seconds instead of
 --   microseconds as its argument.
 --
@@ -98,8 +105,8 @@ mention_retweets cfg ss = do
 -- | Filter out replies/retweets based on the configuration.
 --
 filter_statuses :: Cfg -> [Status] -> [Status]
-filter_statuses cfg statuses =
-  (reply_filter . retweet_filter) statuses
+filter_statuses cfg =
+  reply_filter . retweet_filter
   where
     reply_filter = if ignore_replies cfg
                    then filter (not . reply)
@@ -230,18 +237,47 @@ main = do
     _ <- show_help
     exitWith (ExitFailure exit_no_usernames)
 
+
+  when (daemonize cfg) $ do
+    -- Old PID files can be left around after an unclean shutdown. We
+    -- only care if we're running as a daemon.
+    pidfile_exists <- doesFileExist (pidfile cfg)
+    when pidfile_exists $ do
+      hPutStrLn stderr $ "ERROR: PID file " ++ (pidfile cfg) ++
+                         " already exists. Refusing to start."
+      exitWith (ExitFailure exit_pidfile_exists)
+
+
   -- If we have both a "To" and "From" address, we'll create a
   -- message object to be passed to all of our threads.
   let message = construct_message cfg
-
-  -- Execute run on each username in a new thread.
   let run_curried = run cfg message
-  _ <- mapM (forkIO . run_curried) (get_usernames (usernames cfg))
 
-  _ <- forever $
-    -- This thread (the one executing main) doesn't do anything,
-    -- but when it terminates, so do all the threads we forked.
-    -- As a result, we need to keep this thread on life support.
-    thread_sleep (heartbeat cfg)
+  let run_program = (mapM_ -- Execute run on each username in a new thread.
+                      (forkIO . run_curried)
+                      (get_usernames (usernames cfg)))
+                    >>
+                    -- This thread (the one executing main) doesn't do
+                    -- anything, but when it terminates, so do all the
+                    -- threads we forked.  As a result, we need to
+                    -- keep this thread on life support.  If we were
+                    -- asked to daemonize, do that; otherwise just run
+                    -- the thing.
+                    forever (thread_sleep (heartbeat cfg))
+
+  if (daemonize cfg)
+  then try_daemonize cfg run_program
+  else run_program
 
-  return ()
+  where
+    -- | A exception handler around full_daemonize. If full_daemonize
+    --   doesn't work, we report the error and crash. This is fine; we
+    --   only need the program to be resilient once it actually starts.
+    --
+    try_daemonize :: Cfg -> IO () -> IO ()
+    try_daemonize cfg program =
+      catchIOError
+        (full_daemonize cfg program)
+        (\e -> do
+          hPutStrLn stderr ("ERROR: " ++ (show e))
+          throw e)
index fca9e58eca283abfa8d45912a71dfee3baa03b39..17c7191dead04be9265ec82e91604a730034ebb0 100644 (file)
@@ -113,12 +113,13 @@ instance Monoid OptionalCfg where
 --
 from_rc :: IO OptionalCfg
 from_rc = do
-  etc  <- catchIOError getSysconfDir (\e -> do
-                                        hPutStrLn stderr (show e)
-                                        return "/etc")
-  home <- catchIOError getHomeDirectory (\e -> do
-                                           hPutStrLn stderr (show e)
-                                           return "$(HOME)")
+  etc  <- catchIOError
+            getSysconfDir (\e -> do hPutStrLn stderr ("ERROR: " ++ (show e))
+                                    return "/etc")
+  home <- catchIOError
+            getHomeDirectory (\e -> do hPutStrLn stderr ("ERROR: " ++ (show e))
+                                       return "$(HOME)")
+
   let global_config_path = etc </> "halcyonrc"
   let user_config_path = home </> ".halcyonrc"
   cfg <- DC.load [ DC.Optional global_config_path,
diff --git a/src/Unix.hs b/src/Unix.hs
new file mode 100644 (file)
index 0000000..0de0203
--- /dev/null
@@ -0,0 +1,152 @@
+-- | Non-portable code for daemonizing on unix.
+--
+module Unix
+where
+
+import Control.Concurrent ( ThreadId, myThreadId )
+import Control.Exception ( throwTo )
+import Control.Monad ( unless )
+import System.Directory ( createDirectory, doesDirectoryExist )
+import System.Exit ( ExitCode( ExitSuccess ) )
+import System.FilePath ( dropFileName, dropTrailingPathSeparator )
+import System.IO ( hPutStrLn, stderr )
+import System.IO.Error ( catchIOError )
+import System.Posix (
+  GroupEntry ( groupID ),
+  GroupID,
+  Handler ( Catch ),
+  UserEntry ( userID ),
+  UserID,
+  exitImmediately,
+  getGroupEntryForName,
+  getProcessID,
+  getRealGroupID,
+  getRealUserID,
+  getUserEntryForName,
+  installHandler,
+  removeLink,
+  setFileCreationMask,
+  setGroupID,
+  setOwnerAndGroup,
+  setUserID,
+  sigTERM )
+import System.Posix.Daemonize ( daemonize )
+
+import Configuration (
+  Cfg( pidfile,
+       run_as_group,
+       run_as_user ) )
+
+
+-- | Retrieve the uid associated with the given system user name. We
+--   take a Maybe String as an argument so the user name can be passed
+--   in directly from the config.
+--
+get_user_id :: Maybe String -> IO UserID
+get_user_id Nothing  = getRealUserID
+get_user_id (Just s) = fmap userID (getUserEntryForName s)
+
+
+-- | Retrieve the gid associated with the given system group name. We
+--   take a Maybe String as an argument so the group name can be
+--   passed in directly from the config.
+--
+get_group_id :: Maybe String -> IO GroupID
+get_group_id Nothing  = getRealGroupID
+get_group_id (Just s) = fmap groupID (getGroupEntryForName s)
+
+
+-- | This function will be called in response to a SIGTERM; i.e. when
+--   someone tries to kill our process. We simply delete the PID file
+--   and signal our parent thread to quit (successfully).
+--
+--   If that doesn't work, report the error and quit rudely.
+--
+graceful_shutdown :: Cfg -> ThreadId -> IO ()
+graceful_shutdown cfg main_thread_id = do
+  putStrLn "SIGTERM received, removing PID file and shutting down."
+  catchIOError try_nicely (\e -> do
+                             hPutStrLn stderr ("ERROR: " ++ (show e))
+                             exitImmediately ExitSuccess )
+  where
+    try_nicely = do
+      removeLink (pidfile cfg)
+      throwTo main_thread_id ExitSuccess
+
+
+-- | Create the directory in which we intend to store the PID
+--   file. This will *not* create any parent directories. The PID
+--   directory will have its owner/group changed to the user/group
+--   under which we'll be running. No permissions will be set; the
+--   system's umask must allow owner-write.
+--
+--   This is intended to create one level beneath either /var/run or
+--   /run which often do not survive a reboot.
+--
+--   If the directory already exists, it is left alone; that is, we
+--   don't change its owner/group.
+--
+create_pid_directory :: FilePath -- ^ The directory to contain the PID file.
+                     -> UserID   -- ^ Owner of the new directory if created.
+                     -> GroupID  -- ^ Group of the new directory if created.
+                     -> IO ()
+create_pid_directory pid_directory uid gid = do
+  it_exists <- doesDirectoryExist pid_directory
+  unless it_exists $ do
+    putStrLn $ "Creating PID directory " ++ pid_directory ++ "."
+    createDirectory pid_directory
+    putStrLn $ "Changing owner/group of " ++ pid_directory ++
+               " to " ++ (show uid) ++ "/" ++ (show gid) ++ "."
+    setOwnerAndGroup pid_directory uid gid
+
+
+-- | Write a PID file, install a SIGTERM handler, drop privileges, and
+--   finally do the daemonization dance.
+--
+full_daemonize :: Cfg -> IO () -> IO ()
+full_daemonize cfg program = do
+  uid <- get_user_id (run_as_user cfg)
+  gid <- get_group_id (run_as_group cfg)
+
+  -- This will have to be done as root and the result chowned to our
+  -- user/group, so it must happen before daemonizing.
+  let pid_directory = dropTrailingPathSeparator $ dropFileName $ pidfile cfg
+  create_pid_directory pid_directory uid gid
+
+  -- The call to 'daemonize' will set the umask to zero, but we want
+  -- to retain it. So, we set the umask to zero before 'daemonize'
+  -- can, so that we can record the previous umask value (returned by
+  -- setFileCreationMask).
+  orig_umask <- setFileCreationMask 0
+
+  -- This is the 'daemonize' from System.Posix.Daemonize.
+  daemonize (program' orig_umask uid gid)
+  where
+    -- We need to do all this stuff *after* we daemonize.
+    program' orig_umask uid gid = do
+      -- First we install a signal handler for sigTERM. We need to
+      -- pass the thread ID to the signal handler so it knows which
+      -- process to "exit."
+      tid <- myThreadId
+      _ <- installHandler sigTERM (Catch (graceful_shutdown cfg tid)) Nothing
+
+      -- Next we drop privileges. Group ID has to go first, otherwise
+      -- you ain't root to change groups.
+      setGroupID gid
+      setUserID uid
+
+      -- Now we create the PID file.
+      pid <- getProcessID
+
+      -- The PID file needs to be read-only for anyone but its
+      -- owner. Hopefully the umask accomplishes this!
+      _ <- setFileCreationMask orig_umask
+
+      -- When we later attempt to delete the PID file, it requires
+      -- write permission to the parent directory and not to the PID
+      -- file itself. Therefore, if that's going to work, this has to
+      -- work, even as a limited user.
+      writeFile (pidfile cfg) (show pid)
+
+      -- Finally run the program we were asked to.
+      program