soma_fleet/
openssh_connector.rs1use 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#[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 #[must_use]
38 pub fn host(&self) -> &HostId {
39 &self.host
40 }
41 #[must_use]
43 pub fn revision(&self) -> &TopologyRevision {
44 &self.revision
45 }
46 #[must_use]
48 pub fn destination(&self) -> &str {
49 &self.destination
50 }
51 #[must_use]
53 pub const fn port(&self) -> u16 {
54 self.port
55 }
56 #[must_use]
58 pub fn user(&self) -> Option<&str> {
59 self.user.as_deref()
60 }
61 #[must_use]
63 pub fn identity_file(&self) -> Option<&Path> {
64 self.identity_file.as_deref()
65 }
66 #[must_use]
68 pub fn config_file(&self) -> Option<&Path> {
69 self.config_file.as_deref()
70 }
71 #[must_use]
73 pub fn known_hosts_file(&self) -> Option<&Path> {
74 self.known_hosts_file.as_deref()
75 }
76 #[must_use]
78 pub const fn connect_timeout(&self) -> Duration {
79 self.connect_timeout
80 }
81 #[must_use]
83 pub const fn server_alive_interval(&self) -> Duration {
84 self.server_alive_interval
85 }
86 #[must_use]
88 pub const fn strict_known_hosts(&self) -> bool {
89 self.strict_known_hosts
90 }
91}
92
93#[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 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 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
178pub struct OpenSshConnection {
180 host: HostId,
181 revision: TopologyRevision,
182 session: RwLock<Option<Arc<Session>>>,
183 permits: Semaphore,
184}
185
186impl OpenSshConnection {
187 #[must_use]
189 pub fn host(&self) -> &HostId {
190 &self.host
191 }
192 #[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;