1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91

#![forbid(unsafe_code)]

use super::{Connection, Error, Response, Result};

use std::io::ErrorKind;

enum EstablishedSessionState {
    Exited(Option<u32>),
    TtyAllocFail,
}

/// # Cancel safety
///
/// All methods of this struct is not cancellation safe.
#[derive(Debug)]
pub struct EstablishedSession {
    pub(super) conn: Connection,
    pub(super) session_id: u32,
}
impl EstablishedSession {
    fn check_session_id(&self, session_id: u32) -> Result<()> {
        if self.session_id != session_id {
            Err(Error::UnmatchedSessionId)
        } else {
            Ok(())
        }
    }

    /// Return None if TtyAllocFail, Some(...) if the process exited.
    async fn wait_impl(&mut self) -> Result<EstablishedSessionState> {
        use Response::*;

        let response = match self.conn.read_response().await {
            Result::Ok(response) => response,
            Err(err) => match &err {
                Error::IOError(io_err) if io_err.kind() == ErrorKind::UnexpectedEof => {
                    return Result::Ok(EstablishedSessionState::Exited(None))
                }
                _ => return Err(err),
            },
        };

        match response {
            TtyAllocFail { session_id } => {
                self.check_session_id(session_id)?;
                Result::Ok(EstablishedSessionState::TtyAllocFail)
            }
            ExitMessage {
                session_id,
                exit_value,
            } => {
                self.check_session_id(session_id)?;
                Result::Ok(EstablishedSessionState::Exited(Some(exit_value)))
            }
            response => Err(Error::InvalidServerResponse(
                "Expected Response TtyAllocFail or ExitMessage",
                response,
            )),
        }
    }

    /// Wait for session status to change
    ///
    /// Return `Self` on error so that you can handle the error and restart
    /// the operation.
    ///
    /// If the server close the connection without sending anything,
    /// this function would return `Ok(None)`.
    pub async fn wait(mut self) -> Result<SessionStatus, (Error, Self)> {
        use EstablishedSessionState::*;

        match self.wait_impl().await {
            Ok(Exited(exit_value)) => Ok(SessionStatus::Exited { exit_value }),
            Ok(TtyAllocFail) => Ok(SessionStatus::TtyAllocFail(self)),
            Err(err) => Err((err, self)),
        }
    }
}

#[derive(Debug)]
pub enum SessionStatus {
    /// Remote ssh server failed to allocate a tty, you can now return the tty
    /// to cooked mode.
    ///
    /// This arm includes `EstablishedSession` so that you can call `wait` on it
    /// again and retrieve the exit status and the underlying connection.
    TtyAllocFail(EstablishedSession),

    /// The process on the remote machine has exited with `exit_value`.
    Exited { exit_value: Option<u32> },
}