openssh_mux_client/
shutdown_mux_master.rs

1#![forbid(unsafe_code)]
2
3use crate::{constants, request::Request, Error, ErrorExt, Response, Result};
4
5use std::{io::Read, io::Write, os::unix::net::UnixStream, path::Path};
6
7use serde::{Deserialize, Serialize};
8use ssh_format::{from_bytes, Serializer};
9
10struct Connection {
11    raw_conn: UnixStream,
12    serializer: Serializer,
13}
14
15impl Connection {
16    fn write(&mut self, value: &Request) -> Result<()> {
17        let serializer = &mut self.serializer;
18
19        serializer.reset_counter();
20        // Reserve the header
21        serializer.output.resize(4, 0);
22
23        value.serialize(&mut *serializer)?;
24
25        let header = serializer.create_header(0)?;
26        // Write the header
27        serializer.output[..4].copy_from_slice(&header);
28
29        self.raw_conn.write_all(&serializer.output)?;
30
31        Ok(())
32    }
33
34    fn read_and_deserialize<'a, T>(&'a mut self, size: usize) -> Result<T>
35    where
36        T: Deserialize<'a>,
37    {
38        let buffer = &mut self.serializer.output;
39
40        buffer.resize(size, 0);
41
42        self.raw_conn.read_exact(buffer)?;
43
44        // Ignore any trailing bytes to be forward compatible
45        Ok(from_bytes(buffer)?.0)
46    }
47
48    /// Return size of the response.
49    fn read_header(&mut self) -> Result<u32> {
50        self.read_and_deserialize(4)
51    }
52
53    fn read_response(&mut self) -> Result<Response> {
54        let len = self.read_header()?;
55        self.read_and_deserialize(len as usize)
56    }
57
58    fn check_response_id(request_id: u32, response_id: u32) -> Result<()> {
59        if request_id != response_id {
60            Err(Error::UnmatchedRequestId)
61        } else {
62            Ok(())
63        }
64    }
65
66    fn exchange_hello(mut self) -> Result<Self> {
67        self.write(&Request::Hello {
68            version: constants::SSHMUX_VER,
69        })?;
70
71        let response = self.read_response()?;
72        if let Response::Hello { version } = response {
73            if version != constants::SSHMUX_VER {
74                Err(Error::UnsupportedMuxProtocol)
75            } else {
76                Ok(self)
77            }
78        } else {
79            Err(Error::invalid_server_response(&"Hello message", &response))
80        }
81    }
82
83    fn connect<P: AsRef<Path>>(path: P) -> Result<Self> {
84        Self::new(UnixStream::connect(path)?).exchange_hello()
85    }
86
87    fn new(raw_conn: UnixStream) -> Self {
88        Self {
89            raw_conn,
90            serializer: Serializer::new(Vec::with_capacity(20)),
91        }
92    }
93
94    /// Request the master to stop accepting new multiplexing requests
95    /// and remove its listener socket.
96    fn request_stop_listening(&mut self) -> Result<()> {
97        use Response::*;
98
99        let request_id = 0;
100        self.write(&Request::StopListening { request_id })?;
101
102        match self.read_response()? {
103            Ok { response_id } => {
104                Self::check_response_id(request_id, response_id)?;
105                Result::Ok(())
106            }
107            PermissionDenied {
108                response_id,
109                reason,
110            } => {
111                Self::check_response_id(request_id, response_id)?;
112                Err(Error::PermissionDenied(reason))
113            }
114            Failure {
115                response_id,
116                reason,
117            } => {
118                Self::check_response_id(request_id, response_id)?;
119                Err(Error::RequestFailure(reason))
120            }
121            response => Err(Error::invalid_server_response(
122                &"Ok, PermissionDenied or Failure",
123                &response,
124            )),
125        }
126    }
127}
128
129/// Request the master to stop accepting new multiplexing requests
130/// and remove its listener socket.
131///
132/// **Only suitable to use in `Drop::drop`.**
133pub fn shutdown_mux_master<P: AsRef<Path>>(path: P) -> Result<()> {
134    Connection::connect(path)?.request_stop_listening()
135}
136
137pub(crate) fn shutdown_mux_master_from(raw_conn: UnixStream) -> Result<()> {
138    Connection::new(raw_conn).request_stop_listening()
139}
140
141#[cfg(test)]
142mod tests {
143    use super::shutdown_mux_master;
144
145    #[test]
146    fn test_sync_request_stop_listening() {
147        shutdown_mux_master("/tmp/openssh-mux-client-test.socket").unwrap();
148    }
149}