ssh_mcp/background/
spooler.rs

1use std::env;
2use std::path::{Component, Path, PathBuf};
3use std::time::{Duration, SystemTime};
4
5#[cfg(unix)]
6use std::os::unix::fs::PermissionsExt;
7
8use tokio::fs;
9use tokio::io::{AsyncReadExt, AsyncWriteExt};
10use tracing::warn;
11
12use super::job::{JobState, PersistedJobState};
13use super::{BackgroundError, Result};
14#[cfg(unix)]
15use crate::platform::O_NOFOLLOW_FLAG;
16
17/// Local-only log spooler.
18///
19/// Phase 1 is responsible for local directory existence and deterministic
20/// path generation. Later phases will stream remote stdout/stderr into these
21/// files and use the registry to serve `check-process`.
22#[derive(Debug, Clone)]
23pub struct LocalLogSpooler {
24    base_dir: PathBuf,
25}
26
27impl LocalLogSpooler {
28    pub fn new(base_dir: PathBuf) -> Self {
29        Self { base_dir }
30    }
31
32    pub fn new_default() -> Self {
33        Self::new(env::temp_dir().join("ssh-mcp"))
34    }
35
36    pub fn base_dir(&self) -> &Path {
37        &self.base_dir
38    }
39
40    pub async fn ensure_dir(&self) -> Result<()> {
41        match fs::symlink_metadata(&self.base_dir).await {
42            Ok(meta) => validate_spool_dir_meta(&meta)?,
43            Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
44                match fs::create_dir(&self.base_dir).await {
45                    Ok(()) => {}
46                    Err(e) if e.kind() == std::io::ErrorKind::AlreadyExists => {}
47                    Err(e) => return Err(e.into()),
48                }
49
50                // Re-validate after creation to close TOCTOU window.
51                let meta = fs::symlink_metadata(&self.base_dir).await?;
52                validate_spool_dir_meta(&meta)?;
53            }
54            Err(e) => return Err(e.into()),
55        }
56
57        #[cfg(unix)]
58        {
59            let meta = fs::symlink_metadata(&self.base_dir).await?;
60            validate_spool_dir_meta(&meta)?;
61            let perms = std::fs::Permissions::from_mode(0o700);
62            fs::set_permissions(&self.base_dir, perms).await?;
63        }
64
65        Ok(())
66    }
67
68    pub fn log_path_for(&self, job_id: &str) -> Result<PathBuf> {
69        validate_job_id(job_id)?;
70        Ok(self.base_dir.join(format!("{job_id}.log")))
71    }
72
73    pub fn state_path_for(&self, job_id: &str) -> Result<PathBuf> {
74        validate_job_id(job_id)?;
75        Ok(self.base_dir.join(format!("{job_id}.state")))
76    }
77
78    pub async fn persist_job_state(&self, job: &JobState) -> Result<()> {
79        self.ensure_dir().await?;
80
81        if job.log_path.parent() != Some(self.base_dir()) {
82            return Err(BackgroundError::InvalidState {
83                message: "job log path is outside spool directory",
84            });
85        }
86
87        let path = self.state_path_for(&job.job_id)?;
88        let payload =
89            serde_json::to_vec(&job.to_persisted()).map_err(|_| BackgroundError::InvalidState {
90                message: "failed to serialize persisted job state",
91            })?;
92
93        let mut file = open_spool_write_no_symlink(&path).await?;
94        file.write_all(&payload).await?;
95        file.sync_all().await?;
96        Ok(())
97    }
98
99    pub async fn load_job_state(&self, job_id: &str) -> Result<Option<JobState>> {
100        self.ensure_dir().await?;
101        let path = self.state_path_for(job_id)?;
102
103        let mut file = match open_spool_read_no_symlink(&path).await? {
104            Some(file) => file,
105            None => return Ok(None),
106        };
107
108        let mut payload = Vec::new();
109        file.read_to_end(&mut payload).await?;
110
111        let persisted: PersistedJobState =
112            serde_json::from_slice(&payload).map_err(|_| BackgroundError::InvalidState {
113                message: "failed to parse persisted job state",
114            })?;
115        let job = JobState::from_persisted(persisted)
116            .map_err(|message| BackgroundError::InvalidState { message })?;
117
118        if job.job_id != job_id {
119            return Err(BackgroundError::InvalidState {
120                message: "persisted job id does not match requested job id",
121            });
122        }
123        if job.log_path.parent() != Some(self.base_dir()) {
124            return Err(BackgroundError::InvalidState {
125                message: "persisted log path is outside spool directory",
126            });
127        }
128
129        Ok(Some(job))
130    }
131
132    pub async fn cleanup_old_logs(&self, max_age: Duration) -> Result<usize> {
133        self.ensure_dir().await?;
134
135        let now = SystemTime::now();
136        let mut removed = 0usize;
137
138        let mut entries = match fs::read_dir(&self.base_dir).await {
139            Ok(e) => e,
140            Err(e) => return Err(e.into()),
141        };
142
143        loop {
144            let entry = match entries.next_entry().await {
145                Ok(Some(e)) => e,
146                Ok(None) => break,
147                Err(e) => {
148                    warn!(error = ?e, "failed to read spool directory entry");
149                    continue;
150                }
151            };
152            let path = entry.path();
153
154            let file_name = match entry.file_name().to_str() {
155                Some(s) => s.to_owned(),
156                None => continue,
157            };
158
159            let Some((job_id, ext)) = split_spool_file_name(&file_name) else {
160                continue;
161            };
162            if validate_job_id(job_id).is_err() {
163                continue;
164            }
165            if ext != "log" && ext != "exit" && ext != "state" {
166                continue;
167            }
168
169            let meta = match fs::symlink_metadata(&path).await {
170                Ok(m) => m,
171                Err(e) => {
172                    warn!(path = ?path, error = ?e, "failed to stat spool file");
173                    continue;
174                }
175            };
176
177            let ft = meta.file_type();
178            if ft.is_symlink() || !ft.is_file() {
179                continue;
180            }
181
182            let modified = match meta.modified() {
183                Ok(m) => m,
184                Err(e) => {
185                    warn!(path = ?path, error = ?e, "failed to read mtime");
186                    continue;
187                }
188            };
189
190            let age = match now.duration_since(modified) {
191                Ok(d) => d,
192                Err(e) => {
193                    warn!(path = ?path, error = ?e, "invalid modified time");
194                    continue;
195                }
196            };
197            if age <= max_age {
198                continue;
199            }
200
201            match fs::remove_file(&path).await {
202                Ok(()) => removed += 1,
203                Err(e) => {
204                    warn!(path = ?path, error = ?e, "failed to remove old spool file");
205                }
206            }
207        }
208
209        Ok(removed)
210    }
211}
212
213async fn open_spool_write_no_symlink(path: &Path) -> Result<tokio::fs::File> {
214    match fs::symlink_metadata(path).await {
215        Ok(meta) if meta.file_type().is_symlink() => {
216            return Err(BackgroundError::InvalidState {
217                message: "spool metadata path is a symlink",
218            });
219        }
220        Ok(_) => {}
221        Err(e) if e.kind() == std::io::ErrorKind::NotFound => {}
222        Err(e) => return Err(e.into()),
223    }
224
225    let mut opts = tokio::fs::OpenOptions::new();
226    opts.write(true).create(true).truncate(true);
227
228    #[cfg(unix)]
229    {
230        opts.custom_flags(O_NOFOLLOW_FLAG);
231    }
232
233    match opts.open(path).await {
234        Ok(file) => Ok(file),
235        Err(e) => {
236            if let Ok(meta) = fs::symlink_metadata(path).await
237                && meta.file_type().is_symlink()
238            {
239                return Err(BackgroundError::InvalidState {
240                    message: "spool metadata path is a symlink",
241                });
242            }
243            Err(e.into())
244        }
245    }
246}
247
248async fn open_spool_read_no_symlink(path: &Path) -> Result<Option<tokio::fs::File>> {
249    match fs::symlink_metadata(path).await {
250        Ok(meta) => {
251            if meta.file_type().is_symlink() {
252                return Err(BackgroundError::InvalidState {
253                    message: "spool metadata path is a symlink",
254                });
255            }
256            if !meta.is_file() {
257                return Err(BackgroundError::InvalidState {
258                    message: "spool metadata path is not a regular file",
259                });
260            }
261        }
262        Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(None),
263        Err(e) => return Err(e.into()),
264    }
265
266    let mut opts = tokio::fs::OpenOptions::new();
267    opts.read(true);
268
269    #[cfg(unix)]
270    {
271        opts.custom_flags(O_NOFOLLOW_FLAG);
272    }
273
274    match opts.open(path).await {
275        Ok(file) => Ok(Some(file)),
276        Err(e) => {
277            if let Ok(meta) = fs::symlink_metadata(path).await
278                && meta.file_type().is_symlink()
279            {
280                return Err(BackgroundError::InvalidState {
281                    message: "spool metadata path is a symlink",
282                });
283            }
284            Err(e.into())
285        }
286    }
287}
288
289fn validate_spool_dir_meta(meta: &std::fs::Metadata) -> Result<()> {
290    let ft = meta.file_type();
291    if ft.is_symlink() {
292        return Err(BackgroundError::InvalidState {
293            message: "spool directory is a symlink",
294        });
295    }
296    if !ft.is_dir() {
297        return Err(BackgroundError::InvalidState {
298            message: "spool path exists but is not a directory",
299        });
300    }
301    Ok(())
302}
303
304fn validate_job_id(job_id: &str) -> Result<()> {
305    if job_id.is_empty() || job_id.len() > 128 {
306        return Err(BackgroundError::InvalidJobId {
307            job_id: job_id.to_owned(),
308        });
309    }
310    if job_id.as_bytes().contains(&0) {
311        return Err(BackgroundError::InvalidJobId {
312            job_id: job_id.to_owned(),
313        });
314    }
315
316    // Require a single, normal path component (reject absolute paths, separators, '.', '..').
317    let p = Path::new(job_id);
318    let mut components = p.components();
319    let Some(Component::Normal(_)) = components.next() else {
320        return Err(BackgroundError::InvalidJobId {
321            job_id: job_id.to_owned(),
322        });
323    };
324    if components.next().is_some() {
325        return Err(BackgroundError::InvalidJobId {
326            job_id: job_id.to_owned(),
327        });
328    }
329
330    // Tighten further: allow only ASCII alnum plus '-' and '_' to keep file naming predictable.
331    if !job_id
332        .bytes()
333        .all(|b| b.is_ascii_alphanumeric() || b == b'-' || b == b'_')
334    {
335        return Err(BackgroundError::InvalidJobId {
336            job_id: job_id.to_owned(),
337        });
338    }
339
340    Ok(())
341}
342
343fn split_spool_file_name(name: &str) -> Option<(&str, &str)> {
344    let (stem, ext) = name.rsplit_once('.')?;
345    if stem.is_empty() || ext.is_empty() {
346        return None;
347    }
348    Some((stem, ext))
349}
350
351#[cfg(test)]
352mod tests {
353    use super::*;
354    use std::time::{Instant, SystemTime};
355
356    async fn wait_until_older_than(path: &Path, min_age: Duration) {
357        let start = Instant::now();
358        loop {
359            let meta = tokio::fs::metadata(path).await.expect("metadata");
360            let modified = meta.modified().expect("modified time");
361            let age = SystemTime::now()
362                .duration_since(modified)
363                .unwrap_or_else(|_| Duration::from_secs(0));
364
365            if age >= min_age {
366                return;
367            }
368
369            assert!(
370                start.elapsed() < Duration::from_secs(2),
371                "file did not become old enough: {path:?}"
372            );
373            tokio::time::sleep(Duration::from_millis(5)).await;
374        }
375    }
376
377    #[tokio::test]
378    async fn test_ensure_dir_and_log_path_for() {
379        let tmp = tempfile::TempDir::new().expect("tempdir");
380        let base = tmp.path().join("spool");
381        let spooler = LocalLogSpooler::new(base.clone());
382
383        spooler.ensure_dir().await.expect("ensure_dir");
384        let meta = std::fs::metadata(&base).expect("spool dir metadata");
385        assert!(meta.is_dir());
386
387        let log = spooler.log_path_for("job_123").expect("log_path_for");
388        assert_eq!(log, base.join("job_123.log"));
389        let state = spooler.state_path_for("job_123").expect("state_path_for");
390        assert_eq!(state, base.join("job_123.state"));
391    }
392
393    #[test]
394    fn test_log_path_for_rejects_invalid_job_ids() {
395        let spooler = LocalLogSpooler::new(PathBuf::from("/tmp/ssh-mcp-test"));
396        for job_id in ["", "..", "/abs", "a/b", "a\\b", "job id", "job\n1"] {
397            assert!(spooler.log_path_for(job_id).is_err(), "job_id={job_id}");
398        }
399    }
400
401    #[tokio::test]
402    async fn test_cleanup_old_logs_removes_log_exit_and_state_files_only() {
403        let tmp = tempfile::TempDir::new().expect("tempdir");
404        let base = tmp.path().join("spool");
405        let spooler = LocalLogSpooler::new(base.clone());
406        spooler.ensure_dir().await.expect("ensure_dir");
407
408        let log = base.join("job_1.log");
409        let exit = base.join("job_1.exit");
410        let state = base.join("job_1.state");
411        let keep = base.join("job_1.tmp");
412        tokio::fs::write(&log, "hello\n").await.expect("write log");
413        tokio::fs::write(&exit, "0\n").await.expect("write exit");
414        tokio::fs::write(&state, "{}\n").await.expect("write state");
415        tokio::fs::write(&keep, "x\n").await.expect("write tmp");
416
417        // Avoid flakiness from tight timing windows by waiting until the newest file
418        // is safely older than the max_age threshold.
419        wait_until_older_than(&keep, Duration::from_millis(25)).await;
420        let removed = spooler
421            .cleanup_old_logs(Duration::from_millis(1))
422            .await
423            .expect("cleanup_old_logs");
424
425        assert!(removed >= 3, "expected to remove at least log+exit+state");
426        assert!(!log.exists(), "log should be removed");
427        assert!(!exit.exists(), "exit should be removed");
428        assert!(!state.exists(), "state should be removed");
429        assert!(keep.exists(), "non-log file should be kept");
430    }
431
432    #[tokio::test]
433    async fn test_persist_and_load_job_state_round_trip() {
434        let tmp = tempfile::TempDir::new().expect("tempdir");
435        let base = tmp.path().join("spool");
436        let spooler = LocalLogSpooler::new(base.clone());
437        spooler.ensure_dir().await.expect("ensure_dir");
438
439        let mut job = JobState::new_running(super::super::job::NewRunningJob {
440            job_id: "job_123".to_string(),
441            pid: 4242,
442            log_path: base.join("job_123.log"),
443            command: "wget https://example.test/file".to_string(),
444            connection_id: "test@localhost:22".to_string(),
445        });
446        job.mark_state_lost("stream_error");
447
448        spooler
449            .persist_job_state(&job)
450            .await
451            .expect("persist_job_state");
452
453        let loaded = spooler
454            .load_job_state("job_123")
455            .await
456            .expect("load_job_state")
457            .expect("job should exist");
458
459        assert_eq!(loaded.job_id, job.job_id);
460        assert_eq!(loaded.pid, job.pid);
461        assert_eq!(loaded.status, job.status);
462        assert_eq!(loaded.exit_code, job.exit_code);
463        assert_eq!(loaded.state_reason, job.state_reason);
464        assert_eq!(loaded.command, job.command);
465        assert_eq!(loaded.log_path, job.log_path);
466    }
467}
ssh_mcp/background/spooler.rs

ssh_mcp/background/
spooler.rs
