Skip to main content

soma_fleet/
openssh_connector.rs

1use std::path::{Path, PathBuf};
2use std::sync::Arc;
3use std::time::Duration;
4
5use async_trait::async_trait;
6use openssh::{KnownHosts, Session, SessionBuilder};
7use tokio::sync::{RwLock, Semaphore};
8use tokio_util::sync::CancellationToken;
9
10use crate::{
11    ConnectionFactory, FleetError, FleetResult, HostEndpoint, HostId, HostRecord, TopologyRevision,
12    runtime::secure_runtime_subdir,
13};
14
15const DEFAULT_CONNECT_TIMEOUT: Duration = Duration::from_secs(5);
16const DEFAULT_SERVER_ALIVE_INTERVAL: Duration = Duration::from_secs(15);
17const DEFAULT_EXEC_PERMITS: usize = 4;
18
19/// Inspectable strict OpenSSH connection plan derived from one host record.
20#[derive(Debug, Clone, PartialEq, Eq)]
21pub struct OpenSshConnectPlan {
22    host: HostId,
23    revision: TopologyRevision,
24    destination: String,
25    port: u16,
26    user: Option<String>,
27    identity_file: Option<PathBuf>,
28    config_file: Option<PathBuf>,
29    known_hosts_file: Option<PathBuf>,
30    connect_timeout: Duration,
31    server_alive_interval: Duration,
32    strict_known_hosts: bool,
33}
34
35impl OpenSshConnectPlan {
36    /// Returns the target host identity.
37    #[must_use]
38    pub fn host(&self) -> &HostId {
39        &self.host
40    }
41    /// Returns the topology revision bound to this plan.
42    #[must_use]
43    pub fn revision(&self) -> &TopologyRevision {
44        &self.revision
45    }
46    /// Returns the SSH hostname or configuration alias.
47    #[must_use]
48    pub fn destination(&self) -> &str {
49        &self.destination
50    }
51    /// Returns the SSH port.
52    #[must_use]
53    pub const fn port(&self) -> u16 {
54        self.port
55    }
56    /// Returns the optional SSH user.
57    #[must_use]
58    pub fn user(&self) -> Option<&str> {
59        self.user.as_deref()
60    }
61    /// Returns the optional identity-file path.
62    #[must_use]
63    pub fn identity_file(&self) -> Option<&Path> {
64        self.identity_file.as_deref()
65    }
66    /// Returns the optional SSH config path.
67    #[must_use]
68    pub fn config_file(&self) -> Option<&Path> {
69        self.config_file.as_deref()
70    }
71    /// Returns the optional explicit known-hosts path.
72    #[must_use]
73    pub fn known_hosts_file(&self) -> Option<&Path> {
74        self.known_hosts_file.as_deref()
75    }
76    /// Returns the outer connection timeout.
77    #[must_use]
78    pub const fn connect_timeout(&self) -> Duration {
79        self.connect_timeout
80    }
81    /// Returns the server-alive interval.
82    #[must_use]
83    pub const fn server_alive_interval(&self) -> Duration {
84        self.server_alive_interval
85    }
86    /// Returns whether strict host-key verification is mandatory.
87    #[must_use]
88    pub const fn strict_known_hosts(&self) -> bool {
89        self.strict_known_hosts
90    }
91}
92
93/// Strict OpenSSH connection factory.
94#[derive(Debug, Clone)]
95pub struct OpenSshConnector {
96    connect_timeout: Duration,
97    server_alive_interval: Duration,
98    exec_permits: usize,
99}
100
101impl Default for OpenSshConnector {
102    fn default() -> Self {
103        Self {
104            connect_timeout: DEFAULT_CONNECT_TIMEOUT,
105            server_alive_interval: DEFAULT_SERVER_ALIVE_INTERVAL,
106            exec_permits: DEFAULT_EXEC_PERMITS,
107        }
108    }
109}
110
111impl OpenSshConnector {
112    /// Creates a connector with explicit nonzero bounds.
113    pub fn new(
114        connect_timeout: Duration,
115        server_alive_interval: Duration,
116        exec_permits: usize,
117    ) -> FleetResult<Self> {
118        if connect_timeout.is_zero() || server_alive_interval.is_zero() || exec_permits == 0 {
119            return Err(FleetError::Connection {
120                host: HostId::new("ssh").expect("static host id"),
121                message: "OpenSSH timeouts and execution permits must be nonzero".into(),
122            });
123        }
124        Ok(Self {
125            connect_timeout,
126            server_alive_interval,
127            exec_permits,
128        })
129    }
130
131    /// Derives the exact strict connection plan for one SSH host.
132    pub fn plan(&self, host: &HostRecord) -> FleetResult<OpenSshConnectPlan> {
133        let HostEndpoint::Ssh(endpoint) = host.endpoint() else {
134            return Err(FleetError::Connection {
135                host: host.id().clone(),
136                message: "OpenSSH connector requires an SSH endpoint".into(),
137            });
138        };
139        Ok(OpenSshConnectPlan {
140            host: host.id().clone(),
141            revision: host.revision().clone(),
142            destination: endpoint.host().to_owned(),
143            port: endpoint.port(),
144            user: endpoint.user().map(str::to_owned),
145            identity_file: endpoint.identity_file().map(Path::to_path_buf),
146            config_file: endpoint.config_file().map(Path::to_path_buf),
147            known_hosts_file: endpoint.known_hosts_file().map(Path::to_path_buf),
148            connect_timeout: self.connect_timeout,
149            server_alive_interval: self.server_alive_interval,
150            strict_known_hosts: true,
151        })
152    }
153
154    fn builder(&self, plan: &OpenSshConnectPlan) -> FleetResult<SessionBuilder> {
155        let mut builder = SessionBuilder::default();
156        builder
157            .known_hosts_check(KnownHosts::Strict)
158            .control_directory(secure_runtime_subdir("control")?)
159            .connect_timeout(plan.connect_timeout)
160            .server_alive_interval(plan.server_alive_interval)
161            .port(plan.port);
162        if let Some(user) = &plan.user {
163            builder.user(user.clone());
164        }
165        if let Some(path) = &plan.identity_file {
166            builder.keyfile(path);
167        }
168        if let Some(path) = &plan.config_file {
169            builder.config_file(path);
170        }
171        if let Some(path) = &plan.known_hosts_file {
172            builder.user_known_hosts_file(path);
173        }
174        Ok(builder)
175    }
176}
177
178/// Revision-bound multiplexed OpenSSH session.
179pub struct OpenSshConnection {
180    host: HostId,
181    revision: TopologyRevision,
182    session: RwLock<Option<Arc<Session>>>,
183    permits: Semaphore,
184}
185
186impl OpenSshConnection {
187    /// Returns the host identity.
188    #[must_use]
189    pub fn host(&self) -> &HostId {
190        &self.host
191    }
192    /// Returns the exact topology revision.
193    #[must_use]
194    pub fn revision(&self) -> &TopologyRevision {
195        &self.revision
196    }
197
198    pub(crate) async fn acquire_permit(
199        &self,
200    ) -> Result<tokio::sync::SemaphorePermit<'_>, tokio::sync::AcquireError> {
201        self.permits.acquire().await
202    }
203
204    pub(crate) async fn session(&self) -> FleetResult<Arc<Session>> {
205        self.session
206            .read()
207            .await
208            .as_ref()
209            .map(Arc::clone)
210            .ok_or_else(|| FleetError::Connection {
211                host: self.host.clone(),
212                message: "OpenSSH connection is closed".into(),
213            })
214    }
215
216    async fn close(&self) -> FleetResult<()> {
217        let session = self.session.write().await.take();
218        if let Some(session) = session
219            && let Ok(session) = Arc::try_unwrap(session)
220        {
221            session
222                .close()
223                .await
224                .map_err(|error| FleetError::Connection {
225                    host: self.host.clone(),
226                    message: format!("close failed: {error}"),
227                })?;
228        }
229        Ok(())
230    }
231}
232
233#[async_trait]
234impl ConnectionFactory for OpenSshConnector {
235    type Connection = OpenSshConnection;
236
237    async fn connect(
238        &self,
239        host: &HostRecord,
240        cancellation: &CancellationToken,
241    ) -> FleetResult<Self::Connection> {
242        if cancellation.is_cancelled() {
243            return Err(FleetError::Cancelled);
244        }
245        let plan = self.plan(host)?;
246        let builder = self.builder(&plan)?;
247        let connect = builder.connect_mux(&plan.destination);
248        let session = tokio::select! {
249            () = cancellation.cancelled() => return Err(FleetError::Cancelled),
250            result = tokio::time::timeout(plan.connect_timeout, connect) => match result {
251                Err(_) => return Err(FleetError::DeadlineExceeded),
252                Ok(Err(error)) => return Err(FleetError::Connection {
253                    host: host.id().clone(),
254                    message: format!("strict OpenSSH connect failed: {error}"),
255                }),
256                Ok(Ok(session)) => session,
257            }
258        };
259        Ok(OpenSshConnection {
260            host: host.id().clone(),
261            revision: host.revision().clone(),
262            session: RwLock::new(Some(Arc::new(session))),
263            permits: Semaphore::new(self.exec_permits),
264        })
265    }
266
267    async fn close(&self, connection: &Self::Connection) -> FleetResult<()> {
268        connection.close().await
269    }
270}
271
272#[cfg(test)]
273#[path = "openssh_connector_tests.rs"]
274mod tests;