soma_fleet/
openssh_driver.rs

1use std::sync::Arc;
2use std::time::Duration;
3
4use async_trait::async_trait;
5use tokio::io::AsyncWriteExt;
6use tokio_util::sync::CancellationToken;
7
8use crate::{
9    CommandExecutor, CommandOutput, CommandRequest, ConnectionPool, FleetError, FleetResult,
10    HostEndpoint, HostId, HostRecord, OpenSshConnection, OpenSshConnector, TopologySnapshot,
11    io::drain_bounded,
12};
13
14/// Pooled strict OpenSSH command driver.
15pub struct OpenSshDriver {
16    pool: ConnectionPool<OpenSshConnector>,
17}
18
19impl Default for OpenSshDriver {
20    fn default() -> Self {
21        Self::new(OpenSshConnector::default())
22    }
23}
24
25impl OpenSshDriver {
26    /// Creates a driver with an empty revision-aware connection pool.
27    #[must_use]
28    pub fn new(connector: OpenSshConnector) -> Self {
29        Self {
30            pool: ConnectionPool::new(Arc::new(connector)),
31        }
32    }
33
34    /// Invalidates every cached revision for one host.
35    pub async fn invalidate_host(&self, host: &HostId) -> FleetResult<usize> {
36        self.pool.invalidate_host(host).await
37    }
38
39    /// Evicts cached sessions absent from the supplied snapshot.
40    pub async fn retain_snapshot(&self, snapshot: &TopologySnapshot) -> FleetResult<usize> {
41        self.pool.retain_snapshot(snapshot).await
42    }
43
44    /// Closes all cached sessions.
45    pub async fn shutdown(&self) -> FleetResult<usize> {
46        self.pool.shutdown().await
47    }
48
49    /// Returns the exact pooled connection for forwarding adapters.
50    pub async fn connection(
51        &self,
52        host: &HostRecord,
53        cancellation: &CancellationToken,
54    ) -> FleetResult<Arc<OpenSshConnection>> {
55        self.pool.get_or_connect(host, cancellation).await
56    }
57}
58
59#[async_trait]
60impl CommandExecutor for OpenSshDriver {
61    async fn execute(
62        &self,
63        host: &HostRecord,
64        request: &CommandRequest,
65        cancellation: &CancellationToken,
66    ) -> FleetResult<CommandOutput> {
67        if !matches!(host.endpoint(), HostEndpoint::Ssh(_)) {
68            return Err(FleetError::Command {
69                host: host.id().clone(),
70                message: "OpenSSH driver requires an SSH endpoint".into(),
71            });
72        }
73        if request.working_dir().is_some() {
74            return Err(FleetError::Command {
75                host: host.id().clone(),
76                message:
77                    "remote working directories are unsupported without a typed remote launcher"
78                        .into(),
79            });
80        }
81        if cancellation.is_cancelled() {
82            return Err(FleetError::Cancelled);
83        }
84        request.validate_at(soma_ops::Timestamp::now())?;
85        let connection = self.connection(host, cancellation).await?;
86        if connection.revision() != host.revision() {
87            return Err(FleetError::StaleTopology {
88                host: host.id().clone(),
89                expected: connection.revision().clone(),
90                actual: host.revision().clone(),
91            });
92        }
93        let permit_timeout = remaining(request.deadline())?;
94        let permit = tokio::select! {
95            () = cancellation.cancelled() => return Err(FleetError::Cancelled),
96            result = tokio::time::timeout(permit_timeout, connection.acquire_permit()) => match result {
97                Err(_) => return Err(FleetError::DeadlineExceeded),
98                Ok(Err(_)) => return Err(FleetError::Connection {
99                    host: host.id().clone(),
100                    message: "OpenSSH execution semaphore is closed".into(),
101                }),
102                Ok(Ok(permit)) => permit,
103            }
104        };
105        let session = connection.session().await?;
106        let mut command = session.arc_command(request.program().to_owned());
107        command.args(request.args());
108        command
109            .stdin(if request.stdin().is_some() {
110                openssh::Stdio::piped()
111            } else {
112                openssh::Stdio::null()
113            })
114            .stdout(openssh::Stdio::piped())
115            .stderr(openssh::Stdio::piped());
116        let mut child = command.spawn().await.map_err(|error| FleetError::Command {
117            host: host.id().clone(),
118            message: format!("OpenSSH spawn failed: {error}"),
119        })?;
120        let input = match request.stdin() {
121            Some(bytes) => Some((
122                child.stdin().take().ok_or_else(|| FleetError::Command {
123                    host: host.id().clone(),
124                    message: "OpenSSH stdin pipe unavailable".into(),
125                })?,
126                bytes.to_vec(),
127            )),
128            None => None,
129        };
130        let stdout = child.stdout().take().ok_or_else(|| FleetError::Command {
131            host: host.id().clone(),
132            message: "OpenSSH stdout pipe unavailable".into(),
133        })?;
134        let stderr = child.stderr().take().ok_or_else(|| FleetError::Command {
135            host: host.id().clone(),
136            message: "OpenSSH stderr pipe unavailable".into(),
137        })?;
138        let timeout = remaining(request.deadline())?;
139        let completion = async move {
140            let streams = async {
141                tokio::try_join!(
142                    drain_bounded(stdout, request.max_stdout_bytes()),
143                    drain_bounded(stderr, request.max_stderr_bytes())
144                )
145                .map_err(openssh::Error::ChildIo)
146            };
147            let input = async move {
148                if let Some((mut stdin, bytes)) = input {
149                    stdin
150                        .write_all(&bytes)
151                        .await
152                        .map_err(openssh::Error::ChildIo)?;
153                    stdin.shutdown().await.map_err(openssh::Error::ChildIo)?;
154                }
155                Ok::<_, openssh::Error>(())
156            };
157            let (status, (stdout, stderr), ()) = tokio::try_join!(child.wait(), streams, input)?;
158            Ok::<_, openssh::Error>((status, stdout, stderr))
159        };
160        let mut completion = Box::pin(completion);
161        let result = tokio::select! {
162            () = cancellation.cancelled() => Err(FleetError::RemoteCommandDetached {
163                host: host.id().clone(),
164                reason: "cancellation",
165            }),
166            result = tokio::time::timeout(timeout, &mut completion) => match result {
167                Err(_) => Err(FleetError::RemoteCommandDetached {
168                    host: host.id().clone(),
169                    reason: "deadline",
170                }),
171                Ok(Err(error)) => Err(FleetError::Command {
172                    host: host.id().clone(),
173                    message: format!("OpenSSH command I/O failed: {error}"),
174                }),
175                Ok(Ok((status, (stdout, stdout_truncated), (stderr, stderr_truncated)))) => {
176                    Ok(CommandOutput::new(
177                        stdout,
178                        stderr,
179                        status.code(),
180                        stdout_truncated || stderr_truncated,
181                    ))
182                }
183            }
184        };
185        drop(completion);
186        drop(permit);
187        if result.is_err() {
188            let _ = self.invalidate_host(host.id()).await;
189        }
190        result
191    }
192}
193
194fn remaining(deadline: soma_ops::Timestamp) -> FleetResult<Duration> {
195    let millis = deadline
196        .unix_millis()
197        .saturating_sub(soma_ops::Timestamp::now().unix_millis());
198    if millis <= 0 {
199        Err(FleetError::DeadlineExceeded)
200    } else {
201        Ok(Duration::from_millis(millis as u64))
202    }
203}
204
205#[cfg(test)]
206#[path = "openssh_driver_tests.rs"]
207mod tests;
soma_fleet/openssh_driver.rs

soma_fleet/
openssh_driver.rs
