diff options
Diffstat (limited to 'src')
| -rw-r--r-- | src/api/schema.rs | 2 | ||||
| -rw-r--r-- | src/api/schema/server.rs | 9 | ||||
| -rw-r--r-- | src/api/schema/tests.rs | 1 | ||||
| -rw-r--r-- | src/api/server.rs | 136 | ||||
| -rw-r--r-- | src/app/api.rs | 7 | ||||
| -rw-r--r-- | src/cli/status.rs | 10 | ||||
| -rw-r--r-- | src/pane.rs | 2 | ||||
| -rw-r--r-- | src/platform/mod.rs | 3 | ||||
| -rw-r--r-- | src/platform/ssh_agent.rs | 361 | ||||
| -rw-r--r-- | src/remote.rs | 2 | ||||
| -rw-r--r-- | src/remote/host.rs | 2 | ||||
| -rw-r--r-- | src/remote/ssh_agent.rs | 239 | ||||
| -rw-r--r-- | src/update.rs | 3 |
13 files changed, 775 insertions, 2 deletions
diff --git a/src/api/schema.rs b/src/api/schema.rs index e6c678dc..d7e1a950 100644 --- a/src/api/schema.rs +++ b/src/api/schema.rs @@ -53,6 +53,8 @@ pub enum Method { ServerLiveHandoff(ServerLiveHandoffParams), #[serde(rename = "server.reload_config")] ServerReloadConfig(EmptyParams), + #[serde(rename = "server.ssh_agent.register")] + ServerSshAgentRegister(ServerSshAgentRegisterParams), #[serde(rename = "server.agent_manifests")] ServerAgentManifests(EmptyParams), #[serde(rename = "server.reload_agent_manifests")] diff --git a/src/api/schema/server.rs b/src/api/schema/server.rs index b274bc95..1adb62f5 100644 --- a/src/api/schema/server.rs +++ b/src/api/schema/server.rs @@ -14,6 +14,12 @@ pub struct ServerLiveHandoffParams { } #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)] +pub struct ServerSshAgentRegisterParams { + /// Absolute remote-host agent socket. Registration lasts until this API connection closes. + pub socket_path: String, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)] pub struct ServerCapabilities { pub live_handoff: bool, #[serde(default)] @@ -27,4 +33,7 @@ pub struct ServerCapabilities { /// Whether this server supports endpoint health probes. #[serde(default)] pub health_check: bool, + /// Supports connection-scoped `server.ssh_agent.register` on the local JSON API. + #[serde(default)] + pub ssh_agent_registration: bool, } diff --git a/src/api/schema/tests.rs b/src/api/schema/tests.rs index 4f15b322..3b43133b 100644 --- a/src/api/schema/tests.rs +++ b/src/api/schema/tests.rs @@ -726,6 +726,7 @@ fn success_response_round_trips() { endpoint_protocol_generation: Some(1), surface_interest: true, health_check: true, + ssh_agent_registration: false, }), }, }; diff --git a/src/api/server.rs b/src/api/server.rs index 23817fff..edb778f7 100644 --- a/src/api/server.rs +++ b/src/api/server.rs @@ -74,13 +74,14 @@ fn default_capabilities() -> Option<ServerCapabilities> { endpoint_protocol_generation: Some(crate::protocol::endpoint::ENDPOINT_PROTOCOL_GENERATION), surface_interest: true, health_check: true, + ssh_agent_registration: false, }) } fn start_server_inner( api_tx: ApiRequestSender, event_hub: EventHub, - capabilities: Option<ServerCapabilities>, + mut capabilities: Option<ServerCapabilities>, server_stop: Option<Arc<AtomicBool>>, ) -> std::io::Result<ServerHandle> { let path = socket_path(); @@ -91,6 +92,31 @@ fn start_server_inner( let identity = socket_file_identity(&path)?; info!(path = %path.display(), "api server listening"); + #[cfg(unix)] + let ssh_agents = match crate::platform::ssh_agent::SshAgentRegistry::new( + crate::platform::ssh_agent::socket_path(), + std::env::var_os("SSH_AUTH_SOCK").map(PathBuf::from), + ) { + Ok(registry) => Some(registry), + Err(error) => { + warn!(%error, "SSH agent refresh unavailable; retaining inherited pane environment"); + None + } + }; + + if let Some(capabilities) = capabilities.as_mut() { + capabilities.ssh_agent_registration = { + #[cfg(unix)] + { + ssh_agents.is_some() + } + #[cfg(not(unix))] + { + false + } + }; + } + let running = Arc::new(AtomicBool::new(true)); let listener_running = Arc::clone(&running); let thread = std::thread::spawn(move || { @@ -102,6 +128,8 @@ fn start_server_inner( let capabilities = capabilities.clone(); let server_stop = server_stop.clone(); let connection_running = Arc::clone(&listener_running); + #[cfg(unix)] + let ssh_agents = ssh_agents.clone(); std::thread::spawn(move || { if let Err(err) = handle_connection_with_stop( stream, @@ -110,6 +138,8 @@ fn start_server_inner( &connection_running, capabilities, server_stop.as_ref(), + #[cfg(unix)] + ssh_agents.as_ref(), ) { warn!(err = %err, "api connection failed"); } @@ -153,7 +183,16 @@ fn handle_connection( running: &Arc<AtomicBool>, capabilities: Option<ServerCapabilities>, ) -> std::io::Result<()> { - handle_connection_with_stop(stream, api_tx, event_hub, running, capabilities, None) + handle_connection_with_stop( + stream, + api_tx, + event_hub, + running, + capabilities, + None, + #[cfg(unix)] + None, + ) } fn handle_connection_with_stop( @@ -163,6 +202,7 @@ fn handle_connection_with_stop( running: &Arc<AtomicBool>, capabilities: Option<ServerCapabilities>, server_stop: Option<&Arc<AtomicBool>>, + #[cfg(unix)] ssh_agents: Option<&crate::platform::ssh_agent::SshAgentRegistry>, ) -> std::io::Result<()> { if let Err(err) = stream.set_send_timeout(Some(STREAM_WRITE_TIMEOUT)) { debug!(err = %err, "api connection write timeout unavailable"); @@ -213,6 +253,49 @@ fn handle_connection_with_stop( crate::logging::api_request_started(&request_id, method, changes_ui); match request.method { + #[cfg(unix)] + Method::ServerSshAgentRegister(params) => { + let lease = ssh_agents + .ok_or_else(|| io::Error::other("SSH agent registration is unavailable")) + .and_then(|registry| registry.register(PathBuf::from(params.socket_path))); + let lease = match lease { + Ok(lease) => lease, + Err(error) => { + return write_text_line_allow_disconnect( + &mut stream, + &error_response_json( + request_id, + if error.kind() == io::ErrorKind::InvalidInput { + "invalid_ssh_agent" + } else { + "ssh_agent_unavailable" + }, + error.to_string(), + ), + ) + } + }; + write_json_line( + &mut stream, + &SuccessResponse { + id: request_id, + result: ResponseResult::Ok {}, + }, + )?; + set_local_stream_polling(&mut stream, true)?; + let mut byte = [0]; + while running.load(Ordering::Relaxed) { + match poll_local_stream_read(&mut stream, &mut byte)? { + LocalStreamRead::Pending => { + // SSH can unlink an inherited socket after its bridge's lease closes. + lease.refresh()?; + std::thread::sleep(CONNECTION_POLL_INTERVAL); + } + _ => break, + } + } + Ok(()) + } Method::PaneGraphicsStream(params) => { let result = pane_graphics_stream::serve(stream, request_id.clone(), params, api_tx, running); @@ -404,6 +487,7 @@ pub(crate) fn api_method_name(method: &Method) -> &'static str { Method::ServerStop(_) => "server.stop", Method::ServerLiveHandoff(_) => "server.live_handoff", Method::ServerReloadConfig(_) => "server.reload_config", + Method::ServerSshAgentRegister(_) => "server.ssh_agent.register", Method::ServerAgentManifests(_) => "server.agent_manifests", Method::ServerReloadAgentManifests(_) => "server.reload_agent_manifests", Method::NotificationShow(_) => "notification.show", @@ -1025,6 +1109,53 @@ mod tests { (client, server, path) } + #[test] + fn ssh_agent_registration_lasts_only_for_the_api_connection() { + let directory = unique_test_path("agent-lease"); + fs::create_dir(&directory).unwrap(); + let agent = directory.join("upstream"); + let _agent = UnixListener::bind(&agent).unwrap(); + let stable = directory.join("stable"); + let registry = + crate::platform::ssh_agent::SshAgentRegistry::new(stable.clone(), None).unwrap(); + let (mut client, server, api_path) = local_stream_pair("agent-api"); + let (tx, _rx) = mpsc::unbounded_channel(); + let worker_registry = registry.clone(); + let worker = std::thread::spawn(move || { + handle_connection_with_stop( + server, + &tx, + &EventHub::default(), + &Arc::new(AtomicBool::new(true)), + None, + None, + Some(&worker_registry), + ) + .unwrap(); + }); + write_json_line( + &mut client, + &Request { + id: "agent-lease".into(), + method: Method::ServerSshAgentRegister( + crate::api::schema::ServerSshAgentRegisterParams { + socket_path: agent.to_string_lossy().into_owned(), + }, + ), + }, + ) + .unwrap(); + let response: SuccessResponse = serde_json::from_str(&read_line(&mut client)).unwrap(); + assert!(matches!(response.result, ResponseResult::Ok {})); + assert_eq!(fs::read_link(&stable).unwrap(), agent); + drop(client); + worker.join().unwrap(); + assert!(!stable.exists()); + drop(registry); + fs::remove_file(api_path).unwrap(); + fs::remove_dir_all(directory).unwrap(); + } + fn pane_info( pane_id: &str, agent_status: crate::api::schema::AgentStatus, @@ -1185,6 +1316,7 @@ mod tests { ), surface_interest: true, health_check: true, + ssh_agent_registration: false, }), None, None, diff --git a/src/app/api.rs b/src/app/api.rs index 0900a09d..1b5bad71 100644 --- a/src/app/api.rs +++ b/src/app/api.rs @@ -899,6 +899,13 @@ impl App { result: ResponseResult::Ok {}, } } + Method::ServerSshAgentRegister(_) => { + return responses::encode_error( + request.id, + "connection_local_only", + "SSH agent registration requires a persistent local JSON API connection", + ); + } Method::ServerLiveHandoff(_) => { let response = ErrorResponse { id: request.id, diff --git a/src/cli/status.rs b/src/cli/status.rs index d13be784..826c8ea6 100644 --- a/src/cli/status.rs +++ b/src/cli/status.rs @@ -280,6 +280,7 @@ struct ServerCapabilitiesJson { endpoint_protocol_generation: Option<u32>, surface_interest: bool, health_check: bool, + ssh_agent_registration: bool, } #[derive(Serialize)] @@ -325,6 +326,7 @@ fn server_status_json(server: &ServerRuntimeStatus) -> ServerStatusJson { endpoint_protocol_generation: capabilities.endpoint_protocol_generation, surface_interest: capabilities.surface_interest, health_check: capabilities.health_check, + ssh_agent_registration: capabilities.ssh_agent_registration, }), compatible: protocol.map(|value| value == crate::protocol::PROTOCOL_VERSION), endpoint_compatible: capabilities.as_ref().and_then(|capabilities| { @@ -422,11 +424,19 @@ mod tests { endpoint_protocol_generation: endpoint_generation, surface_interest: true, health_check: true, + ssh_agent_registration: false, }), } } #[test] + fn status_exposes_ssh_agent_registration() { + let server = running_server(Some("test"), None); + let value = serde_json::to_value(server_status_json(&server)).unwrap(); + assert_eq!(value["capabilities"]["ssh_agent_registration"], false); + } + + #[test] fn stale_compatible_server_does_not_require_restart() { let server = running_server( Some("0.0.0-old"), diff --git a/src/pane.rs b/src/pane.rs index 099580c3..485b226c 100644 --- a/src/pane.rs +++ b/src/pane.rs @@ -146,6 +146,8 @@ impl PaneLaunchEnv { } fn apply_pane_launch_env(cmd: &mut CommandBuilder, launch_env: &PaneLaunchEnv) { + #[cfg(unix)] + crate::platform::ssh_agent::apply_pane_env(cmd); cmd.env_remove("CODEX_THREAD_ID"); // OMP sets OMPCODE for shells it spawns. A pane launched from inside OMP // must not inherit it or its root agent would look like a nested session. diff --git a/src/platform/mod.rs b/src/platform/mod.rs index 488c5bee..864516da 100644 --- a/src/platform/mod.rs +++ b/src/platform/mod.rs @@ -3,6 +3,9 @@ //! Centralizes OS-dependent behavior behind a clean boundary so core //! modules don't scatter `#[cfg]` branches through product logic. +#[cfg(unix)] +pub(crate) mod ssh_agent; + pub(crate) struct HostShutdownMonitor { task: Option<tokio::task::JoinHandle<()>>, } 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(); + } +} diff --git a/src/remote.rs b/src/remote.rs index a79ac774..39c8bd58 100644 --- a/src/remote.rs +++ b/src/remote.rs @@ -4,6 +4,8 @@ mod host; mod process; mod restart_policy; mod saved; +#[cfg(unix)] +mod ssh_agent; pub(crate) use args::*; pub(crate) use attach::*; diff --git a/src/remote/host.rs b/src/remote/host.rs index 08f5d122..825ae10c 100644 --- a/src/remote/host.rs +++ b/src/remote/host.rs @@ -20,6 +20,8 @@ pub(crate) fn run_remote_client_bridge(args: &[String]) -> io::Result<()> { } }; ensure_remote_server_running()?; + #[cfg(unix)] + let _ssh_agent = super::ssh_agent::Registration::start(); let socket_path = crate::server::socket_paths::client_socket_path(); let stream = crate::ipc::connect_local_stream(&socket_path).map_err(|err| { diff --git a/src/remote/ssh_agent.rs b/src/remote/ssh_agent.rs new file mode 100644 index 00000000..639ce8e6 --- /dev/null +++ b/src/remote/ssh_agent.rs @@ -0,0 +1,239 @@ +use std::io::{self, Write}; +use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use crate::api::client::{parse_response_value, ApiClientError}; +use crate::api::schema::{Method, Request, ResponseResult, ServerSshAgentRegisterParams}; +use crate::ipc::{LocalStream, LocalStreamRead, LocalStreamReadCount}; + +pub(super) struct Registration { + stop: Arc<AtomicBool>, + thread: Option<std::thread::JoinHandle<()>>, +} + +impl Registration { + pub(super) fn start() -> Option<Self> { + let path = std::env::var("SSH_AUTH_SOCK") + .ok() + .filter(|path| !path.is_empty())?; + Self::start_at(path, crate::api::socket_path()) + } + + fn start_at(path: String, socket_path: PathBuf) -> Option<Self> { + let mut stream = match connect(&path, &socket_path) { + Ok(None) => return None, + Ok(stream) => stream, + Err(error) => { + tracing::debug!(%error, "SSH agent refresh unavailable; retrying while attached"); + None + } + }; + let stop = Arc::new(AtomicBool::new(false)); + let worker_stop = stop.clone(); + let thread = std::thread::spawn(move || { + let mut byte = [0]; + while !worker_stop.load(Ordering::Relaxed) { + if let Some(connection) = stream.as_mut() { + if !matches!( + crate::ipc::poll_local_stream_read(connection, &mut byte), + Ok(LocalStreamRead::Pending) + ) { + stream = None; + } + } else { + // Retry both initial API readiness and connections lost during handoff. + match connect(&path, &socket_path) { + Ok(None) => break, + result => stream = result.ok().flatten(), + } + } + std::thread::sleep(Duration::from_millis(100)); + } + }); + Some(Self { + stop, + thread: Some(thread), + }) + } +} + +impl Drop for Registration { + fn drop(&mut self) { + self.stop.store(true, Ordering::Relaxed); + if let Some(thread) = self.thread.take() { + let _ = thread.join(); + } + } +} + +fn connect(path: &str, socket_path: &Path) -> io::Result<Option<LocalStream>> { + let timeout = Duration::from_millis(500); + let status = crate::api::read_runtime_status_at(socket_path, timeout)?.ok_or_else(|| { + io::Error::new( + io::ErrorKind::NotConnected, + "SSH agent status API is not ready", + ) + })?; + if !status + .capabilities + .is_some_and(|capabilities| capabilities.ssh_agent_registration) + { + return Ok(None); + } + let mut stream = crate::ipc::connect_local_stream(socket_path)?; + let request = Request { + id: "remote:ssh-agent".into(), + method: Method::ServerSshAgentRegister(ServerSshAgentRegisterParams { + socket_path: path.into(), + }), + }; + serde_json::to_writer(&mut stream, &request)?; + stream.write_all(b"\n")?; + crate::ipc::set_local_stream_polling(&mut stream, true)?; + let deadline = Instant::now() + timeout; + let mut response = Vec::new(); + let mut byte = [0]; + while Instant::now() < deadline && response.len() < 4096 { + match crate::ipc::poll_local_stream_read_count(&mut stream, &mut byte)? { + LocalStreamReadCount::Data(_) if byte[0] == b'\n' => { + let response = match parse_response_value(serde_json::from_slice(&response)?) { + Ok(response) => response, + Err(ApiClientError::ErrorResponse(response)) + if response.error.code == "invalid_ssh_agent" => + { + tracing::debug!(error = %response.error.message, "SSH agent registration rejected"); + return Ok(None); + } + Err(error) => return Err(io::Error::other(error)), + }; + return match response.result { + ResponseResult::Ok {} => Ok(Some(stream)), + _ => Err(io::Error::other( + "unexpected SSH agent registration response", + )), + }; + } + LocalStreamReadCount::Data(_) => response.push(byte[0]), + LocalStreamReadCount::Closed => { + return Err(io::Error::other("SSH agent registration closed")) + } + LocalStreamReadCount::Pending => std::thread::sleep(Duration::from_millis(10)), + } + } + Err(io::Error::new( + io::ErrorKind::TimedOut, + "SSH agent registration did not complete", + )) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::io::{BufRead, BufReader, Read}; + use std::os::unix::net::UnixListener; + + #[test] + fn registration_stops_when_the_server_rejects_the_agent() { + let socket_path = + std::env::temp_dir().join(format!("herdr-agent-rejected-{}.sock", std::process::id())); + let listener = UnixListener::bind(&socket_path).unwrap(); + let server = std::thread::spawn(move || { + for expected in ["ping", "server.ssh_agent.register"] { + let (mut stream, _) = listener.accept().unwrap(); + stream + .set_read_timeout(Some(Duration::from_secs(5))) + .unwrap(); + let mut line = String::new(); + BufReader::new(&mut stream).read_line(&mut line).unwrap(); + let request: serde_json::Value = serde_json::from_str(&line).unwrap(); + assert_eq!(request["method"], expected); + let response = if expected == "ping" { + serde_json::json!({"id": request["id"], "result": { + "type": "pong", "version": "test", "protocol": crate::protocol::PROTOCOL_VERSION, + "capabilities": {"live_handoff": false, "ssh_agent_registration": true} + }}) + } else { + serde_json::json!({"id": request["id"], "error": { + "code": "invalid_ssh_agent", "message": "invalid agent path" + }}) + }; + writeln!(stream, "{response}").unwrap(); + } + }); + let registration = + Registration::start_at("relative-agent.sock".into(), socket_path.clone()); + server.join().unwrap(); + std::fs::remove_file(socket_path).unwrap(); + assert!( + registration.is_none(), + "a rejected agent must not start a retry worker" + ); + } + + #[test] + fn registration_retries_when_the_api_is_initially_missing() { + let socket_path = + std::env::temp_dir().join(format!("herdr-agent-retry-{}.sock", std::process::id())); + let registration = Registration::start_at("/test/agent.sock".into(), socket_path.clone()) + .expect("missing API must not permanently disable registration"); + let listener = UnixListener::bind(&socket_path).unwrap(); + listener.set_nonblocking(true).unwrap(); + for (attempt, expected) in [ + "ping", + "server.ssh_agent.register", + "ping", + "server.ssh_agent.register", + ] + .into_iter() + .enumerate() + { + let deadline = Instant::now() + Duration::from_secs(5); + let mut stream = loop { + match listener.accept() { + Ok((stream, _)) => { + stream.set_nonblocking(false).unwrap(); + break stream; + } + Err(error) if error.kind() == io::ErrorKind::WouldBlock => { + assert!(Instant::now() < deadline, "registration did not retry"); + std::thread::sleep(Duration::from_millis(10)); + } + Err(error) => panic!("{error}"), + } + }; + stream + .set_read_timeout(Some(Duration::from_secs(5))) + .unwrap(); + let mut line = String::new(); + BufReader::new(&mut stream).read_line(&mut line).unwrap(); + let request: serde_json::Value = serde_json::from_str(&line).unwrap(); + assert_eq!(request["method"], expected); + if attempt == 1 { + writeln!(stream, "{}", serde_json::json!({"id": request["id"], "error": { + "code": "ssh_agent_unavailable", "message": "replacement server owns the address" + }})).unwrap(); + continue; + } + let result = if expected == "ping" { + serde_json::json!({"type": "pong", "version": "test", "protocol": crate::protocol::PROTOCOL_VERSION, + "capabilities": {"live_handoff": false, "ssh_agent_registration": true}}) + } else { + serde_json::json!({"type": "ok"}) + }; + writeln!( + stream, + "{}", + serde_json::json!({"id": request["id"], "result": result}) + ) + .unwrap(); + if expected == "server.ssh_agent.register" { + drop(registration); + assert_eq!(stream.read(&mut [0]).unwrap(), 0); + break; + } + } + std::fs::remove_file(socket_path).unwrap(); + } +} diff --git a/src/update.rs b/src/update.rs index 4c128c57..958be6e1 100644 --- a/src/update.rs +++ b/src/update.rs @@ -2853,6 +2853,7 @@ mod tests { ), surface_interest: true, health_check: true, + ssh_agent_registration: false, }), }; let missing_baseline = crate::api::RuntimeStatus { @@ -2927,6 +2928,7 @@ mod tests { ), surface_interest: true, health_check: true, + ssh_agent_registration: false, }), }, }; @@ -3185,6 +3187,7 @@ mod tests { ), surface_interest: true, health_check: true, + ssh_agent_registration: false, }), }, }; |