openssh_mux_client/
session.rs

1#![forbid(unsafe_code)]
2
3use super::{Connection, Error, ErrorExt, Response, Result};
4
5use std::io::ErrorKind;
6
7enum EstablishedSessionState {
8    Exited(Option<u32>),
9    TtyAllocFail,
10}
11
12/// NOTE that once `EstablishedSession` is dropped, any data written to
13/// `stdin` will not be send to the remote process and
14/// `stdout` and `stderr` would eof immediately.
15///
16/// # Cancel safety
17///
18/// All methods of this struct is not cancellation safe.
19#[derive(Debug)]
20pub struct EstablishedSession {
21    pub(super) conn: Connection,
22    pub(super) session_id: u32,
23}
24impl EstablishedSession {
25    fn check_session_id(&self, session_id: u32) -> Result<()> {
26        if self.session_id != session_id {
27            Err(Error::UnmatchedSessionId)
28        } else {
29            Ok(())
30        }
31    }
32
33    /// Return None if TtyAllocFail, Some(...) if the process exited.
34    async fn wait_impl(&mut self) -> Result<EstablishedSessionState> {
35        use Response::*;
36
37        let response = match self.conn.read_response().await {
38            Result::Ok(response) => response,
39            Err(err) => match &err {
40                Error::IOError(io_err) if io_err.kind() == ErrorKind::UnexpectedEof => {
41                    return Result::Ok(EstablishedSessionState::Exited(None))
42                }
43                _ => return Err(err),
44            },
45        };
46
47        match response {
48            TtyAllocFail { session_id } => {
49                self.check_session_id(session_id)?;
50                Result::Ok(EstablishedSessionState::TtyAllocFail)
51            }
52            ExitMessage {
53                session_id,
54                exit_value,
55            } => {
56                self.check_session_id(session_id)?;
57                Result::Ok(EstablishedSessionState::Exited(Some(exit_value)))
58            }
59            response => Err(Error::invalid_server_response(
60                &"TtyAllocFail or ExitMessage",
61                &response,
62            )),
63        }
64    }
65
66    /// Wait for session status to change
67    ///
68    /// Return `Self` on error so that you can handle the error and restart
69    /// the operation.
70    ///
71    /// If the server close the connection without sending anything,
72    /// this function would return `Ok(None)`.
73    pub async fn wait(mut self) -> Result<SessionStatus, (Error, Self)> {
74        use EstablishedSessionState::*;
75
76        match self.wait_impl().await {
77            Ok(Exited(exit_value)) => Ok(SessionStatus::Exited { exit_value }),
78            Ok(TtyAllocFail) => Ok(SessionStatus::TtyAllocFail(self)),
79            Err(err) => Err((err, self)),
80        }
81    }
82}
83
84#[derive(Debug)]
85pub enum SessionStatus {
86    /// Remote ssh server failed to allocate a tty, you can now return the tty
87    /// to cooked mode.
88    ///
89    /// This arm includes `EstablishedSession` so that you can call `wait` on it
90    /// again and retrieve the exit status and the underlying connection.
91    TtyAllocFail(EstablishedSession),
92
93    /// The process on the remote machine has exited with `exit_value`.
94    Exited { exit_value: Option<u32> },
95}