From c9a905e0ab69317448f261377ab9031fb83443b4 Mon Sep 17 00:00:00 2001 From: Michael Orlitzky Date: Wed, 16 Jul 2014 17:10:27 -0400 Subject: [PATCH] Add the Unix module and enable daemonization. Add some dependencies for the Unix module. --- .ghci | 2 + halcyon.cabal | 8 +- src/ExitCodes.hs | 13 ++- src/Main.hs | 60 +++++++++++--- src/OptionalConfiguration.hs | 13 +-- src/Unix.hs | 152 +++++++++++++++++++++++++++++++++++ 6 files changed, 228 insertions(+), 20 deletions(-) create mode 100644 src/Unix.hs diff --git a/.ghci b/.ghci index 9d0bab9..8d30952 100644 --- 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. diff --git a/halcyon.cabal b/halcyon.cabal index 877d6ab..da45e30 100644 --- a/halcyon.cabal +++ b/halcyon.cabal @@ -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 diff --git a/src/ExitCodes.hs b/src/ExitCodes.hs index 7597a54..65a5198 100644 --- a/src/ExitCodes.hs +++ b/src/ExitCodes.hs @@ -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 diff --git a/src/Main.hs b/src/Main.hs index 3c33d1d..1d82ee4 100644 --- a/src/Main.hs +++ b/src/Main.hs @@ -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) diff --git a/src/OptionalConfiguration.hs b/src/OptionalConfiguration.hs index fca9e58..17c7191 100644 --- a/src/OptionalConfiguration.hs +++ b/src/OptionalConfiguration.hs @@ -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 index 0000000..0de0203 --- /dev/null +++ b/src/Unix.hs @@ -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 -- 2.43.2