ssh_mcp/background/
stream.rs

1use std::path::{Path, PathBuf};
2use std::sync::Arc;
3
4use russh::ChannelMsg;
5use russh::client;
6use tokio::io::AsyncWriteExt;
7use tracing::warn;
8
9use crate::background::{JobRegistry, LocalLogSpooler, Result};
10#[cfg(unix)]
11use crate::platform::O_NOFOLLOW_FLAG;
12
13const STATE_REASON_LIMIT_CHARS: usize = 160;
14
15enum JobTerminalUpdate {
16    Exit(i32),
17    StateLost(String),
18}
19
20fn truncate_state_reason(input: &str) -> String {
21    input.chars().take(STATE_REASON_LIMIT_CHARS).collect()
22}
23
24fn clamp_exit_status(exit_status: u32) -> i32 {
25    if exit_status > 255 {
26        255
27    } else {
28        exit_status as i32
29    }
30}
31
32async fn open_append_no_symlink(path: &Path) -> Result<tokio::fs::File> {
33    match tokio::fs::symlink_metadata(path).await {
34        Ok(meta) if meta.file_type().is_symlink() => {
35            return Err(std::io::Error::new(
36                std::io::ErrorKind::InvalidInput,
37                "log path is a symlink (refusing to follow it)",
38            )
39            .into());
40        }
41        Ok(_) => {}
42        Err(e) if e.kind() == std::io::ErrorKind::NotFound => {}
43        Err(e) => return Err(e.into()),
44    }
45
46    let mut opts = tokio::fs::OpenOptions::new();
47    opts.create(true).append(true);
48
49    #[cfg(unix)]
50    {
51        opts.custom_flags(O_NOFOLLOW_FLAG);
52    }
53
54    match opts.open(path).await {
55        Ok(f) => Ok(f),
56        Err(e) => {
57            if let Ok(meta) = tokio::fs::symlink_metadata(path).await
58                && meta.file_type().is_symlink()
59            {
60                return Err(std::io::Error::new(
61                    std::io::ErrorKind::InvalidInput,
62                    "log path is a symlink (refusing to follow it)",
63                )
64                .into());
65            }
66            Err(e.into())
67        }
68    }
69}
70
71pub struct OutputStreamer {
72    job_id: String,
73    local_log_path: PathBuf,
74    registry: Arc<JobRegistry>,
75    spooler: Arc<LocalLogSpooler>,
76}
77
78impl OutputStreamer {
79    pub fn new(
80        job_id: String,
81        local_log_path: PathBuf,
82        registry: Arc<JobRegistry>,
83        spooler: Arc<LocalLogSpooler>,
84    ) -> Self {
85        Self {
86            job_id,
87            local_log_path,
88            registry,
89            spooler,
90        }
91    }
92
93    pub async fn stream_channel(
94        self,
95        mut channel: russh::Channel<client::Msg>,
96        initial_stdout: Vec<u8>,
97    ) -> Result<Option<i32>> {
98        let res = self
99            .stream_channel_inner(&mut channel, initial_stdout)
100            .await;
101
102        match res {
103            Ok(Some(exit_code)) => {
104                self.update_job_terminal(JobTerminalUpdate::Exit(exit_code))
105                    .await;
106                Ok(Some(exit_code))
107            }
108            Ok(None) => {
109                self.update_job_terminal(JobTerminalUpdate::StateLost(
110                    "background_channel_closed_without_exit_status".to_string(),
111                ))
112                .await;
113                Ok(None)
114            }
115            Err(e) => {
116                self.update_job_terminal(JobTerminalUpdate::StateLost(format!(
117                    "background_stream_error: {}",
118                    truncate_state_reason(&e.to_string())
119                )))
120                .await;
121                Err(e)
122            }
123        }
124    }
125
126    async fn stream_channel_inner(
127        &self,
128        channel: &mut russh::Channel<client::Msg>,
129        initial_stdout: Vec<u8>,
130    ) -> Result<Option<i32>> {
131        let file = open_append_no_symlink(&self.local_log_path).await?;
132        let mut file = tokio::io::BufWriter::new(file);
133
134        if !initial_stdout.is_empty() {
135            file.write_all(&initial_stdout).await?;
136            file.flush().await?;
137        }
138
139        let mut exit_code: Option<i32> = None;
140        let mut saw_close_or_eof = false;
141
142        loop {
143            let next = if exit_code.is_some() && !saw_close_or_eof {
144                match tokio::time::timeout(std::time::Duration::from_millis(200), channel.wait())
145                    .await
146                {
147                    Ok(v) => v,
148                    Err(_) => break,
149                }
150            } else {
151                channel.wait().await
152            };
153
154            let Some(msg) = next else {
155                break;
156            };
157
158            match msg {
159                ChannelMsg::Data { data } => {
160                    file.write_all(data.as_ref()).await?;
161                    file.flush().await?;
162                }
163                ChannelMsg::ExtendedData { data, .. } => {
164                    // Phase 2 behavior used `2>&1` remote redirection; preserve a combined stream.
165                    file.write_all(data.as_ref()).await?;
166                    file.flush().await?;
167                }
168                ChannelMsg::ExitStatus { exit_status } => {
169                    exit_code = Some(clamp_exit_status(exit_status));
170                }
171                ChannelMsg::ExitSignal { signal_name, .. } => {
172                    // Map signal to a conventional shell exit code (128 + signal).
173                    let code = match signal_name {
174                        russh::Sig::HUP => 129,
175                        russh::Sig::INT => 130,
176                        russh::Sig::QUIT => 131,
177                        russh::Sig::ILL => 132,
178                        russh::Sig::ABRT => 134,
179                        russh::Sig::FPE => 136,
180                        russh::Sig::KILL => 137,
181                        russh::Sig::USR1 => 138,
182                        russh::Sig::SEGV => 139,
183                        russh::Sig::PIPE => 141,
184                        russh::Sig::ALRM => 142,
185                        russh::Sig::TERM => 143,
186                        russh::Sig::Custom(_) => 128,
187                    };
188                    exit_code = Some(code);
189                }
190                ChannelMsg::Close | ChannelMsg::Eof => {
191                    saw_close_or_eof = true;
192                    if exit_code.is_some() {
193                        break;
194                    }
195                }
196                _ => {}
197            }
198        }
199
200        file.flush().await?;
201        let inner = file.into_inner();
202        inner.sync_all().await?;
203
204        Ok(exit_code)
205    }
206
207    async fn update_job_terminal(&self, update: JobTerminalUpdate) {
208        let Some(job) = self.registry.get(&self.job_id).await else {
209            return;
210        };
211
212        let mut guard = job.lock().await;
213        match update {
214            JobTerminalUpdate::Exit(code) => guard.mark_exit(code),
215            JobTerminalUpdate::StateLost(reason) => guard.mark_state_lost(reason),
216        }
217        let persisted = guard.clone();
218        drop(guard);
219
220        if let Err(e) = self.spooler.persist_job_state(&persisted).await {
221            warn!(job_id = ?self.job_id, error = ?e, "failed to persist terminal job state");
222        }
223    }
224}
225
226#[cfg(test)]
227mod tests {
228    use super::*;
229
230    #[test]
231    fn test_clamp_exit_status_saturates() {
232        assert_eq!(clamp_exit_status(0), 0);
233        assert_eq!(clamp_exit_status(255), 255);
234        assert_eq!(clamp_exit_status(256), 255);
235        assert_eq!(clamp_exit_status(u32::MAX), 255);
236    }
237
238    #[tokio::test]
239    async fn test_open_append_no_symlink_writes() {
240        let tmp = tempfile::TempDir::new().expect("tempdir");
241        let path = tmp.path().join("out.log");
242
243        let mut file = open_append_no_symlink(&path)
244            .await
245            .expect("open_append_no_symlink");
246        file.write_all(b"hello\n").await.expect("write");
247        file.sync_all().await.expect("sync");
248
249        let content = tokio::fs::read_to_string(&path).await.expect("read");
250        assert!(content.contains("hello"));
251    }
252
253    #[cfg(unix)]
254    #[test]
255    fn test_open_append_no_symlink_rejects_symlink() {
256        use std::os::unix::fs::symlink;
257
258        let rt = tokio::runtime::Runtime::new().expect("runtime");
259        rt.block_on(async {
260            let tmp = tempfile::TempDir::new().expect("tempdir");
261            let target = tmp.path().join("target.log");
262            tokio::fs::write(&target, "x\n")
263                .await
264                .expect("write target");
265
266            let link = tmp.path().join("link.log");
267            symlink(&target, &link).expect("symlink");
268
269            let err = open_append_no_symlink(&link)
270                .await
271                .expect_err("symlink should be rejected");
272            let msg = err.to_string();
273            assert!(msg.contains("symlink"), "unexpected error: {msg}");
274        });
275    }
276}
ssh_mcp/background/stream.rs

ssh_mcp/background/
stream.rs
