Skip to content

Commit

Permalink
Duplicate all public functions with _and_unset_env
Browse files Browse the repository at this point in the history 
They need to be `unsafe` because `std::env::remove_var` is unsafe.
  • Loading branch information
@tbu-
tbu- committed Jan 15, 2025
1 parent 5e11d80 commit d9ed7a9
Showing 1 changed file with 79 additions and 124 deletions.
203 changes: 79 additions & 124 deletions src/lib.rs
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,4 @@
#![doc(html_root_url = "https://docs.rs/sd-notify/0.1.0")]
#![deny(missing_docs)]

//! Lightweight crate for interacting with `systemd`.
//!
Expand All @@ -26,6 +25,7 @@ use std::convert::TryFrom;
use std::env;
use std::fmt::{self, Display, Formatter, Write};
use std::fs;
use std::ffi::OsStr;
use std::io::{self, ErrorKind};
#[cfg(feature = "fdstore")]
use std::os::fd::BorrowedFd;
Expand Down Expand Up @@ -113,13 +113,16 @@ pub fn booted() -> io::Result<bool> {
Ok(m.is_dir())
}

// Constants for env variable names so that we don't typo them.
const LISTEN_FDNAMES: &str = "LISTEN_FDNAMES";
const LISTEN_FDS: &str = "LISTEN_FDS";
const LISTEN_PID: &str = "LISTEN_PID";
const NOTIFY_SOCKET: &str = "NOTIFY_SOCKET";
const WATCHDOG_PID: &str = "WATCHDOG_PID";
const WATCHDOG_USEC: &str = "WATCHDOG_USEC";

