1use super::{
6 take_utf8_capped, TunnelChannel, SshClientTrait, ConnectionConfig, ExecutionOutput,
7 TransferResult,
8 };
9 use crate::errors::{SshCliError, SshCliResult};
10 use async_trait::async_trait;
11 use std::path::Path;
12 use std::time::{Duration, Instant};
13 use zeroize::Zeroizing;
14
15 pub use crate::ssh::client_handler::ClientHandler;
17
18 pub struct SshClient {
20 pub session: russh::client::Handle<ClientHandler>,
22 cfg: ConnectionConfig,
23 forwarded: tokio::sync::Mutex<crate::ssh::client_handler::ForwardedSource>,
30 }
31
32 impl std::fmt::Debug for SshClient {
33 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
34 f.debug_struct("SshClient")
35 .field("host", &self.cfg.host)
36 .field("port", &self.cfg.port)
37 .field("user", &self.cfg.username)
38 .field("timeout_ms", &self.cfg.timeout_ms)
39 .finish()
40 }
41 }
42
43 impl SshClient {
44 #[must_use]
46 pub fn timeout_ms(&self) -> u64 {
47 self.cfg.timeout_ms.get()
48 }
49 }
50
51 fn map_exit_status(exit_status: u32) -> i32 {
52 i32::try_from(exit_status).unwrap_or(-1)
53 }
54
55 fn process_exec_message(
56 msg: russh::ChannelMsg,
57 stdout_bytes: &mut Vec<u8>,
58 stderr_bytes: &mut Vec<u8>,
59 exit_code: &mut Option<i32>,
60 byte_cap: usize,
61 truncated_stdout: &mut bool,
62 truncated_stderr: &mut bool,
63 ) -> bool {
64 use russh::ChannelMsg;
65
66 match msg {
67 ChannelMsg::Data { data } => {
68 super::append_capped(stdout_bytes, data.as_ref(), byte_cap, truncated_stdout);
70 }
71 ChannelMsg::ExtendedData { data, ext } => {
72 if ext == 1 {
74 super::append_capped(stderr_bytes, data.as_ref(), byte_cap, truncated_stderr);
75 } else {
76 tracing::debug!(ext, "extended data ignored");
77 }
78 }
79 ChannelMsg::ExitStatus { exit_status } => {
80 *exit_code = Some(map_exit_status(exit_status));
84 }
86 ChannelMsg::ExitSignal {
87 signal_name,
88 core_dumped,
89 error_message,
90 ..
91 } => {
92 tracing::warn!(
93 ?signal_name,
94 core_dumped,
95 %error_message,
96 "remote process terminated by signal"
97 );
98 }
100 ChannelMsg::Eof => {
101 tracing::debug!("EOF on SSH channel");
102 }
103 ChannelMsg::Close => {
104 tracing::debug!("SSH channel closed by server");
105 return true;
106 }
107 _ => {}
108 }
109
110 false
111 }
112
113 use crate::ssh::scp_wire::{
118 apply_local_mode, format_scp_t_line, format_scp_upload_header_with_mode,
119 interpret_scp_status, parse_scp_header, parse_scp_t_line, partial_download_path,
120 remote_scp_command, scp_mode_from_metadata, scp_read_data, scp_read_until_newline,
121 scp_wait_status, system_time_secs, SCP_OK,
122 };
123
124 impl SshClient {
125 pub async fn connect(cfg: ConnectionConfig) -> SshCliResult<Self> {
136 let auth = crate::ssh::client_connect::connect_authenticated(cfg).await?;
137 Ok(Self {
138 session: auth.session,
139 cfg: auth.cfg,
140 forwarded: tokio::sync::Mutex::new(auth.forwarded),
141 })
142 }
143
144 pub async fn run_command(
146 &mut self,
147 command: &str,
148 max_chars: usize,
149 stdin_data: Option<Vec<u8>>,
150 ) -> SshCliResult<ExecutionOutput> {
151 self.run_command_internal(command, max_chars, true, stdin_data)
152 .await
153 }
154
155 async fn run_command_internal(
156 &mut self,
157 command: &str,
158 max_chars: usize,
159 abort_on_timeout: bool,
160 stdin_data: Option<Vec<u8>>,
161 ) -> SshCliResult<ExecutionOutput> {
162 let start = Instant::now();
163 let timeout = Duration::from_millis(self.cfg.timeout_ms.get());
164
165 let job_marker = if abort_on_timeout {
171 crate::ssh::packing::new_remote_job_marker()
172 } else {
173 None
174 };
175 let marked_command = job_marker
176 .as_deref()
177 .map(|marker| crate::ssh::packing::wrap_with_abort_marker(command, marker));
178 let command: &str = marked_command.as_deref().unwrap_or(command);
179
180 let stdin_data: Option<Zeroizing<Vec<u8>>> = stdin_data.map(Zeroizing::new);
185 let result = tokio::time::timeout(timeout, async {
186 let mut channel = self
187 .session
188 .channel_open_session()
189 .await
190 .map_err(|e| SshCliError::channel_msg(format!("open session: {e}")))?;
191
192 channel
193 .exec(true, command)
194 .await
195 .map_err(|e| SshCliError::channel_msg(format!("exec: {e}")))?;
196
197 if let Some(ref bytes) = stdin_data {
199 channel
200 .data(bytes.as_slice())
201 .await
202 .map_err(|e| SshCliError::channel_msg(format!("stdin channel: {e}")))?;
203 channel
204 .eof()
205 .await
206 .map_err(|e| SshCliError::channel_msg(format!("eof channel: {e}")))?;
207 }
208 drop(stdin_data);
210
211 let byte_cap = super::exec_capture_byte_cap(max_chars);
213 let initial = byte_cap.min(8 * 1024);
214 let mut stdout_bytes: Vec<u8> = Vec::with_capacity(initial);
215 let mut stderr_bytes: Vec<u8> = Vec::with_capacity(initial);
216 let mut exit_code: Option<i32> = None;
217 let mut byte_trunc_stdout = false;
218 let mut byte_trunc_stderr = false;
219
220 while let Some(msg) = channel.wait().await {
221 if crate::signals::should_stop() {
223 return Err(SshCliError::Config(
224 "operation cancelled by signal".to_string(),
225 ));
226 }
227 if process_exec_message(
228 msg,
229 &mut stdout_bytes,
230 &mut stderr_bytes,
231 &mut exit_code,
232 byte_cap,
233 &mut byte_trunc_stdout,
234 &mut byte_trunc_stderr,
235 ) {
236 break;
237 }
238 }
239
240 Ok::<_, SshCliError>((
241 stdout_bytes,
242 stderr_bytes,
243 exit_code,
244 byte_trunc_stdout,
245 byte_trunc_stderr,
246 ))
247 })
248 .await;
249
250 let (stdout_bytes, stderr_bytes, exit_code, byte_trunc_stdout, byte_trunc_stderr) =
251 match result {
252 Ok(Ok(t)) => t,
253 Ok(Err(err)) => return Err(err),
254 Err(_) => {
255 if let Some(marker) = job_marker.as_deref() {
256 let abort_cmd = crate::ssh::packing::pack_abort_pkill(marker);
257 tracing::warn!(
258 marker = %marker,
259 "local timeout; attempting best-effort remote abort of this invocation only"
260 );
261 let _ = self.try_remote_abort(&abort_cmd).await;
262 }
263 return Err(SshCliError::SshTimeout(self.cfg.timeout_ms.get()));
264 }
265 };
266
267 let (stdout_truncado, trunc_stdout_chars) =
269 take_utf8_capped(stdout_bytes, max_chars);
270 let (stderr_truncado, trunc_stderr_chars) =
271 take_utf8_capped(stderr_bytes, max_chars);
272 let truncated_stdout = trunc_stdout_chars || byte_trunc_stdout;
273 let truncated_stderr = trunc_stderr_chars || byte_trunc_stderr;
274
275 let duration_ms = u64::try_from(start.elapsed().as_millis()).unwrap_or(u64::MAX);
276
277 Ok(ExecutionOutput {
278 stdout: stdout_truncado,
279 stderr: stderr_truncado,
280 exit_code,
281 truncated_stdout,
282 truncated_stderr,
283 duration_ms,
284 })
285 }
286 }