rustferry_ssh/
process_session.rs

1//! OpenSSH-backed execution for one framed snapshot session.
2
3use std::{
4    fs::File,
5    io::{self, Read, Write},
6    panic::{self, AssertUnwindSafe},
7    process::{Child, ChildStderr, ChildStdin, ChildStdout, Command, ExitStatus, Stdio},
8    sync::{
9        Arc,
10        atomic::{AtomicU8, Ordering},
11        mpsc::{self, Receiver, RecvTimeoutError, SyncSender, TrySendError},
12    },
13    thread,
14    time::{Duration, Instant},
15};
16
17use camino::Utf8Path;
18use rustferry_remote::{
19    CancellationToken, ProtocolPath, RemoteBuildEvent, SnapshotArtifactDescriptor,
20    WorkerDataPlaneFrameKind, write_worker_data_plane_frame,
21};
22use thiserror::Error;
23
24use crate::{
25    session::{
26        CreateOnlyArtifactSpool, SnapshotSessionClientError, SnapshotSessionOutcome,
27        SnapshotSessionProgress, SnapshotSessionRequest, run_snapshot_session_deferred,
28    },
29    transport::{ProcessSshRunner, SshInvocation, SshTransportError, terminate_child_bounded},
30};
31
32const SESSION_IO_POLL_INTERVAL: Duration = Duration::from_millis(10);
33const SESSION_CANCEL_WRITE_TIMEOUT: Duration = Duration::from_millis(100);
34const SESSION_EXIT_OBSERVATION_TIMEOUT: Duration = Duration::from_millis(100);
35const SESSION_STDERR_BUFFER_BYTES: usize = 16 * 1024;
36const SESSION_PIPE_CHUNK_BYTES: usize = 16 * 1024;
37
38/// Stable distinction between a validated session failure and process transport failure.
39#[derive(Clone, Debug, Eq, Error, PartialEq)]
40pub enum SshSnapshotSessionError {
41    /// Framing, peer identity, artifact verification, publication, or cleanup failed.
42    #[error(transparent)]
43    Session(#[from] SnapshotSessionClientError),
44    /// OpenSSH spawn, pipe, cancellation, timeout, or exit status failed.
45    #[error(transparent)]
46    Transport(#[from] SshTransportError),
47}
48
49impl ProcessSshRunner {
50    /// Run one full-duplex framed snapshot build through the fixed OpenSSH invocation.
51    ///
52    /// Successful protocol completion and a zero OpenSSH exit authorize the artifact
53    /// guard. The caller must still invoke [`CreateOnlyArtifactSpool::commit`] after
54    /// its own supporting-output checks.
55    ///
56    /// # Errors
57    ///
58    /// Distinguishes validated session failures from secret-free transport failures.
59    #[allow(clippy::too_many_arguments)]
60    pub fn run_snapshot_session<EventSink, Verifier, VerifyError>(
61        &self,
62        invocation: &SshInvocation,
63        request: SnapshotSessionRequest<'_>,
64        artifact_spool: &mut CreateOnlyArtifactSpool,
65        cancellation: &CancellationToken,
66        on_event: EventSink,
67        verify_artifact: Verifier,
68    ) -> Result<SnapshotSessionOutcome, SshSnapshotSessionError>
69    where
70        EventSink: FnMut(RemoteBuildEvent),
71        Verifier: FnOnce(
72            &mut File,
73            &Utf8Path,
74            &SnapshotArtifactDescriptor,
75        ) -> Result<ProtocolPath, VerifyError>,
76    {
77        let deadline = Instant::now() + invocation.timeout();
78        if cancellation.is_cancelled() {
79            return transport_error(artifact_spool, SshTransportError::Cancelled);
80        }
81        if invocation.revalidate_identity_file().is_err() {
82            return transport_error(artifact_spool, SshTransportError::IdentityFileChanged);
83        }
84        if invocation.revalidate_trust_snapshot().is_err() {
85            return transport_error(artifact_spool, SshTransportError::TrustSnapshotChanged);
86        }
87
88        let Ok(mut child) = Command::new(invocation.program())
89            .args(invocation.arguments())
90            .stdin(Stdio::piped())
91            .stdout(Stdio::piped())
92            .stderr(Stdio::piped())
93            .spawn()
94        else {
95            return transport_error(artifact_spool, SshTransportError::SpawnFailed);
96        };
97        let Some(stdin) = child.stdin.take() else {
98            terminate_child_bounded(child, deadline);
99            return transport_error(artifact_spool, SshTransportError::IoFailed);
100        };
101        let Some(stdout) = child.stdout.take() else {
102            terminate_child_bounded(child, deadline);
103            return transport_error(artifact_spool, SshTransportError::IoFailed);
104        };
105        let Some(stderr) = child.stderr.take() else {
106            terminate_child_bounded(child, deadline);
107            return transport_error(artifact_spool, SshTransportError::IoFailed);
108        };
109
110        let abort = AbortController::new(cancellation.clone(), deadline);
111        if spawn_stderr_drain(stderr).is_err() {
112            terminate_child_bounded(child, deadline);
113            return transport_error(artifact_spool, SshTransportError::IoFailed);
114        }
115        let Ok(stdout_events) = spawn_stdout_reader(stdout) else {
116            terminate_child_bounded(child, deadline);
117            return transport_error(artifact_spool, SshTransportError::IoFailed);
118        };
119        let Ok(mut writer) = InterruptibleWriter::spawn(stdin, abort.clone()) else {
120            terminate_child_bounded(child, deadline);
121            return transport_error(artifact_spool, SshTransportError::IoFailed);
122        };
123        let mut reader = InterruptibleReader::new(stdout_events, abort.clone());
124        let mut progress = SnapshotSessionProgress::default();
125
126        let session = panic::catch_unwind(AssertUnwindSafe(|| {
127            run_snapshot_session_deferred(
128                &mut reader,
129                &mut writer,
130                request,
131                artifact_spool,
132                on_event,
133                verify_artifact,
134                &mut progress,
135            )
136        }));
137        let session = match session {
138            Ok(session) => session,
139            Err(payload) => {
140                if progress.can_send_cancel() && writer.best_effort_cancel() {
141                    progress.mark_cancel_sent();
142                }
143                drop(writer);
144                terminate_child_bounded(child, deadline);
145                let _ = artifact_spool.cleanup_failure();
146                panic::resume_unwind(payload);
147            }
148        };
149        let outcome = match session {
150            Ok(outcome) => outcome,
151            Err(error) => {
152                if progress.can_send_cancel() && writer.best_effort_cancel() {
153                    progress.mark_cancel_sent();
154                }
155                drop(writer);
156                return finish_failed_session(child, deadline, &abort, error);
157            }
158        };
159
160        if writer.close().is_err() {
161            let error = classify_abort(&abort).unwrap_or(SshTransportError::IoFailed);
162            terminate_child_bounded(child, deadline);
163            return transport_error(artifact_spool, error);
164        }
165        wait_for_successful_exit(child, deadline, &abort)
166            .or_else(|error| transport_error(artifact_spool, error))?;
167        if let Err(error) = artifact_spool.authorize_commit() {
168            return session_error(artifact_spool, error);
169        }
170        Ok(outcome)
171    }
172}
173
174fn session_error<T>(
175    spool: &mut CreateOnlyArtifactSpool,
176    error: SnapshotSessionClientError,
177) -> Result<T, SshSnapshotSessionError> {
178    spool.cleanup_failure()?;
179    Err(SshSnapshotSessionError::Session(error))
180}
181
182fn transport_error<T>(
183    spool: &mut CreateOnlyArtifactSpool,
184    error: SshTransportError,
185) -> Result<T, SshSnapshotSessionError> {
186    spool.cleanup_failure()?;
187    Err(SshSnapshotSessionError::Transport(error))
188}
189
190fn finish_failed_session<T>(
191    mut child: Child,
192    deadline: Instant,
193    abort: &AbortController,
194    session_error: SnapshotSessionClientError,
195) -> Result<T, SshSnapshotSessionError> {
196    if !matches!(session_error, SnapshotSessionClientError::DataPlane(_)) {
197        terminate_child_bounded(child, deadline);
198        return Err(SshSnapshotSessionError::Session(session_error));
199    }
200    if let Some(error) = classify_cancel_or_timeout(abort) {
201        terminate_child_bounded(child, deadline);
202        return Err(SshSnapshotSessionError::Transport(error));
203    }
204    if let Some(status) = observe_exit_status(&mut child, deadline, &session_error) {
205        return if status.success() {
206            Err(SshSnapshotSessionError::Session(session_error))
207        } else {
208            Err(SshSnapshotSessionError::Transport(
209                SshTransportError::ProcessFailed {
210                    status: status.code(),
211                },
212            ))
213        };
214    }
215    let transport = classify_abort(abort);
216    terminate_child_bounded(child, deadline);
217    match transport {
218        Some(error) => Err(SshSnapshotSessionError::Transport(error)),
219        None => Err(SshSnapshotSessionError::Session(session_error)),
220    }
221}
222
223fn classify_cancel_or_timeout(abort: &AbortController) -> Option<SshTransportError> {
224    if abort.cancellation.is_cancelled() {
225        abort.mark(AbortReason::Cancelled);
226    } else if Instant::now() >= abort.deadline {
227        abort.mark(AbortReason::TimedOut);
228    }
229    match abort.reason() {
230        AbortReason::Cancelled => Some(SshTransportError::Cancelled),
231        AbortReason::TimedOut => Some(SshTransportError::TimedOut),
232        AbortReason::Active | AbortReason::IoFailed => None,
233    }
234}
235
236fn observe_exit_status(
237    child: &mut Child,
238    deadline: Instant,
239    error: &SnapshotSessionClientError,
240) -> Option<ExitStatus> {
241    let observe_until = if matches!(error, SnapshotSessionClientError::DataPlane(_)) {
242        deadline.min(Instant::now() + SESSION_EXIT_OBSERVATION_TIMEOUT)
243    } else {
244        Instant::now()
245    };
246    loop {
247        match child.try_wait() {
248            Ok(Some(status)) => return Some(status),
249            Ok(None) => {}
250            Err(_) => return None,
251        }
252        let remaining = observe_until.saturating_duration_since(Instant::now());
253        if remaining.is_zero() {
254            return None;
255        }
256        thread::sleep(SESSION_IO_POLL_INTERVAL.min(remaining));
257    }
258}
259
260fn wait_for_successful_exit(
261    mut child: Child,
262    deadline: Instant,
263    abort: &AbortController,
264) -> Result<(), SshTransportError> {
265    loop {
266        if let Some(error) = classify_abort(abort) {
267            terminate_child_bounded(child, deadline);
268            return Err(error);
269        }
270        match child.try_wait() {
271            Ok(Some(status)) if status.success() => return Ok(()),
272            Ok(Some(status)) => {
273                return Err(SshTransportError::ProcessFailed {
274                    status: status.code(),
275                });
276            }
277            Ok(None) => {}
278            Err(_) => {
279                terminate_child_bounded(child, deadline);
280                return Err(SshTransportError::IoFailed);
281            }
282        }
283        let remaining = deadline.saturating_duration_since(Instant::now());
284        if remaining.is_zero() {
285            abort.mark(AbortReason::TimedOut);
286            terminate_child_bounded(child, deadline);
287            return Err(SshTransportError::TimedOut);
288        }
289        thread::sleep(SESSION_IO_POLL_INTERVAL.min(remaining));
290    }
291}
292
293#[derive(Clone)]
294struct AbortController {
295    state: Arc<AtomicU8>,
296    cancellation: CancellationToken,
297    deadline: Instant,
298}
299
300#[derive(Clone, Copy, Debug, Eq, PartialEq)]
301#[repr(u8)]
302enum AbortReason {
303    Active = 0,
304    Cancelled = 1,
305    TimedOut = 2,
306    IoFailed = 3,
307}
308
309impl AbortController {
310    fn new(cancellation: CancellationToken, deadline: Instant) -> Self {
311        Self {
312            state: Arc::new(AtomicU8::new(AbortReason::Active as u8)),
313            cancellation,
314            deadline,
315        }
316    }
317
318    fn checkpoint(&self) -> io::Result<()> {
319        if self.cancellation.is_cancelled() {
320            self.mark(AbortReason::Cancelled);
321        } else if Instant::now() >= self.deadline {
322            self.mark(AbortReason::TimedOut);
323        }
324        match self.reason() {
325            AbortReason::Active => Ok(()),
326            AbortReason::Cancelled => Err(io::Error::new(
327                io::ErrorKind::ConnectionAborted,
328                "snapshot session cancelled",
329            )),
330            AbortReason::TimedOut => Err(io::Error::new(
331                io::ErrorKind::TimedOut,
332                "snapshot session timed out",
333            )),
334            AbortReason::IoFailed => Err(io::Error::new(
335                io::ErrorKind::BrokenPipe,
336                "snapshot session pipe failed",
337            )),
338        }
339    }
340
341    fn wait_interval(&self) -> Duration {
342        SESSION_IO_POLL_INTERVAL.min(self.deadline.saturating_duration_since(Instant::now()))
343    }
344
345    fn mark(&self, reason: AbortReason) {
346        let _ = self.state.compare_exchange(
347            AbortReason::Active as u8,
348            reason as u8,
349            Ordering::AcqRel,
350            Ordering::Acquire,
351        );
352    }
353
354    fn reason(&self) -> AbortReason {
355        match self.state.load(Ordering::Acquire) {
356            1 => AbortReason::Cancelled,
357            2 => AbortReason::TimedOut,
358            3 => AbortReason::IoFailed,
359            _ => AbortReason::Active,
360        }
361    }
362}
363
364fn classify_abort(abort: &AbortController) -> Option<SshTransportError> {
365    if abort.cancellation.is_cancelled() {
366        abort.mark(AbortReason::Cancelled);
367    } else if Instant::now() >= abort.deadline {
368        abort.mark(AbortReason::TimedOut);
369    }
370    match abort.reason() {
371        AbortReason::Active => None,
372        AbortReason::Cancelled => Some(SshTransportError::Cancelled),
373        AbortReason::TimedOut => Some(SshTransportError::TimedOut),
374        AbortReason::IoFailed => Some(SshTransportError::IoFailed),
375    }
376}
377
378enum StdoutEvent {
379    Bytes(Vec<u8>),
380    Eof,
381    Failed,
382}
383
384fn spawn_stdout_reader(mut stdout: ChildStdout) -> Result<Receiver<StdoutEvent>, ()> {
385    let (sender, receiver) = mpsc::sync_channel(2);
386    thread::Builder::new()
387        .name("rustferry-ssh-session-stdout".to_owned())
388        .spawn(move || {
389            let mut buffer = [0_u8; SESSION_PIPE_CHUNK_BYTES];
390            loop {
391                match stdout.read(&mut buffer) {
392                    Ok(0) => {
393                        let _ = sender.send(StdoutEvent::Eof);
394                        return;
395                    }
396                    Ok(count) => {
397                        if sender
398                            .send(StdoutEvent::Bytes(buffer[..count].to_vec()))
399                            .is_err()
400                        {
401                            return;
402                        }
403                    }
404                    Err(_) => {
405                        let _ = sender.send(StdoutEvent::Failed);
406                        return;
407                    }
408                }
409            }
410        })
411        .map(drop)
412        .map_err(|_| ())?;
413    Ok(receiver)
414}
415
416fn spawn_stderr_drain(mut stderr: ChildStderr) -> io::Result<()> {
417    thread::Builder::new()
418        .name("rustferry-ssh-session-stderr".to_owned())
419        .spawn(move || {
420            let mut buffer = [0_u8; SESSION_STDERR_BUFFER_BYTES];
421            while let Ok(count) = stderr.read(&mut buffer) {
422                if count == 0 {
423                    return;
424                }
425            }
426        })
427        .map(drop)
428}
429
430struct InterruptibleReader {
431    events: Receiver<StdoutEvent>,
432    pending: Vec<u8>,
433    offset: usize,
434    eof: bool,
435    abort: AbortController,
436}
437
438impl InterruptibleReader {
439    fn new(events: Receiver<StdoutEvent>, abort: AbortController) -> Self {
440        Self {
441            events,
442            pending: Vec::new(),
443            offset: 0,
444            eof: false,
445            abort,
446        }
447    }
448}
449
450impl Read for InterruptibleReader {
451    fn read(&mut self, buffer: &mut [u8]) -> io::Result<usize> {
452        if buffer.is_empty() {
453            return Ok(0);
454        }
455        loop {
456            if self.offset < self.pending.len() {
457                let count = buffer.len().min(self.pending.len() - self.offset);
458                buffer[..count].copy_from_slice(&self.pending[self.offset..self.offset + count]);
459                self.offset += count;
460                if self.offset == self.pending.len() {
461                    self.pending.clear();
462                    self.offset = 0;
463                }
464                return Ok(count);
465            }
466            if self.eof {
467                return Ok(0);
468            }
469            self.abort.checkpoint()?;
470            match self.events.recv_timeout(self.abort.wait_interval()) {
471                Ok(StdoutEvent::Bytes(bytes)) => self.pending = bytes,
472                Ok(StdoutEvent::Eof) => self.eof = true,
473                Ok(StdoutEvent::Failed) | Err(RecvTimeoutError::Disconnected) => {
474                    self.abort.mark(AbortReason::IoFailed);
475                    return Err(io::Error::new(
476                        io::ErrorKind::BrokenPipe,
477                        "snapshot session stdout failed",
478                    ));
479                }
480                Err(RecvTimeoutError::Timeout) => {}
481            }
482        }
483    }
484}
485
486enum StdinCommand {
487    Write {
488        bytes: Vec<u8>,
489        flush: bool,
490        result: SyncSender<io::Result<()>>,
491    },
492    Close {
493        result: SyncSender<io::Result<()>>,
494    },
495}
496
497struct InterruptibleWriter {
498    commands: SyncSender<StdinCommand>,
499    abort: AbortController,
500}
501
502impl InterruptibleWriter {
503    fn spawn(mut stdin: ChildStdin, abort: AbortController) -> Result<Self, ()> {
504        let (commands, receiver) = mpsc::sync_channel(1);
505        thread::Builder::new()
506            .name("rustferry-ssh-session-stdin".to_owned())
507            .spawn(move || {
508                while let Ok(command) = receiver.recv() {
509                    match command {
510                        StdinCommand::Write {
511                            bytes,
512                            flush,
513                            result,
514                        } => {
515                            let written = stdin
516                                .write_all(&bytes)
517                                .and_then(|()| if flush { stdin.flush() } else { Ok(()) });
518                            let failed = written.is_err();
519                            let _ = result.send(written);
520                            if failed {
521                                return;
522                            }
523                        }
524                        StdinCommand::Close { result } => {
525                            let flushed = stdin.flush();
526                            drop(stdin);
527                            let _ = result.send(flushed);
528                            return;
529                        }
530                    }
531                }
532            })
533            .map(drop)
534            .map_err(|_| ())?;
535        Ok(Self { commands, abort })
536    }
537
538    fn request(&self, bytes: Vec<u8>, flush: bool) -> io::Result<()> {
539        self.abort.checkpoint()?;
540        let (sender, receiver) = mpsc::sync_channel(1);
541        self.send_command(StdinCommand::Write {
542            bytes,
543            flush,
544            result: sender,
545        })?;
546        self.wait_result(&receiver)
547    }
548
549    fn send_command(&self, mut command: StdinCommand) -> io::Result<()> {
550        loop {
551            self.abort.checkpoint()?;
552            match self.commands.try_send(command) {
553                Ok(()) => return Ok(()),
554                Err(TrySendError::Full(returned)) => command = returned,
555                Err(TrySendError::Disconnected(_)) => {
556                    self.abort.mark(AbortReason::IoFailed);
557                    return self.abort.checkpoint();
558                }
559            }
560            thread::sleep(self.abort.wait_interval());
561        }
562    }
563
564    fn wait_result(&self, receiver: &Receiver<io::Result<()>>) -> io::Result<()> {
565        loop {
566            self.abort.checkpoint()?;
567            match receiver.recv_timeout(self.abort.wait_interval()) {
568                Ok(Ok(())) => return Ok(()),
569                Ok(Err(_)) | Err(RecvTimeoutError::Disconnected) => {
570                    self.abort.mark(AbortReason::IoFailed);
571                    return self.abort.checkpoint();
572                }
573                Err(RecvTimeoutError::Timeout) => {}
574            }
575        }
576    }
577
578    fn best_effort_cancel(&self) -> bool {
579        let mut bytes = Vec::new();
580        if write_worker_data_plane_frame(&mut bytes, WorkerDataPlaneFrameKind::Cancel, 3, &[])
581            .is_err()
582        {
583            return false;
584        }
585        let deadline = Instant::now() + SESSION_CANCEL_WRITE_TIMEOUT;
586        let (sender, receiver) = mpsc::sync_channel(1);
587        let mut command = StdinCommand::Write {
588            bytes,
589            flush: true,
590            result: sender,
591        };
592        loop {
593            match self.commands.try_send(command) {
594                Ok(()) => break,
595                Err(TrySendError::Full(returned)) => command = returned,
596                Err(TrySendError::Disconnected(_)) => return false,
597            }
598            let remaining = deadline.saturating_duration_since(Instant::now());
599            if remaining.is_zero() {
600                return false;
601            }
602            thread::sleep(SESSION_IO_POLL_INTERVAL.min(remaining));
603        }
604        loop {
605            let remaining = deadline.saturating_duration_since(Instant::now());
606            if remaining.is_zero() {
607                return false;
608            }
609            match receiver.recv_timeout(SESSION_IO_POLL_INTERVAL.min(remaining)) {
610                Ok(Ok(())) => return true,
611                Ok(Err(_)) | Err(RecvTimeoutError::Disconnected) => return false,
612                Err(RecvTimeoutError::Timeout) => {}
613            }
614        }
615    }
616
617    fn close(self) -> io::Result<()> {
618        let (sender, receiver) = mpsc::sync_channel(1);
619        self.send_command(StdinCommand::Close { result: sender })?;
620        self.wait_result(&receiver)
621    }
622}
623
624impl Write for InterruptibleWriter {
625    fn write(&mut self, buffer: &[u8]) -> io::Result<usize> {
626        self.request(buffer.to_vec(), false)?;
627        Ok(buffer.len())
628    }
629
630    fn flush(&mut self) -> io::Result<()> {
631        self.request(Vec::new(), true)
632    }
633}
rustferry_ssh/process_session.rs

rustferry_ssh/
process_session.rs
