restart_manager/
application_restart.rs

1//! Process-global application restart registration.
2
3use std::ffi::{OsStr, OsString};
4use std::fmt;
5use std::sync::atomic::{AtomicBool, Ordering};
6
7use crate::input::contains_nul;
8use crate::session::map_sys_error;
9use crate::{Error, ErrorKind, Result};
10
11const MAX_RESTART_ARGUMENT_UNITS: usize = 1024;
12static REGISTRATION_OWNED: AtomicBool = AtomicBool::new(false);
13
14/// Command-line arguments and positive restart policies for this application.
15///
16/// The executable name must not be included. The default permits restart
17/// after every supported reason.
18#[derive(Clone, PartialEq, Eq)]
19pub struct ApplicationRestartOptions {
20    arguments: OsString,
21    restart_on_crash: bool,
22    restart_on_hang: bool,
23    restart_on_update: bool,
24    restart_on_reboot: bool,
25}
26
27impl Default for ApplicationRestartOptions {
28    fn default() -> Self {
29        Self {
30            arguments: OsString::new(),
31            restart_on_crash: true,
32            restart_on_hang: true,
33            restart_on_update: true,
34            restart_on_reboot: true,
35        }
36    }
37}
38
39impl fmt::Debug for ApplicationRestartOptions {
40    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
41        formatter
42            .debug_struct("ApplicationRestartOptions")
43            .field("arguments", &"<redacted>")
44            .field("restart_on_crash", &self.restart_on_crash)
45            .field("restart_on_hang", &self.restart_on_hang)
46            .field("restart_on_update", &self.restart_on_update)
47            .field("restart_on_reboot", &self.restart_on_reboot)
48            .finish()
49    }
50}
51
52impl ApplicationRestartOptions {
53    /// Creates options with the supplied restart arguments and all reasons enabled.
54    #[must_use]
55    pub fn new(arguments: impl Into<OsString>) -> Self {
56        Self {
57            arguments: arguments.into(),
58            ..Self::default()
59        }
60    }
61
62    /// Replaces the restart arguments.
63    #[must_use]
64    pub fn with_arguments(mut self, arguments: impl Into<OsString>) -> Self {
65        self.arguments = arguments.into();
66        self
67    }
68
69    /// Chooses whether restart is allowed after an unhandled crash.
70    #[must_use]
71    pub const fn with_restart_on_crash(mut self, enabled: bool) -> Self {
72        self.restart_on_crash = enabled;
73        self
74    }
75
76    /// Chooses whether restart is allowed after an application hang.
77    #[must_use]
78    pub const fn with_restart_on_hang(mut self, enabled: bool) -> Self {
79        self.restart_on_hang = enabled;
80        self
81    }
82
83    /// Chooses whether restart is allowed after an update.
84    #[must_use]
85    pub const fn with_restart_on_update(mut self, enabled: bool) -> Self {
86        self.restart_on_update = enabled;
87        self
88    }
89
90    /// Chooses whether restart is allowed after a system reboot.
91    #[must_use]
92    pub const fn with_restart_on_reboot(mut self, enabled: bool) -> Self {
93        self.restart_on_reboot = enabled;
94        self
95    }
96
97    /// Returns whether crash restart is allowed.
98    #[must_use]
99    pub const fn restart_on_crash(&self) -> bool {
100        self.restart_on_crash
101    }
102
103    /// Returns whether hang restart is allowed.
104    #[must_use]
105    pub const fn restart_on_hang(&self) -> bool {
106        self.restart_on_hang
107    }
108
109    /// Returns whether update restart is allowed.
110    #[must_use]
111    pub const fn restart_on_update(&self) -> bool {
112        self.restart_on_update
113    }
114
115    /// Returns whether reboot restart is allowed.
116    #[must_use]
117    pub const fn restart_on_reboot(&self) -> bool {
118        self.restart_on_reboot
119    }
120
121    fn native_flags(&self) -> u32 {
122        ((!self.restart_on_crash) as u32)
123            | (((!self.restart_on_hang) as u32) << 1)
124            | (((!self.restart_on_update) as u32) << 2)
125            | (((!self.restart_on_reboot) as u32) << 3)
126    }
127}
128
129/// An exclusive crate-managed lease for the current process registration.
130pub struct ApplicationRestartRegistration {
131    options: ApplicationRestartOptions,
132    armed: bool,
133}
134
135impl fmt::Debug for ApplicationRestartRegistration {
136    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
137        formatter
138            .debug_struct("ApplicationRestartRegistration")
139            .field("options", &self.options)
140            .field("armed", &self.armed)
141            .finish()
142    }
143}
144
145impl ApplicationRestartRegistration {
146    /// Registers this process for application restart.
147    pub fn register(options: ApplicationRestartOptions) -> Result<Self> {
148        validate_arguments(&options.arguments)?;
149        if REGISTRATION_OWNED
150            .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
151            .is_err()
152        {
153            return Err(Error::new(
154                ErrorKind::ApplicationRestartInUse,
155                None,
156                "this process already owns an application restart registration",
157            ));
158        }
159        if let Err(error) =
160            crate::sys::register_application_restart(&options.arguments, options.native_flags())
161        {
162            REGISTRATION_OWNED.store(false, Ordering::Release);
163            return Err(map_sys_error(error));
164        }
165        Ok(Self {
166            options,
167            armed: true,
168        })
169    }
170
171    /// Updates the command line and policies without releasing the lease.
172    pub fn update(&mut self, options: ApplicationRestartOptions) -> Result<()> {
173        validate_arguments(&options.arguments)?;
174        crate::sys::register_application_restart(&options.arguments, options.native_flags())
175            .map_err(map_sys_error)?;
176        self.options = options;
177        Ok(())
178    }
179
180    /// Explicitly unregisters this process.
181    ///
182    /// On failure the value remains armed and its destructor makes one
183    /// best-effort retry.
184    pub fn unregister(mut self) -> Result<()> {
185        match crate::sys::unregister_application_restart() {
186            Ok(()) => {
187                self.disarm();
188                Ok(())
189            }
190            Err(error) => Err(map_sys_error(error)),
191        }
192    }
193
194    fn disarm(&mut self) {
195        self.armed = false;
196        REGISTRATION_OWNED.store(false, Ordering::Release);
197    }
198}
199
200impl Drop for ApplicationRestartRegistration {
201    fn drop(&mut self) {
202        if self.armed {
203            let _ = crate::sys::unregister_application_restart();
204            self.disarm();
205        }
206    }
207}
208
209fn validate_arguments(arguments: &OsStr) -> Result<()> {
210    if contains_nul(arguments) {
211        return Err(Error::new(
212            ErrorKind::InvalidInput,
213            None,
214            "application restart arguments may not contain an embedded NUL",
215        ));
216    }
217    if utf16_unit_count(arguments) > MAX_RESTART_ARGUMENT_UNITS {
218        return Err(Error::new(
219            ErrorKind::InvalidInput,
220            None,
221            "application restart arguments exceed 1024 UTF-16 code units",
222        ));
223    }
224    Ok(())
225}
226
227#[cfg(windows)]
228fn utf16_unit_count(value: &OsStr) -> usize {
229    use std::os::windows::ffi::OsStrExt;
230    value.encode_wide().count()
231}
232
233#[cfg(not(windows))]
234fn utf16_unit_count(value: &OsStr) -> usize {
235    value.to_string_lossy().encode_utf16().count()
236}
237
238#[cfg(test)]
239mod tests {
240    use super::*;
241
242    #[test]
243    fn options_use_positive_policies_and_redact_arguments() {
244        let options = ApplicationRestartOptions::new("secret")
245            .with_restart_on_crash(false)
246            .with_restart_on_reboot(false);
247        assert!(!options.restart_on_crash());
248        assert!(options.restart_on_hang());
249        assert!(options.restart_on_update());
250        assert!(!options.restart_on_reboot());
251        assert_eq!(options.native_flags(), 0x09);
252        let debug = format!("{options:?}");
253        assert!(debug.contains("<redacted>"));
254        assert!(!debug.contains("secret"));
255    }
256
257    #[test]
258    fn arguments_validate_nul_and_utf16_limit() {
259        assert!(validate_arguments(OsStr::new("bad\0argument")).is_err());
260        assert!(validate_arguments(OsStr::new(&"x".repeat(1024))).is_ok());
261        assert!(validate_arguments(OsStr::new(&"x".repeat(1025))).is_err());
262    }
263
264    #[cfg(windows)]
265    #[test]
266    fn all_builders_registration_debug_and_drop_are_exercised() {
267        let options = ApplicationRestartOptions::default()
268            .with_arguments("--updated")
269            .with_restart_on_hang(false)
270            .with_restart_on_update(false)
271            .with_restart_on_reboot(false);
272        assert!(options.restart_on_crash());
273        assert!(!options.restart_on_hang());
274        assert!(!options.restart_on_update());
275        assert!(!options.restart_on_reboot());
276
277        let registration = ApplicationRestartRegistration::register(options).unwrap();
278        let debug = format!("{registration:?}");
279        assert!(debug.contains("ApplicationRestartRegistration"));
280        assert!(debug.contains("<redacted>"));
281        drop(registration);
282
283        ApplicationRestartRegistration::register(ApplicationRestartOptions::default())
284            .unwrap()
285            .unregister()
286            .unwrap();
287    }
288}
restart_manager/application_restart.rs

restart_manager/
application_restart.rs
