soma_fleet/
forward.rs

1use std::os::unix::fs::{FileTypeExt, MetadataExt, PermissionsExt};
2use std::path::{Path, PathBuf};
3use std::sync::Arc;
4use std::sync::atomic::{AtomicU64, Ordering};
5use std::time::Duration;
6
7use openssh::{ForwardType, Session, Socket};
8use tokio_util::sync::CancellationToken;
9
10use crate::{
11    FleetError, FleetResult, HostRecord, OpenSshConnection, request::validate_absolute_path,
12    runtime::secure_runtime_subdir,
13};
14
15const SOCKET_WAIT: Duration = Duration::from_secs(2);
16const SOCKET_POLL: Duration = Duration::from_millis(20);
17static FORWARD_SEQUENCE: AtomicU64 = AtomicU64::new(1);
18
19/// Returns a private deterministic-shape local path for a forwarded Unix socket.
20pub fn forwarded_socket_path(host: &HostRecord) -> FleetResult<PathBuf> {
21    let sequence = FORWARD_SEQUENCE.fetch_add(1, Ordering::Relaxed);
22    let revision = host.revision().as_str();
23    let prefix = revision.get(..16).unwrap_or(revision);
24    Ok(secure_runtime_subdir("forward")?
25        .join(format!("{prefix}-{}-{sequence}.sock", std::process::id())))
26}
27
28/// RAII guard over one local-to-remote Unix socket forward.
29pub struct ForwardedUnixSocket {
30    session: Arc<Session>,
31    local_path: PathBuf,
32    remote_path: PathBuf,
33    closed: bool,
34}
35
36impl ForwardedUnixSocket {
37    /// Opens and secures a local Unix socket forwarding to a remote absolute path.
38    pub async fn open(
39        connection: &OpenSshConnection,
40        host: &HostRecord,
41        remote_path: impl Into<PathBuf>,
42        cancellation: &CancellationToken,
43    ) -> FleetResult<Self> {
44        if cancellation.is_cancelled() {
45            return Err(FleetError::Cancelled);
46        }
47        if connection.revision() != host.revision() {
48            return Err(FleetError::StaleTopology {
49                host: host.id().clone(),
50                expected: connection.revision().clone(),
51                actual: host.revision().clone(),
52            });
53        }
54        let remote_path = validate_absolute_path(remote_path.into())?;
55        let local_path = forwarded_socket_path(host)?;
56        if let Ok(metadata) = std::fs::symlink_metadata(&local_path) {
57            if metadata.file_type().is_symlink() {
58                return Err(FleetError::Connection {
59                    host: host.id().clone(),
60                    message: "forward socket path is a symbolic link".into(),
61                });
62            }
63            std::fs::remove_file(&local_path).map_err(|error| FleetError::Connection {
64                host: host.id().clone(),
65                message: format!("remove stale forward socket failed: {error}"),
66            })?;
67        }
68        let session = connection.session().await?;
69        let request = session.request_port_forward(
70            ForwardType::Local,
71            Socket::UnixSocket {
72                path: local_path.as_path().into(),
73            },
74            Socket::UnixSocket {
75                path: remote_path.as_path().into(),
76            },
77        );
78        tokio::select! {
79            () = cancellation.cancelled() => return Err(FleetError::Cancelled),
80            result = tokio::time::timeout(SOCKET_WAIT, request) => match result {
81                Err(_) => return Err(FleetError::DeadlineExceeded),
82                Ok(Err(error)) => return Err(FleetError::Connection {
83                    host: host.id().clone(),
84                    message: format!("open Unix socket forward failed: {error}"),
85                }),
86                Ok(Ok(())) => {}
87            }
88        }
89        if let Err(error) = secure_socket(&local_path, host).await {
90            let _ = session
91                .close_port_forward(
92                    ForwardType::Local,
93                    Socket::UnixSocket {
94                        path: local_path.as_path().into(),
95                    },
96                    Socket::UnixSocket {
97                        path: remote_path.as_path().into(),
98                    },
99                )
100                .await;
101            let _ = std::fs::remove_file(&local_path);
102            return Err(error);
103        }
104        Ok(Self {
105            session,
106            local_path,
107            remote_path,
108            closed: false,
109        })
110    }
111
112    /// Returns the private local socket path.
113    #[must_use]
114    pub fn path(&self) -> &Path {
115        &self.local_path
116    }
117
118    /// Explicitly closes the forward and removes the local socket.
119    pub async fn close(mut self) -> FleetResult<()> {
120        self.closed = true;
121        let result = self
122            .session
123            .close_port_forward(
124                ForwardType::Local,
125                Socket::UnixSocket {
126                    path: self.local_path.as_path().into(),
127                },
128                Socket::UnixSocket {
129                    path: self.remote_path.as_path().into(),
130                },
131            )
132            .await;
133        let _ = std::fs::remove_file(&self.local_path);
134        result.map_err(|error| FleetError::Connection {
135            host: crate::HostId::new("forward").expect("static host id"),
136            message: format!("close Unix socket forward failed: {error}"),
137        })
138    }
139}
140
141impl Drop for ForwardedUnixSocket {
142    fn drop(&mut self) {
143        if self.closed {
144            return;
145        }
146        let _ = std::fs::remove_file(&self.local_path);
147        if let Ok(handle) = tokio::runtime::Handle::try_current() {
148            let session = Arc::clone(&self.session);
149            let local_path = self.local_path.clone();
150            let remote_path = self.remote_path.clone();
151            handle.spawn(async move {
152                let _ = session
153                    .close_port_forward(
154                        ForwardType::Local,
155                        Socket::UnixSocket {
156                            path: local_path.as_path().into(),
157                        },
158                        Socket::UnixSocket {
159                            path: remote_path.as_path().into(),
160                        },
161                    )
162                    .await;
163            });
164        }
165    }
166}
167
168async fn secure_socket(path: &Path, host: &HostRecord) -> FleetResult<()> {
169    let deadline = tokio::time::Instant::now() + SOCKET_WAIT;
170    loop {
171        match std::fs::symlink_metadata(path) {
172            Ok(metadata) => {
173                if metadata.file_type().is_symlink() || !metadata.file_type().is_socket() {
174                    return Err(FleetError::Connection {
175                        host: host.id().clone(),
176                        message: "forward path is not a real Unix socket".into(),
177                    });
178                }
179                let uid = rustix::process::getuid().as_raw();
180                if metadata.uid() != uid {
181                    return Err(FleetError::Connection {
182                        host: host.id().clone(),
183                        message: "forward socket is not owned by the current user".into(),
184                    });
185                }
186                tokio::fs::set_permissions(path, std::fs::Permissions::from_mode(0o600))
187                    .await
188                    .map_err(|error| FleetError::Connection {
189                        host: host.id().clone(),
190                        message: format!("chmod 0600 forward socket failed: {error}"),
191                    })?;
192                return Ok(());
193            }
194            Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
195                if tokio::time::Instant::now() >= deadline {
196                    return Err(FleetError::Connection {
197                        host: host.id().clone(),
198                        message: "forward socket did not appear before timeout".into(),
199                    });
200                }
201                tokio::time::sleep(SOCKET_POLL).await;
202            }
203            Err(error) => {
204                return Err(FleetError::Connection {
205                    host: host.id().clone(),
206                    message: format!("inspect forward socket failed: {error}"),
207                });
208            }
209        }
210    }
211}
212
213#[cfg(test)]
214#[path = "forward_tests.rs"]
215mod tests;
soma_fleet/forward.rs

soma_fleet/
forward.rs
