Skip to main content

daemonize_me/
daemon.rs

1use std::any::Any;
2use std::convert::TryFrom;
3use std::ffi::{CString, OsStr, OsString};
4use std::fs::File;
5use std::io::prelude::*;
6use std::path::{Path, PathBuf};
7use std::process::exit;
8
9use nix::sys::stat::{umask, Mode};
10#[cfg(target_os = "macos")]
11use nix::unistd::{
12    chdir, chown, close, dup2, fork, getpid, setgid, setsid, setuid, ForkResult, Gid, Pid, Uid,
13};
14#[cfg(not(target_os = "macos"))]
15use nix::unistd::{
16    chdir, chown, fork, getpid, initgroups, setgid, setsid, setuid, ForkResult, Gid, Pid, Uid,
17};
18
19use crate::ffi::{set_proc_name, PasswdRecord};
20use crate::group::Group;
21use crate::stdio::{redirect_stdio, Stdio};
22use crate::user::User;
23use crate::DaemonError::{InvalidGroup, InvalidUser};
24use crate::{DaemonError, Result};
25
26/// Basic daemonization consists of:
27/// forking the process, getting a new Session ID (sid), setting the umask, changing the standard io streams
28/// to files and finally dropping privileges.
29///
30/// **NOTE:** Beware there is no escalation back if dropping privileges
31pub struct Daemon<'a> {
32    pub(crate) chdir: PathBuf,
33    pub(crate) pid_file: Option<PathBuf>,
34    pub(crate) chown_pid_file: bool,
35    pub(crate) user: Option<User>,
36    pub(crate) group: Option<Group>,
37    pub(crate) umask: u16,
38    // stdin is practically always null
39    pub(crate) stdin: Stdio,
40    pub(crate) stdout: Stdio,
41    pub(crate) stderr: Stdio,
42    pub(crate) name: Option<OsString>,
43    pub(crate) before_fork_hook: Option<fn(pid: i32)>,
44    pub(crate) after_fork_parent_hook: Option<fn(parent_pid: i32, child_pid: i32) -> !>,
45    pub(crate) after_fork_child_hook: Option<fn(parent_pid: i32, child_pid: i32) -> ()>,
46    pub(crate) after_init_hook_data: Option<&'a dyn Any>,
47    pub(crate) after_init_hook: Option<fn(Option<&'a dyn Any>)>,
48}
49
50impl<'a> Daemon<'a> {
51    pub fn new() -> Self {
52        Daemon {
53            chdir: Path::new("/").to_owned(),
54            pid_file: None,
55            chown_pid_file: false,
56            user: None,
57            group: None,
58            umask: 0o027,
59            stdin: Stdio::devnull(),
60            stdout: Stdio::devnull(),
61            stderr: Stdio::devnull(),
62            name: None,
63            before_fork_hook: None,
64            after_fork_parent_hook: None,
65            after_fork_child_hook: None,
66            after_init_hook_data: None,
67            after_init_hook: None,
68        }
69    }
70
71    /// Give your daemon a pid file
72    ///
73    /// By default, no pid file is created.
74    ///
75    /// # Arguments
76    /// * `path` - path to the file suggested `/var/run/my_program_name.pid`
77    /// * `chmod` - if set a chmod of the file to the user and group passed will be attempted (**this being true makes setting an user and group mandatory**)
78    pub fn pid_file<T: AsRef<Path>>(mut self, path: T, chmod: Option<bool>) -> Self {
79        self.pid_file = Some(path.as_ref().to_owned());
80        self.chown_pid_file = chmod.unwrap_or(false);
81        self
82    }
83
84    /// As the last step the code will change the working directory to this one defaults to `/`
85    pub fn work_dir<T: AsRef<Path>>(mut self, path: T) -> Self {
86        self.chdir = path.as_ref().to_owned();
87        self
88    }
89
90    /// The code will attempt to drop privileges with `setuid` to the provided user
91    ///
92    /// **NOTE:** If you provide a user, you must also provide a group.
93    pub fn user<T: Into<User>>(mut self, user: T) -> Self {
94        self.user = Some(user.into());
95        self
96    }
97
98    /// The code will attempt to drop privileges with `setgid` to the provided group
99    ///
100    /// **NOTE:** You must provide a group if you provide an user.
101    pub fn group<T: Into<Group>>(mut self, group: T) -> Self {
102        self.group = Some(group.into());
103        self
104    }
105
106    pub fn group_copy_user(mut self) -> Result<Self> {
107        if let Some(user) = &self.user {
108            self.group = Some(Group::try_from(&user.name)?);
109            Ok(self)
110        } else {
111            Err(InvalidUser)
112        }
113    }
114
115    /// umask for the process, defaults to `0o027`
116    pub fn umask(mut self, mask: u16) -> Self {
117        self.umask = mask;
118        self
119    }
120
121    pub fn stdin<T: Into<Stdio>>(mut self, stdio: T) -> Self {
122        self.stdin = stdio.into();
123        self
124    }
125
126    /// Determines where standard output will be piped to since daemons have no console attached
127    ///
128    /// It's highly recommended to set this to a file if you want to see output.
129    pub fn stdout<T: Into<Stdio>>(mut self, stdio: T) -> Self {
130        self.stdout = stdio.into();
131        self
132    }
133
134    /// Determines where standard error will be piped to since daemons have no console attached
135    ///
136    /// It's highly recommended to set this to a file if you want to see output.
137    pub fn stderr<T: Into<Stdio>>(mut self, stdio: T) -> Self {
138        self.stderr = stdio.into();
139        self
140    }
141
142    /// Set the daemon process name
143    ///
144    /// For example, this is what shows up in `ps`.
145    pub fn name(mut self, name: &OsStr) -> Self {
146        self.name = Some(OsString::from(name));
147        self
148    }
149
150    /// Hook called before the fork with the current pid as argument
151    pub fn setup_pre_fork_hook(mut self, pre_fork_hook: fn(pid: i32)) -> Self {
152        self.before_fork_hook = Some(pre_fork_hook);
153        self
154    }
155
156    /// Hook called after the fork with the parent pid as argument
157    ///
158    /// Can be used to continue some work on the parent after the fork.
159    /// **NOTE:** This hook must not return! For instance, you could call `std::process::exit()`.
160    ///
161    /// # Examples
162    ///
163    /// ```
164    /// # use daemonize_me::Daemon;
165    ///
166    /// fn post_fork_parent(ppid: i32, cpid: i32) -> ! {
167    ///     println!("Parent pid: {}, Child pid {}", ppid, cpid);
168    ///     println!("Exiting parent now");
169    ///     std::process::exit(0);
170    /// }
171    ///
172    /// let daemon = Daemon::new()
173    ///     .setup_post_fork_parent_hook(post_fork_parent)
174    ///     .start();
175    /// ```
176    pub fn setup_post_fork_parent_hook(
177        mut self,
178        post_fork_parent_hook: fn(parent_pid: i32, child_pid: i32) -> !,
179    ) -> Self {
180        self.after_fork_parent_hook = Some(post_fork_parent_hook);
181        self
182    }
183
184    /// Hook called after the fork with the parent and child pid as arguments
185    ///
186    /// # Examples
187    ///
188    /// ```
189    /// # use daemonize_me::Daemon;
190    ///
191    /// fn post_fork_child(ppid: i32, cpid: i32) {
192    ///     println!("Parent pid: {}, Child pid {}", ppid, cpid);
193    ///     println!("This hook is called in the child");
194    ///     // Child hook must return
195    ///     return
196    /// }
197    ///
198    /// let daemon = Daemon::new()
199    ///     .setup_post_fork_child_hook(post_fork_child)
200    ///     .start();
201    /// ```
202    pub fn setup_post_fork_child_hook(
203        mut self,
204        post_fork_child_hook: fn(parent_pid: i32, child_pid: i32) -> (),
205    ) -> Self {
206        self.after_fork_child_hook = Some(post_fork_child_hook);
207        self
208    }
209
210    /// Hook called right before returning control to the caller, that is, right after `start()`
211    ///
212    /// # Examples
213    ///
214    /// ```
215    /// # use std::any::Any;
216    /// # use daemonize_me::Daemon;
217    ///
218    /// fn after_init(_: Option<&dyn Any>) {
219    ///     println!("Initialized the daemon!");
220    ///     return
221    /// }
222    ///
223    /// let daemon = Daemon::new()
224    ///     .setup_post_init_hook(after_init, None)
225    ///     .start();
226    /// ```
227    pub fn setup_post_init_hook(
228        mut self,
229        post_fork_child_hook: fn(ctx: Option<&'a dyn Any>),
230        data: Option<&'a dyn Any>,
231    ) -> Self {
232        self.after_init_hook = Some(post_fork_child_hook);
233        self.after_init_hook_data = data;
234        self
235    }
236
237    /// Using the parameters set, daemonize the process
238    pub fn start(self) -> Result<()> {
239        let mut pid: Pid;
240        let parent_pid = getpid();
241        // resolve options to concrete values to please the borrow checker
242        let has_pid_file = self.pid_file.is_some();
243        let pid_file_path = match self.pid_file {
244            Some(path) => path.clone(),
245            None => Path::new("").to_path_buf(),
246        };
247
248        // If the hook is set call it with the parent pid
249        if let Some(hook) = self.before_fork_hook {
250            hook(parent_pid.as_raw());
251        }
252
253        // Fork and if the process is the parent exit gracefully
254        // if the  process is the child just continue execution
255        // this was made unsafe by the nix upstream in between versions
256        // thus the unsafe block is required here
257        unsafe {
258            match fork() {
259                Ok(ForkResult::Parent { child: cpid }) => {
260                    if let Some(hook) = self.after_fork_parent_hook {
261                        hook(parent_pid.as_raw(), cpid.as_raw());
262                    } else {
263                        exit(0)
264                    }
265                }
266                Ok(ForkResult::Child) => {
267                    // Set up stream redirection as early as possible
268                    redirect_stdio(&self.stdin, &self.stdout, &self.stderr)?;
269                    pid = getpid();
270                    if let Some(hook) = self.after_fork_child_hook {
271                        hook(parent_pid.as_raw(), pid.as_raw());
272                    }
273                    ()
274                }
275                Err(_) => return Err(DaemonError::Fork),
276            }
277        }
278
279        if self.chown_pid_file && (self.user.is_none() || self.group.is_none()) {
280            return Err(DaemonError::InvalidUserGroupPair);
281        } else if (self.user.is_some() || self.group.is_some())
282            && (self.user.is_none() || self.group.is_none())
283        {
284            return Err(DaemonError::InvalidUserGroupPair);
285        }
286
287        if let Some(proc_name) = &self.name {
288            match set_proc_name(proc_name.as_ref()) {
289                Ok(()) => (),
290                Err(e) => return Err(e),
291            }
292        }
293        // Set the umask either to 0o027 (rwxr-x---) or provided value
294        let umask_mode = match Mode::from_bits(self.umask as _) {
295            Some(mode) => mode,
296            None => return Err(DaemonError::InvalidUmaskBits),
297        };
298        umask(umask_mode);
299
300        // Set the sid so the process isn't session orphan
301        if let Err(_) = setsid() {
302            return Err(DaemonError::SetSid);
303        };
304        if let Err(_) = chdir::<Path>(self.chdir.as_path()) {
305            return Err(DaemonError::ChDir);
306        };
307        pid = getpid();
308
309        // create pid file and if configured to, chmod it
310        if has_pid_file {
311            // chmod of the pid file is deferred to after checking for the presence of the user and group
312            let pid_file = &pid_file_path;
313            match File::create(pid_file) {
314                Ok(mut fp) => {
315                    if let Err(_) = fp.write_all(pid.to_string().as_ref()) {
316                        return Err(DaemonError::WritePid);
317                    }
318                }
319                Err(_) => return Err(DaemonError::WritePid),
320            };
321        }
322
323        // Drop privileges and chown the requested files
324        if self.user.is_some() && self.group.is_some() {
325            let user = match self.user {
326                Some(user) => Uid::from_raw(user.id),
327                None => return Err(InvalidUser),
328            };
329
330            let uname = match PasswdRecord::lookup_record_by_id(user.as_raw()) {
331                Ok(record) => record.pw_name,
332                Err(_) => return Err(DaemonError::InvalidUser),
333            };
334
335            let gr = match self.group {
336                Some(grp) => Gid::from_raw(grp.id),
337                None => return Err(InvalidGroup),
338            };
339
340            if self.chown_pid_file && has_pid_file {
341                match chown(&pid_file_path, Some(user), Some(gr)) {
342                    Ok(_) => (),
343                    Err(_) => return Err(DaemonError::ChownPid),
344                };
345            }
346
347            match setgid(gr) {
348                Ok(_) => (),
349                Err(_) => return Err(DaemonError::SetGid),
350            };
351            #[cfg(not(target_os = "macos"))]
352            {
353                let u_cstr = match CString::new(uname) {
354                    Ok(cstr) => cstr,
355                    Err(_) => return Err(DaemonError::SetGid),
356                };
357                match initgroups(&u_cstr, gr) {
358                    Ok(_) => (),
359                    Err(_) => return Err(DaemonError::InitGroups),
360                };
361            }
362            match setuid(user) {
363                Ok(_) => (),
364                Err(_) => return Err(DaemonError::SetUid),
365            }
366        };
367        // chdir
368        let chdir_path = self.chdir.to_owned();
369        match chdir::<Path>(chdir_path.as_ref()) {
370            Ok(_) => (),
371            Err(_) => return Err(DaemonError::ChDir),
372        };
373
374        // Now this process should be a daemon, we run the hook and return or just return
375        if let Some(hook) = self.after_init_hook {
376            hook(self.after_init_hook_data);
377            Ok(())
378        } else {
379            Ok(())
380        }
381    }
382}