fix: preserve macos services across logout (#4202)

refs #4100

4 files changed, 328 insertions, 2 deletions

diff --git a/src/platform/macos.rs b/src/platform/macos.rs
index 4b11a618..8c529342 100644
--- a/src/platform/macos.rs
+++ b/src/platform/macos.rs
@@ -23,6 +23,9 @@ pub(crate) use super::unix_common::{
     ClientStreamReader, StatusCommandGuard,
 };
 
+mod bootstrap;
+pub(crate) use bootstrap::{configure_server_daemon_context, prepare_server_process};
+
 #[cfg(test)]
 mod config_file_tests;
 
diff --git a/src/platform/macos/bootstrap.rs b/src/platform/macos/bootstrap.rs
new file mode 100644
index 00000000..8cd333a2
--- /dev/null
+++ b/src/platform/macos/bootstrap.rs
@@ -0,0 +1,306 @@
+use std::ffi::{CStr, OsStr};
+use std::io;
+use std::process::Command;
+
+use libc::{c_char, c_int, c_void, mach_port_t, uid_t};
+
+const SERVER_CONTEXT_ENV: &str = "HERDR_MACOS_SERVER_CONTEXT";
+const USER_CONTEXT: &str = "user";
+const TASK_BOOTSTRAP_PORT: c_int = 4;
+
+type GetRoot = unsafe extern "C" fn(mach_port_t, *mut mach_port_t) -> c_int;
+type LookupUser =
+    unsafe extern "C" fn(mach_port_t, *const c_char, uid_t, *mut mach_port_t) -> c_int;
+type SetSpecialPort = unsafe extern "C" fn(mach_port_t, c_int, mach_port_t) -> c_int;
+type Deallocate = unsafe extern "C" fn(mach_port_t, mach_port_t) -> c_int;
+
+unsafe extern "C" {
+    static mut mach_task_self_: mach_port_t;
+    fn task_set_special_port(task: mach_port_t, which: c_int, port: mach_port_t) -> c_int;
+    fn mach_port_deallocate(task: mach_port_t, port: mach_port_t) -> c_int;
+}
+
+pub(crate) fn configure_server_daemon_context(command: &mut Command) {
+    // Handoff inherits the source context, including intentionally direct launches.
+    // Older replacements also cannot consume this marker and would leak it.
+    if command.get_args().any(|arg| arg == "--handoff-import") {
+        command.env_remove(SERVER_CONTEXT_ENV);
+    } else {
+        command.env(SERVER_CONTEXT_ENV, USER_CONTEXT);
+    }
+}
+
+/// Called after exec, before logging, worker threads, or pane creation.
+pub(crate) fn prepare_server_process(handoff_import: bool) -> io::Result<bool> {
+    let requested = std::env::var_os(SERVER_CONTEXT_ENV);
+    // This is a one-shot launch request, not environment for panes or plugins.
+    std::env::remove_var(SERVER_CONTEXT_ENV);
+    if !needs_user_context(requested.as_deref(), handoff_import) {
+        return Ok(false);
+    }
+    let api = BootstrapApi::load()?;
+    api.adopt_user_context()?;
+    Ok(true)
+}
+
+fn needs_user_context(requested: Option<&OsStr>, handoff_import: bool) -> bool {
+    // A replacement must preserve its source's policy, not silently rehome a
+    // direct server. Fixed sources already pass on the durable context by fork.
+    !handoff_import && requested == Some(OsStr::new(USER_CONTEXT))
+}
+
+struct BootstrapApi {
+    bootstrap: *mut mach_port_t,
+    task: mach_port_t,
+    get_root: GetRoot,
+    lookup_user: LookupUser,
+    set_special_port: SetSpecialPort,
+    deallocate: Deallocate,
+}
+
+impl BootstrapApi {
+    fn load() -> io::Result<Self> {
+        // These bootstrap interfaces are private macOS APIs, also used by tmux.
+        // Resolve all symbols before changing anything so their removal cannot
+        // prevent Herdr from launching with its inherited context.
+        let bootstrap = symbol(c"bootstrap_port")?.cast::<mach_port_t>();
+        let get_root = symbol(c"bootstrap_get_root")?;
+        let lookup_user = symbol(c"bootstrap_look_up_per_user")?;
+        Ok(Self {
+            bootstrap,
+            task: unsafe { mach_task_self_ },
+            get_root: unsafe { std::mem::transmute::<*mut c_void, GetRoot>(get_root) },
+            lookup_user: unsafe { std::mem::transmute::<*mut c_void, LookupUser>(lookup_user) },
+            set_special_port: task_set_special_port,
+            deallocate: mach_port_deallocate,
+        })
+    }
+
+    fn adopt_user_context(&self) -> io::Result<()> {
+        let inherited = unsafe { *self.bootstrap };
+        let mut root = SendRight::new(self);
+        check("bootstrap_get_root", unsafe {
+            (self.get_root)(inherited, &mut root.port)
+        })?;
+        root.validate()?;
+
+        let mut user = SendRight::new(self);
+        check("bootstrap_look_up_per_user", unsafe {
+            (self.lookup_user)(root.port, std::ptr::null(), libc::getuid(), &mut user.port)
+        })?;
+        user.validate()?;
+
+        // setsid() does not detach from Aqua's Mach bootstrap namespace. Select
+        // the per-user endpoint before any children inherit the logout-scoped one.
+        check("task_set_bootstrap_port", unsafe {
+            (self.set_special_port)(self.task, TASK_BOOTSTRAP_PORT, user.port)
+        })?;
+        unsafe {
+            *self.bootstrap = user.port;
+        }
+        // The global now owns the lookup's send right, in addition to the task's
+        // retained reference. Only release the old global after both are updated.
+        user.port = 0;
+        drop(SendRight {
+            api: self,
+            port: inherited,
+        });
+        Ok(())
+    }
+}
+
+fn symbol(name: &CStr) -> io::Result<*mut c_void> {
+    let address = unsafe { libc::dlsym(libc::RTLD_DEFAULT, name.as_ptr()) };
+    if address.is_null() {
+        Err(io::Error::new(
+            io::ErrorKind::Unsupported,
+            format!(
+                "macOS service-context API unavailable: {}",
+                name.to_string_lossy()
+            ),
+        ))
+    } else {
+        Ok(address)
+    }
+}
+
+fn check(operation: &str, status: c_int) -> io::Result<()> {
+    if status == 0 {
+        Ok(())
+    } else {
+        Err(io::Error::other(format!(
+            "{operation} failed (Mach error {status})"
+        )))
+    }
+}
+
+struct SendRight<'a> {
+    api: &'a BootstrapApi,
+    port: mach_port_t,
+}
+
+impl<'a> SendRight<'a> {
+    fn new(api: &'a BootstrapApi) -> Self {
+        Self { api, port: 0 }
+    }
+
+    fn validate(&self) -> io::Result<()> {
+        if self.port == 0 || self.port == mach_port_t::MAX {
+            Err(io::Error::other("macOS returned an invalid bootstrap port"))
+        } else {
+            Ok(())
+        }
+    }
+}
+
+impl Drop for SendRight<'_> {
+    fn drop(&mut self) {
+        if self.port != 0 && self.port != mach_port_t::MAX {
+            unsafe {
+                (self.api.deallocate)(self.api.task, self.port);
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::cell::RefCell;
+
+    #[derive(Default)]
+    struct Calls {
+        fail: Option<&'static str>,
+        installed: Vec<mach_port_t>,
+        released: Vec<mach_port_t>,
+    }
+
+    thread_local! {
+        static CALLS: RefCell<Calls> = RefCell::default();
+    }
+
+    unsafe extern "C" fn get_root(_: mach_port_t, out: *mut mach_port_t) -> c_int {
+        if CALLS.with(|calls| calls.borrow().fail == Some("root")) {
+            return 5;
+        }
+        unsafe { *out = 20 };
+        0
+    }
+
+    unsafe extern "C" fn lookup_user(
+        _: mach_port_t,
+        _: *const c_char,
+        _: uid_t,
+        out: *mut mach_port_t,
+    ) -> c_int {
+        let failure = CALLS.with(|calls| calls.borrow().fail);
+        if failure == Some("lookup") {
+            return 5;
+        }
+        unsafe { *out = if failure == Some("invalid") { 0 } else { 30 } };
+        0
+    }
+
+    unsafe extern "C" fn set_special_port(_: mach_port_t, _: c_int, port: mach_port_t) -> c_int {
+        CALLS.with(|calls| {
+            let mut calls = calls.borrow_mut();
+            if calls.fail == Some("install") {
+                5
+            } else {
+                calls.installed.push(port);
+                0
+            }
+        })
+    }
+
+    unsafe extern "C" fn deallocate(_: mach_port_t, port: mach_port_t) -> c_int {
+        CALLS.with(|calls| calls.borrow_mut().released.push(port));
+        0
+    }
+
+    fn fake_api(bootstrap: &mut mach_port_t) -> BootstrapApi {
+        BootstrapApi {
+            bootstrap,
+            task: 1,
+            get_root,
+            lookup_user,
+            set_special_port,
+            deallocate,
+        }
+    }
+
+    #[test]
+    fn adoption_transfers_user_right_and_releases_old_and_root_rights() {
+        CALLS.with(|calls| *calls.borrow_mut() = Calls::default());
+        let mut bootstrap = 10;
+        fake_api(&mut bootstrap).adopt_user_context().unwrap();
+        assert_eq!(bootstrap, 30);
+        CALLS.with(|calls| {
+            let calls = calls.borrow();
+            assert_eq!(calls.installed, [30]);
+            assert_eq!(calls.released, [10, 20]);
+        });
+    }
+
+    #[test]
+    fn failed_adoption_keeps_inherited_context_and_releases_temporary_rights() {
+        for (failure, released) in [
+            ("root", vec![]),
+            ("lookup", vec![20]),
+            ("invalid", vec![20]),
+            ("install", vec![30, 20]),
+        ] {
+            CALLS.with(|calls| {
+                *calls.borrow_mut() = Calls {
+                    fail: Some(failure),
+                    ..Calls::default()
+                }
+            });
+            let mut bootstrap = 10;
+            assert!(fake_api(&mut bootstrap).adopt_user_context().is_err());
+            assert_eq!(bootstrap, 10, "{failure}");
+            CALLS.with(|calls| {
+                let calls = calls.borrow();
+                assert!(calls.installed.is_empty(), "{failure}");
+                assert_eq!(calls.released, released, "{failure}");
+            });
+        }
+    }
+
+    #[test]
+    fn missing_private_symbol_is_a_recoverable_error() {
+        let error = symbol(c"herdr_nonexistent_bootstrap_symbol").unwrap_err();
+        assert_eq!(error.kind(), io::ErrorKind::Unsupported);
+    }
+
+    #[test]
+    fn direct_and_handoff_servers_keep_inherited_context() {
+        assert!(!needs_user_context(None, false));
+        assert!(!needs_user_context(Some(OsStr::new("unknown")), false));
+        assert!(needs_user_context(Some(OsStr::new(USER_CONTEXT)), false));
+        assert!(!needs_user_context(None, true));
+        assert!(!needs_user_context(Some(OsStr::new(USER_CONTEXT)), true));
+    }
+
+    #[test]
+    fn handoff_commands_remove_even_an_inherited_launch_marker() {
+        let mut command = Command::new("herdr");
+        command.args(["server", "--handoff-import", "socket", "token"]);
+        command.env(SERVER_CONTEXT_ENV, USER_CONTEXT);
+        crate::platform::detach_server_daemon_command(&mut command);
+        assert!(command
+            .get_envs()
+            .any(|(key, value)| key == SERVER_CONTEXT_ENV && value.is_none()));
+    }
+
+    #[test]
+    fn only_detached_server_commands_request_user_context() {
+        let direct = Command::new("herdr");
+        assert!(!direct.get_envs().any(|(key, _)| key == SERVER_CONTEXT_ENV));
+        let mut detached = Command::new("herdr");
+        crate::platform::detach_server_daemon_command(&mut detached);
+        assert!(detached.get_envs().any(|(key, value)| {
+            key == SERVER_CONTEXT_ENV && value == Some(OsStr::new(USER_CONTEXT))
+        }));
+    }
+}
diff --git a/src/platform/mod.rs b/src/platform/mod.rs
index 70047fef..eb84e0b1 100644
--- a/src/platform/mod.rs
+++ b/src/platform/mod.rs
@@ -153,10 +153,18 @@ pub fn launch_server_daemon_command(command: &mut std::process::Command) -> std:
     command.spawn().map(|child| child.id())
 }
 
+#[cfg(not(target_os = "macos"))]
+pub(crate) fn prepare_server_process(_handoff_import: bool) -> std::io::Result<bool> {
+    Ok(false)
+}
+
 #[cfg(any(target_os = "linux", target_os = "macos"))]
 pub fn detach_server_daemon_command(command: &mut std::process::Command) {
     use std::os::unix::process::CommandExt;
 
+    #[cfg(target_os = "macos")]
+    macos::configure_server_daemon_context(command);
+
     unsafe {
         command.pre_exec(|| {
             if libc::setsid() < 0 {
diff --git a/src/server/headless/bootstrap.rs b/src/server/headless/bootstrap.rs
index 638b1b2d..e5ff400d 100644
--- a/src/server/headless/bootstrap.rs
+++ b/src/server/headless/bootstrap.rs
@@ -2,11 +2,20 @@ use super::*;
 
 /// Run the headless server. This is the entry point called from main.rs.
 pub fn run_server() -> io::Result<()> {
+    let args: Vec<String> = std::env::args().collect();
+    let handoff_import = args.get(2).map(String::as_str) == Some("--handoff-import");
+    let process_context = crate::platform::prepare_server_process(handoff_import);
     init_logging();
+    match process_context {
+        Ok(true) => info!("server using persistent user service context"),
+        Ok(false) => {}
+        Err(err) => {
+            warn!(%err, "could not select persistent user service context; retaining inherited context")
+        }
+    }
     crate::platform::raise_server_nofile_limit();
 
-    let args: Vec<String> = std::env::args().collect();
-    if args.get(2).map(String::as_str) == Some("--handoff-import") {
+    if handoff_import {
         let socket_path = args
             .get(3)
             .map(PathBuf::from)