1 files changed, 361 insertions, 0 deletions
diff --git a/src/platform/ssh_agent.rs b/src/platform/ssh_agent.rs
new file mode 100644
index 00000000..9caddd73
--- /dev/null
+++ b/src/platform/ssh_agent.rs
@@ -0,0 +1,361 @@
+//! Session-local indirection for SSH agents whose sockets belong to an attachment.
+
+use std::fs;
+use std::io;
+use std::os::unix::fs::{symlink, FileTypeExt, MetadataExt};
+use std::path::{Path, PathBuf};
+use std::sync::{Arc, Mutex};
+use std::time::{Duration, Instant};
+
+use interprocess::local_socket::{ConnectOptions, GenericFilePath, ToFsName};
+use interprocess::ConnectWaitMode;
+
+use crate::ipc::LocalStream;
+
+const PROBE_INTERVAL: Duration = Duration::from_secs(1);
+
+#[derive(Clone)]
+pub(crate) struct SshAgentRegistry(Arc<Mutex<State>>);
+
+struct State {
+    path: PathBuf,
+    fallback: Option<PathBuf>,
+    agents: Vec<(u64, PathBuf)>,
+    next_id: u64,
+    identity: Option<(u64, u64)>,
+    last_probe: Option<Instant>,
+}
+
+pub(crate) struct SshAgentLease {
+    registry: SshAgentRegistry,
+    id: u64,
+}
+
+pub(crate) fn socket_path() -> PathBuf {
+    agent_path_for(&crate::api::socket_path())
+}
+
+fn agent_path_for(api_path: &Path) -> PathBuf {
+    let mut path = api_path.as_os_str().to_os_string();
+    path.push(".agent");
+    path.into()
+}
+
+fn usable_socket(path: &Path) -> bool {
+    fs::metadata(path).is_ok_and(|metadata| {
+        // The API is user-private; do not redirect that user's panes to another user's agent.
+        metadata.file_type().is_socket() && metadata.uid() == unsafe { libc::geteuid() }
+    })
+}
+
+fn live_socket(path: &Path) -> bool {
+    if !usable_socket(path) {
+        return false;
+    }
+    let Ok(name) = path.to_fs_name::<GenericFilePath>() else {
+        return false;
+    };
+    // Never wait for a full accept queue or retain a forwarded SSH channel after probing.
+    let Ok(LocalStream::UdSocket(stream)) = ConnectOptions::new()
+        .name(name)
+        .wait_mode(ConnectWaitMode::Timeout(Duration::ZERO))
+        .nonblocking_stream(true)
+        .connect_sync()
+    else {
+        return false;
+    };
+    // Linux can report an unconnected socket writable after connect returns EAGAIN.
+    stream.inner().peer_addr().is_ok()
+}
+
+impl SshAgentRegistry {
+    pub(crate) fn new(path: PathBuf, inherited: Option<PathBuf>) -> io::Result<Self> {
+        let inherited = inherited.filter(|path| !path.as_os_str().is_empty());
+        let managed = inherited.is_some()
+            || fs::symlink_metadata(&path).is_ok_and(|metadata| metadata.file_type().is_symlink());
+        // A handoff keeps existing pane environments, including this stable pathname.
+        let fallback = fs::read_link(&path)
+            .ok()
+            .filter(|target| target != &path && usable_socket(target))
+            .or_else(|| inherited.filter(|target| target != &path && usable_socket(target)));
+        let mut state = State {
+            path,
+            fallback,
+            agents: Vec::new(),
+            next_id: 0,
+            identity: None,
+            last_probe: None,
+        };
+        if managed {
+            state.publish()?;
+        }
+        Ok(Self(Arc::new(Mutex::new(state))))
+    }
+
+    pub(crate) fn register(&self, path: PathBuf) -> io::Result<SshAgentLease> {
+        let mut state = self
+            .0
+            .lock()
+            .map_err(|_| io::Error::other("SSH agent registry poisoned"))?;
+        if !path.is_absolute() || path == state.path || !usable_socket(&path) {
+            return Err(io::Error::new(
+                io::ErrorKind::InvalidInput,
+                "SSH agent must be an absolute, user-owned socket",
+            ));
+        }
+        // Canonicalizing /var to /private/var can exceed macOS's Unix socket path limit.
+        let id = state.next_id;
+        state.next_id += 1;
+        state.agents.push((id, path));
+        if let Err(error) = state.publish() {
+            state.agents.retain(|(candidate, _)| *candidate != id);
+            return Err(error);
+        }
+        Ok(SshAgentLease {
+            registry: self.clone(),
+            id,
+        })
+    }
+}
+
+impl State {
+    fn publish(&mut self) -> io::Result<()> {
+        if let Some(identity) = self.identity {
+            let metadata = fs::symlink_metadata(&self.path)?;
+            if identity != (metadata.dev(), metadata.ino()) {
+                return Err(io::Error::other(
+                    "SSH agent address belongs to a replacement server",
+                ));
+            }
+        }
+        // Keep a working agent rather than letting probes or a second client replace it.
+        let unavailable = self.path.with_extension("unavailable");
+        self.last_probe = Some(Instant::now());
+        let target = self
+            .fallback
+            .as_deref()
+            .filter(|path| live_socket(path))
+            .or_else(|| {
+                self.agents
+                    .iter()
+                    .map(|(_, path)| path.as_path())
+                    .find(|path| live_socket(path))
+            })
+            .unwrap_or(&unavailable);
+        if self.identity.is_some() && fs::read_link(&self.path).ok().as_deref() == Some(target) {
+            return Ok(());
+        }
+        let temporary = self
+            .path
+            .with_extension(format!("{}.new", std::process::id()));
+        symlink(target, &temporary)?;
+        if let Err(error) = fs::rename(&temporary, &self.path) {
+            let _ = fs::remove_file(&temporary);
+            return Err(error);
+        }
+        let metadata = fs::symlink_metadata(&self.path)?;
+        self.identity = Some((metadata.dev(), metadata.ino()));
+        Ok(())
+    }
+}
+
+impl Drop for State {
+    fn drop(&mut self) {
+        if fs::symlink_metadata(&self.path)
+            .is_ok_and(|metadata| self.identity == Some((metadata.dev(), metadata.ino())))
+        {
+            let _ = fs::remove_file(&self.path);
+        }
+    }
+}
+
+impl SshAgentLease {
+    pub(crate) fn refresh(&self) -> io::Result<()> {
+        self.refresh_at(Instant::now())
+    }
+
+    fn refresh_at(&self, now: Instant) -> io::Result<()> {
+        let mut state = self
+            .registry
+            .0
+            .lock()
+            .map_err(|_| io::Error::other("SSH agent registry poisoned"))?;
+        // Share the probe budget across attachments, not one SSH channel per polling client.
+        if state
+            .last_probe
+            .is_none_or(|last| now.saturating_duration_since(last) >= PROBE_INTERVAL)
+        {
+            state.publish()?;
+        }
+        Ok(())
+    }
+}
+
+impl Drop for SshAgentLease {
+    fn drop(&mut self) {
+        if let Ok(mut state) = self.registry.0.lock() {
+            state.agents.retain(|(id, _)| *id != self.id);
+            if let Err(error) = state.publish() {
+                tracing::warn!(%error, "could not refresh SSH agent after attachment ended");
+            }
+        }
+    }
+}
+
+pub(crate) fn apply_pane_env(command: &mut portable_pty::CommandBuilder) {
+    let path = socket_path();
+    if fs::symlink_metadata(&path).is_ok_and(|metadata| metadata.file_type().is_symlink()) {
+        command.env("SSH_AUTH_SOCK", path);
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::os::unix::net::UnixListener;
+
+    #[test]
+    fn server_without_an_agent_leaves_local_pane_agent_setup_alone() {
+        let directory = std::env::temp_dir().join(format!("herdr-no-agent-{}", std::process::id()));
+        fs::create_dir(&directory).unwrap();
+        let stable = directory.join("agent");
+        for inherited in [None, Some(PathBuf::new())] {
+            let registry = SshAgentRegistry::new(stable.clone(), inherited).unwrap();
+            assert!(
+                fs::symlink_metadata(&stable).is_err(),
+                "a local server without an agent must not advertise an agent address to panes"
+            );
+            drop(registry);
+        }
+        fs::remove_dir_all(directory).unwrap();
+    }
+
+    #[test]
+    fn closed_agent_listener_does_not_block_a_live_replacement() {
+        let directory =
+            std::env::temp_dir().join(format!("herdr-dead-agent-{}", std::process::id()));
+        fs::create_dir(&directory).unwrap();
+        let stable = directory.join("agent");
+        let a = directory.join("a");
+        let b = directory.join("b");
+        let listener_a = UnixListener::bind(&a).unwrap();
+        let _listener_b = UnixListener::bind(&b).unwrap();
+        let registry = SshAgentRegistry::new(stable.clone(), Some(a.clone())).unwrap();
+        let (mut probe, _) = listener_a.accept().unwrap();
+        probe.set_nonblocking(true).unwrap();
+        assert_eq!(
+            std::io::Read::read(&mut probe, &mut [0]).unwrap(),
+            0,
+            "agent checks must not hold forwarded SSH channels open"
+        );
+        let lease_b = registry.register(b.clone()).unwrap();
+        assert_eq!(fs::read_link(&stable).unwrap(), a);
+        drop(listener_a);
+        assert!(fs::metadata(&a).unwrap().file_type().is_socket());
+        lease_b.refresh_at(Instant::now() + PROBE_INTERVAL).unwrap();
+        assert_eq!(fs::read_link(&stable).unwrap(), b);
+        drop(lease_b);
+        drop(registry);
+        fs::remove_dir_all(directory).unwrap();
+    }
+
+    #[cfg(target_os = "linux")]
+    #[test]
+    fn unestablished_agent_connection_does_not_block_a_live_replacement() {
+        use std::os::fd::AsRawFd;
+        use std::os::unix::net::UnixStream;
+        let directory =
+            std::env::temp_dir().join(format!("herdr-busy-agent-{}", std::process::id()));
+        fs::create_dir(&directory).unwrap();
+        let a = directory.join("a");
+        let b = directory.join("b");
+        let stable = directory.join("agent");
+        let listener_a = UnixListener::bind(&a).unwrap();
+        // Linux allows one queued connection with a zero backlog.
+        assert_eq!(unsafe { libc::listen(listener_a.as_raw_fd(), 0) }, 0);
+        let _queued = UnixStream::connect(&a).unwrap();
+        let _listener_b = UnixListener::bind(&b).unwrap();
+        let registry = SshAgentRegistry::new(stable.clone(), Some(a)).unwrap();
+        let lease_b = registry.register(b.clone()).unwrap();
+        assert_eq!(fs::read_link(&stable).unwrap(), b);
+        drop(lease_b);
+        drop(registry);
+        fs::remove_dir_all(directory).unwrap();
+    }
+
+    #[test]
+    fn registration_preserves_the_supplied_socket_address() {
+        let directory =
+            std::env::temp_dir().join(format!("herdr-agent-path-{}", std::process::id()));
+        fs::create_dir(&directory).unwrap();
+        symlink(".", directory.join("alias")).unwrap();
+        let _listener = UnixListener::bind(directory.join("upstream")).unwrap();
+        let supplied = directory.join("alias/upstream");
+        let stable = directory.join("agent");
+        let registry = SshAgentRegistry::new(stable.clone(), None).unwrap();
+        let lease = registry.register(supplied.clone()).unwrap();
+        assert_eq!(fs::read_link(&stable).unwrap(), supplied);
+        drop(lease);
+        drop(registry);
+        fs::remove_dir_all(directory).unwrap();
+    }
+
+    #[test]
+    fn socket_overrides_have_independent_agent_addresses() {
+        let directory =
+            std::env::temp_dir().join(format!("herdr-agent-overrides-{}", std::process::id()));
+        fs::create_dir(&directory).unwrap();
+        let a = directory.join("a");
+        let b = directory.join("b");
+        let _a_listener = UnixListener::bind(&a).unwrap();
+        let _b_listener = UnixListener::bind(&b).unwrap();
+        let stable_a = agent_path_for(&directory.join("first.sock"));
+        let stable_b = agent_path_for(&directory.join("second.sock"));
+        let registry_a = SshAgentRegistry::new(stable_a.clone(), Some(a.clone())).unwrap();
+        let registry_b = SshAgentRegistry::new(stable_b.clone(), Some(b.clone())).unwrap();
+        assert_eq!(fs::read_link(&stable_a).unwrap(), a);
+        assert_eq!(fs::read_link(&stable_b).unwrap(), b);
+        drop(registry_a);
+        assert_eq!(fs::read_link(&stable_b).unwrap(), b);
+        drop(registry_b);
+        fs::remove_dir_all(directory).unwrap();
+    }
+
+    #[test]
+    fn reconnect_and_overlapping_attachments_keep_a_stable_agent_address() {
+        let directory =
+            std::env::temp_dir().join(format!("herdr-ssh-agent-{}", std::process::id()));
+        fs::create_dir(&directory).unwrap();
+        let stable = directory.join("agent");
+        let a = directory.join("a");
+        let b = directory.join("b");
+        let probe = directory.join("probe");
+        let _a_listener = UnixListener::bind(&a).unwrap();
+        let _b_listener = UnixListener::bind(&b).unwrap();
+        let _probe_listener = UnixListener::bind(&probe).unwrap();
+        let registry = SshAgentRegistry::new(stable.clone(), Some(a.clone())).unwrap();
+        let temporary = registry.register(a.clone()).unwrap();
+        drop(temporary);
+        assert_eq!(fs::read_link(&stable).unwrap(), a);
+        let lease_a = registry.register(a.clone()).unwrap();
+        let lease_probe = registry.register(probe).unwrap();
+        assert_eq!(fs::read_link(&stable).unwrap(), a);
+        drop(lease_probe);
+        let lease_b = registry.register(b.clone()).unwrap();
+        assert_eq!(fs::read_link(&stable).unwrap(), a);
+        fs::remove_file(&a).unwrap();
+        lease_b.refresh_at(Instant::now() + PROBE_INTERVAL).unwrap();
+        assert_eq!(fs::read_link(&stable).unwrap(), b);
+        drop(lease_a);
+        assert_eq!(fs::read_link(&stable).unwrap(), b);
+        drop(lease_b);
+        assert!(!stable.exists());
+        assert!(fs::symlink_metadata(&stable).is_ok());
+        let _lease_b = registry.register(b.clone()).unwrap();
+        assert_eq!(fs::read_link(&stable).unwrap(), b);
+        assert!(registry.register(stable).is_err());
+        drop(_lease_b);
+        drop(registry);
+        fs::remove_dir_all(directory).unwrap();
+    }
+}