/// Sends the service manager a list of state changes.
///
/// The `unset_env` parameter should generally not be set, see the Safety section.
/// If the `unset_env` parameter is set, the `NOTIFY_SOCKET` environment variable
/// will be unset before returning. Further calls to `sd_notify` will fail, but
/// child processes will no longer inherit the variable.
///
/// The notification mechanism involves sending a datagram to a Unix domain socket.
/// See [`sd_notify(3)`][sd_notify] for details.
///
Expand All @@ -134,27 +137,22 @@ pub fn booted() -> io::Result<bool> {
///
/// If you wish to send file descriptors, use the `notify_with_fds` function.
///
/// # Safety
///
/// Although this function is not marked as `unsafe`, it is `unsafe` to call
/// with `unset_env` set to true. See the documentation on
/// [std::env::remove_var] for further information on this unsafety.
///
/// # Example
///
/// ```no_run
/// # use sd_notify::NotifyState;
/// #
/// let _ = sd_notify::notify(false, &[NotifyState::Ready]);
/// let _ = sd_notify::notify(&[NotifyState::Ready]);
/// ```
pub fn notify(unset_env: bool, state: &[NotifyState]) -> io::Result<()> {
let mut msg = String::new();
let Some(sock) = connect_notify_socket(unset_env)? else {
pub fn notify(state: &[NotifyState]) -> io::Result<()> {
let Some(socket_path) = env::var_os(NOTIFY_SOCKET) else {
return Ok(());
};
let mut msg = String::new();
for s in state {
let _ = writeln!(msg, "{}", s);
writeln!(msg, "{}", s).unwrap();
}
let sock = connected_unix_datagram(&socket_path)?;
let len = sock.send(msg.as_bytes())?;
if len != msg.len() {
Err(io::Error::new(ErrorKind::WriteZero, "incomplete write"))
Expand All @@ -163,13 +161,16 @@ pub fn notify(unset_env: bool, state: &[NotifyState]) -> io::Result<()> {
}
}

pub unsafe fn notify_and_unset_env(state: &[NotifyState]) -> io::Result<()> {
let result = notify(state);
unsafe {
env::remove_var(NOTIFY_SOCKET);
}
result
}

/// Sends the service manager a list of state changes with file descriptors.
///
/// The `unset_env` parameter should generally not be set, see the Safety section.
/// If the `unset_env` parameter is set, the `NOTIFY_SOCKET` environment variable
/// will be unset before returning. Further calls to `sd_notify` will fail, but
/// child processes will no longer inherit the variable.
///
/// The notification mechanism involves sending a datagram to a Unix domain socket.
/// See [`sd_pid_notify_with_fds(3)`][sd_pid_notify_with_fds] for details.
///
Expand All @@ -182,12 +183,6 @@ pub fn notify(unset_env: bool, state: &[NotifyState]) -> io::Result<()> {
/// and does not increase the send buffer size. It's still useful, though, in
/// usual situations.
///
/// # Safety
///
/// Although this function is not marked as `unsafe`, it is `unsafe` to call
/// with `unset_env` set to true. See the documentation on
/// [std::env::remove_var] for further information on this unsafety.
///
/// # Example
///
/// ```no_run
Expand All @@ -200,19 +195,19 @@ pub fn notify(unset_env: bool, state: &[NotifyState]) -> io::Result<()> {
/// ```
#[cfg(feature = "fdstore")]
pub fn notify_with_fds(
unset_env: bool,
state: &[NotifyState],
fds: &[BorrowedFd<'_>],
) -> io::Result<()> {
use sendfd::SendWithFd;

let mut msg = String::new();
let Some(sock) = connect_notify_socket(unset_env)? else {
let Some(socket_path) = env::var_os(NOTIFY_SOCKET) else {
return Ok(());
};
let mut msg = String::new();
for s in state {
let _ = writeln!(msg, "{}", s);
writeln!(msg, "{}", s).unwrap();
}
let sock = connected_unix_datagram(&socket_path);
let len = sock.send_with_fd(msg.as_bytes(), borrowed_fd_slice(fds))?;
if len != msg.len() {
Err(io::Error::new(ErrorKind::WriteZero, "incomplete write"))
Expand All @@ -221,6 +216,18 @@ pub fn notify_with_fds(
}
}

#[cfg(feature = "fdstore")]
pub unsafe fn notify_with_fds_and_unset_env(
state: &[NotifyState],
fds: &[BorrowedFd<'_>],
) -> io::Result<()> {
let result = notify_with_fds(state, fds);
unsafe {
env::remove_var(NOTIFY_SOCKET);
}
result
}

#[cfg(feature = "fdstore")]
fn borrowed_fd_slice<'a>(s: &'a [BorrowedFd<'_>]) -> &'a [RawFd] {
// SAFETY: BorrowedFd is #[repr(transparent)] over RawFd (memory safety)
Expand All @@ -229,18 +236,10 @@ fn borrowed_fd_slice<'a>(s: &'a [BorrowedFd<'_>]) -> &'a [RawFd] {
unsafe { std::mem::transmute(s) }
}

fn connect_notify_socket(unset_env: bool) -> io::Result<Option<UnixDatagram>> {
let Some(socket_path) = env::var_os("NOTIFY_SOCKET") else {
return Ok(None);
};

if unset_env {
env::remove_var("NOTIFY_SOCKET");
}

let sock = UnixDatagram::unbound()?;
sock.connect(socket_path)?;
Ok(Some(sock))
fn connected_unix_datagram(path: &OsStr) -> io::Result<UnixDatagram> {
let result = UnixDatagram::unbound()?;
result.connect(path)?;
Ok(result)
}

/// Checks for file descriptors passed by the service manager for socket
Expand All @@ -261,21 +260,8 @@ fn connect_notify_socket(unset_env: bool) -> io::Result<Option<UnixDatagram>> {
/// ```no_run
/// let socket = sd_notify::listen_fds().map(|mut fds| fds.next().expect("missing fd"));
/// ```
pub fn listen_fds() -> io::Result<impl Iterator<Item = RawFd>> {
listen_fds_internal(true)
}

fn listen_fds_internal(unset_env: bool) -> io::Result<impl ExactSizeIterator<Item = RawFd>> {
let _guard1 = UnsetEnvGuard {
name: "LISTEN_PID",
unset_env,
};
let _guard2 = UnsetEnvGuard {
name: "LISTEN_FDS",
unset_env,
};

let listen_pid = if let Ok(pid) = env::var("LISTEN_PID") {
pub fn listen_fds() -> io::Result<impl ExactSizeIterator<Item = RawFd>> {
let listen_pid = if let Ok(pid) = env::var(LISTEN_PID) {
pid
} else {
return Ok(0..0);
Expand All @@ -287,7 +273,7 @@ fn listen_fds_internal(unset_env: bool) -> io::Result<impl ExactSizeIterator<Ite
return Ok(0..0);
}

let listen_fds = if let Ok(fds) = env::var("LISTEN_FDS") {
let listen_fds = if let Ok(fds) = env::var(LISTEN_FDS) {
fds
} else {
return Ok(0..0);
Expand All @@ -311,46 +297,45 @@ fn listen_fds_internal(unset_env: bool) -> io::Result<impl ExactSizeIterator<Ite
Ok(listen_fds)
}

pub unsafe fn listen_fds_and_unset_env() -> io::Result<impl ExactSizeIterator<Item = RawFd>> {
let result = listen_fds();
unsafe {
env::remove_var(LISTEN_PID);
env::remove_var(LISTEN_FDS);
}
result
}
/// Checks for file descriptors passed by the service manager for socket
/// activation.
///
/// The function returns an iterator over file descriptors, starting from
/// `SD_LISTEN_FDS_START`. The number of descriptors is obtained from the
/// `LISTEN_FDS` environment variable.
///
/// The `unset_env` parameter should generally not be set, see the Safety section.
/// If the `unset_env` parameter is set, the `LISTEN_PID`, `LISTEN_FDS` and
/// `LISTEN_FDNAMES` environment variable will be unset before returning.
/// Child processes will not see the fdnames passed to this process. This is
/// usually not necessary, as a process should only use the `LISTEN_FDS`
/// variable if it is the PID given in `LISTEN_PID`.
///
/// Before returning, the file descriptors are set as `O_CLOEXEC`.
///
/// See [`sd_listen_fds_with_names(3)`][sd_listen_fds_with_names] for details.
///
/// [sd_listen_fds_with_names]: https://www.freedesktop.org/software/systemd/man/sd_listen_fds.html
///
/// # Safety
///
/// Although this function is not marked as `unsafe`, it is `unsafe` to call
/// with `unset_env` set to true. See the documentation on
/// [std::env::remove_var] for further information on this unsafety.
///
/// # Example
///
/// ```no_run
/// let socket = sd_notify::listen_fds().map(|mut fds| fds.next().expect("missing fd"));
/// ```
pub fn listen_fds_with_names(
unset_env: bool,
) -> io::Result<impl ExactSizeIterator<Item = (RawFd, String)>> {
let listen_fds = listen_fds_internal(unset_env)?;
let _guard = UnsetEnvGuard {
name: "LISTEN_FDNAMES",
unset_env,
};
zip_fds_with_names(listen_fds, env::var("LISTEN_FDNAMES").ok())
pub fn listen_fds_with_names() -> io::Result<impl ExactSizeIterator<Item = (RawFd, String)>> {
let listen_fds = listen_fds()?;
zip_fds_with_names(listen_fds, env::var(LISTEN_FDNAMES).ok())
}

pub unsafe fn listen_fds_with_names_and_unset_env() -> io::Result<impl ExactSizeIterator<Item = (RawFd, String)>> {
let result = listen_fds_with_names();
unsafe {
env::remove_var(LISTEN_PID);
env::remove_var(LISTEN_FDS);
env::remove_var(LISTEN_FDNAMES);
}
result
}

/// Internal helper that is independent of listen_fds function, for testing purposes.
Expand Down Expand Up @@ -386,19 +371,6 @@ fn zip_fds_with_names(
}
}

struct UnsetEnvGuard {
name: &'static str,
unset_env: bool,
}

impl Drop for UnsetEnvGuard {
fn drop(&mut self) {
if self.unset_env {
env::remove_var(self.name);
}
}
}

fn fd_cloexec(fd: u32) -> io::Result<()> {
let fd = RawFd::try_from(fd).map_err(|_| io::Error::from_raw_os_error(ffi::EBADF))?;
let flags = unsafe { ffi::fcntl(fd, ffi::F_GETFD, 0) };
Expand All @@ -417,21 +389,10 @@ fn fd_cloexec(fd: u32) -> io::Result<()> {

/// Asks the service manager for enabled watchdog.
///
/// The `unset_env` parameter should generally not be set, see the Safety section.
/// If the `unset_env` parameter is set, the `WATCHDOG_USEC` and `WATCHDOG_PID` environment variables
/// will be unset before returning. Further calls to `watchdog_enabled` will fail, but
/// child processes will no longer inherit the variable.
///
/// See [`sd_watchdog_enabled(3)`][sd_watchdog_enabled] for details.
///
/// [sd_watchdog_enabled]: https://www.freedesktop.org/software/systemd/man/sd_watchdog_enabled.html
///
/// # Safety
///
/// Although this function is not marked as `unsafe`, it is `unsafe` to call
/// with `unset_env` set to true. See the documentation on
/// [std::env::remove_var] for further information on this unsafety.
///
/// # Example
///
/// ```no_run
Expand All @@ -440,26 +401,11 @@ fn fd_cloexec(fd: u32) -> io::Result<()> {
/// let mut usec = 0;
/// let enabled = sd_notify::watchdog_enabled(true, &mut usec);
/// ```
pub fn watchdog_enabled(unset_env: bool, usec: &mut u64) -> bool {
struct Guard {
unset_env: bool,
}

impl Drop for Guard {
fn drop(&mut self) {
if self.unset_env {
env::remove_var("WATCHDOG_USEC");
env::remove_var("WATCHDOG_PID");
}
}
}

let _guard = Guard { unset_env };

let s = env::var("WATCHDOG_USEC")
pub fn watchdog_enabled(usec: &mut u64) -> bool {
let s = env::var(WATCHDOG_USEC)
.ok()
.and_then(|s| u64::from_str(&s).ok());
let p = env::var("WATCHDOG_PID")
let p = env::var(WATCHDOG_PID)
.ok()
.and_then(|s| u32::from_str(&s).ok());

Expand All @@ -472,6 +418,15 @@ pub fn watchdog_enabled(unset_env: bool, usec: &mut u64) -> bool {
}
}

pub unsafe fn watchdog_enabled_and_unset_env(usec: &mut u64) -> bool {
let result = watchdog_enabled(usec);
unsafe {
env::remove_var(WATCHDOG_USEC);
env::remove_var(WATCHDOG_PID);
}
result
}

#[cfg(test)]
mod tests {
use super::NotifyState;
Expand Down

0 comments on commit d9ed7a9

Please sign in to comment.