1use std::ffi::OsStr;
4#[cfg(all(test, unix))]
5use std::ffi::OsString;
6#[cfg(all(test, unix))]
7use std::fs;
8use std::io::{self, Read, Write};
9#[cfg(all(test, unix))]
10use std::os::unix::ffi::{OsStrExt, OsStringExt};
11use std::path::{Path, PathBuf};
12use std::time::Duration;
13
14use crate::ClientError;
15use rmux_ipc::{connect_blocking, BlockingLocalStream, LocalEndpoint};
16use rmux_proto::{
17 encode_frame, AttachSessionResponse, ControlMode, ControlModeResponse, FrameDecoder,
18 HandshakeRequest, Request, Response, RmuxError,
19};
20
21const READ_BUFFER_SIZE: usize = 8192;
23const SOCKET_CONNECT_TIMEOUT: Duration = Duration::from_secs(5);
25const SOCKET_WRITE_TIMEOUT: Duration = Duration::from_secs(5);
27const SOCKET_RESPONSE_TIMEOUT: Duration = Duration::from_secs(15);
29
30#[cfg(all(test, unix))]
31const FALLBACK_SOCKET_ROOT: &str = "/tmp";
32#[cfg(all(test, unix))]
33const SOCKET_DIR_PREFIX: &str = "rmux";
34
35pub fn default_socket_path() -> Result<PathBuf, ClientError> {
40 rmux_ipc::default_endpoint()
41 .map(LocalEndpoint::into_path)
42 .map_err(ClientError::Io)
43}
44
45pub fn socket_path_for_label(label: impl AsRef<OsStr>) -> Result<PathBuf, ClientError> {
47 rmux_ipc::endpoint_for_label(label)
48 .map(LocalEndpoint::into_path)
49 .map_err(ClientError::Io)
50}
51
52pub fn resolve_socket_path(
58 socket_name: Option<&OsStr>,
59 socket_path: Option<&Path>,
60) -> Result<PathBuf, ClientError> {
61 rmux_ipc::resolve_endpoint(socket_name, socket_path)
62 .map(LocalEndpoint::into_path)
63 .map_err(ClientError::Io)
64}
65
66#[allow(clippy::large_enum_variant)]
70#[derive(Debug)]
71pub enum ConnectResult {
72 Connected(Connection),
74 Absent,
76}
77
78pub fn connect_or_absent(socket_path: &Path) -> Result<ConnectResult, ClientError> {
86 connect_or_absent_with_timeout_using(
87 socket_path,
88 SOCKET_CONNECT_TIMEOUT,
89 connect_stream_with_timeout,
90 )
91}
92
93pub fn connect(socket_path: &Path) -> Result<Connection, ClientError> {
95 connect_with_timeout_using(
96 socket_path,
97 SOCKET_CONNECT_TIMEOUT,
98 connect_stream_with_timeout,
99 )
100}
101
102#[derive(Debug)]
104pub struct Connection {
105 stream: BlockingLocalStream,
106 decoder: FrameDecoder,
107 handshake_capabilities: Option<Vec<String>>,
108}
109
110#[allow(clippy::large_enum_variant)]
113#[derive(Debug)]
114pub enum AttachTransition {
115 Upgraded(AttachSessionUpgrade),
117 Rejected(Response),
119}
120
121#[allow(clippy::large_enum_variant)]
124#[derive(Debug)]
125pub enum ControlTransition {
126 Upgraded(ControlModeUpgrade),
128 Rejected(Response),
130}
131
132#[derive(Debug)]
134pub struct AttachSessionUpgrade {
135 response: AttachSessionResponse,
136 stream: BlockingLocalStream,
137 initial_bytes: Vec<u8>,
138}
139
140#[derive(Debug)]
142pub struct ControlModeUpgrade {
143 pub(crate) response: ControlModeResponse,
144 pub(crate) stream: BlockingLocalStream,
145}
146
147impl AttachSessionUpgrade {
148 #[must_use]
150 pub const fn response(&self) -> &AttachSessionResponse {
151 &self.response
152 }
153
154 #[must_use]
156 pub fn into_stream(self) -> BlockingLocalStream {
157 self.stream
158 }
159
160 #[must_use]
163 pub fn into_parts(self) -> (BlockingLocalStream, Vec<u8>) {
164 (self.stream, self.initial_bytes)
165 }
166}
167
168impl ControlModeUpgrade {
169 #[must_use]
171 pub const fn response(&self) -> &ControlModeResponse {
172 &self.response
173 }
174
175 #[must_use]
177 pub const fn mode(&self) -> ControlMode {
178 self.response.mode
179 }
180
181 #[must_use]
183 pub fn into_stream(self) -> BlockingLocalStream {
184 self.stream
185 }
186}
187
188impl Connection {
189 pub(crate) fn new(stream: BlockingLocalStream) -> Result<Self, ClientError> {
190 set_read_timeout(&stream, Some(SOCKET_RESPONSE_TIMEOUT)).map_err(ClientError::Io)?;
191 set_write_timeout(&stream, Some(SOCKET_WRITE_TIMEOUT)).map_err(ClientError::Io)?;
192
193 Ok(Self {
194 stream,
195 decoder: FrameDecoder::new(),
196 handshake_capabilities: None,
197 })
198 }
199
200 pub fn roundtrip(&mut self, request: &Request) -> Result<Response, ClientError> {
206 self.write_request(request)?;
207 self.read_response()
208 }
209
210 pub fn supports_capability(&mut self, capability: &str) -> Result<bool, ClientError> {
217 if let Some(capabilities) = &self.handshake_capabilities {
218 return Ok(capabilities.iter().any(|supported| supported == capability));
219 }
220
221 match self.roundtrip(&Request::Handshake(HandshakeRequest::current()))? {
222 Response::Handshake(response) => {
223 self.handshake_capabilities = Some(response.capabilities);
224 Ok(self
225 .handshake_capabilities
226 .as_ref()
227 .expect("handshake capabilities were just cached")
228 .iter()
229 .any(|supported| supported == capability))
230 }
231 Response::Error(error) => {
232 if matches!(&error.error, RmuxError::UnsupportedWireVersion { .. }) {
233 return Err(ClientError::Protocol(error.error));
234 }
235 self.handshake_capabilities = Some(Vec::new());
236 Ok(false)
237 }
238 _ => {
239 self.handshake_capabilities = Some(Vec::new());
240 Ok(false)
241 }
242 }
243 }
244
245 pub(crate) fn roundtrip_without_read_timeout(
250 &mut self,
251 request: &Request,
252 ) -> Result<Response, ClientError> {
253 let previous_timeout = read_timeout(&self.stream).map_err(ClientError::Io)?;
254 set_read_timeout(&self.stream, None).map_err(ClientError::Io)?;
255 let result = self.roundtrip(request);
256 set_read_timeout(&self.stream, previous_timeout).map_err(ClientError::Io)?;
257 result
258 }
259
260 pub(crate) fn write_request(&mut self, request: &Request) -> Result<(), ClientError> {
261 let frame = encode_frame(request).map_err(ClientError::Protocol)?;
262 self.stream.write_all(&frame).map_err(ClientError::Io)
263 }
264
265 pub(crate) fn read_response(&mut self) -> Result<Response, ClientError> {
266 let mut buffer = [0u8; READ_BUFFER_SIZE];
267
268 loop {
269 match self.decoder.next_frame::<Response>() {
270 Ok(Some(response)) => return Ok(response),
271 Ok(None) => {}
272 Err(error) => return Err(ClientError::Protocol(error)),
273 }
274
275 let bytes_read = match self.stream.read(&mut buffer) {
276 Ok(bytes_read) => bytes_read,
277 Err(error) if error.kind() == io::ErrorKind::Interrupted => continue,
278 Err(error) => return Err(ClientError::Io(error)),
279 };
280
281 if bytes_read == 0 {
282 return Err(ClientError::UnexpectedEof);
283 }
284
285 self.decoder.push_bytes(&buffer[..bytes_read]);
286 }
287 }
288
289 pub(crate) fn stream_mut(&mut self) -> &mut BlockingLocalStream {
290 &mut self.stream
291 }
292
293 pub(crate) fn into_attach_upgrade(
294 self,
295 response: AttachSessionResponse,
296 ) -> Result<AttachSessionUpgrade, ClientError> {
297 set_read_timeout(&self.stream, None).map_err(ClientError::Io)?;
298 set_write_timeout(&self.stream, None).map_err(ClientError::Io)?;
299 let initial_bytes = self.decoder.remaining_bytes().to_vec();
300
301 Ok(AttachSessionUpgrade {
302 response,
303 stream: self.stream,
304 initial_bytes,
305 })
306 }
307
308 pub(crate) fn into_control_upgrade(
309 self,
310 response: ControlModeResponse,
311 ) -> Result<ControlModeUpgrade, ClientError> {
312 set_read_timeout(&self.stream, None).map_err(ClientError::Io)?;
313 set_write_timeout(&self.stream, None).map_err(ClientError::Io)?;
314
315 Ok(ControlModeUpgrade {
316 response,
317 stream: self.stream,
318 })
319 }
320}
321
322pub(crate) fn read_response_frame_exact(
323 stream: &mut BlockingLocalStream,
324) -> Result<Response, ClientError> {
325 let mut decoder = FrameDecoder::new();
326 let mut byte = [0_u8; 1];
327
328 loop {
329 match decoder.next_frame::<Response>() {
330 Ok(Some(response)) => return Ok(response),
331 Ok(None) => {}
332 Err(error) => return Err(ClientError::Protocol(error)),
333 }
334
335 read_exact_or_eof(stream, &mut byte)?;
336 decoder.push_bytes(&byte);
337 }
338}
339
340fn read_exact_or_eof(
341 stream: &mut BlockingLocalStream,
342 buffer: &mut [u8],
343) -> Result<(), ClientError> {
344 match stream.read_exact(buffer) {
345 Ok(()) => Ok(()),
346 Err(error) if error.kind() == io::ErrorKind::UnexpectedEof => {
347 Err(ClientError::UnexpectedEof)
348 }
349 Err(error) => Err(ClientError::Io(error)),
350 }
351}
352
353#[cfg(all(test, unix))]
354fn socket_path_from_parts(
355 rmux_tmpdir: Option<&OsStr>,
356 user_id: u32,
357 label: &OsStr,
358) -> io::Result<PathBuf> {
359 let root = socket_root_from_parts(rmux_tmpdir)?;
360 let base = root.join(format!("{SOCKET_DIR_PREFIX}-{user_id}"));
361 let mut path = base.into_os_string().into_vec();
362 path.push(b'/');
363 path.extend_from_slice(label.as_bytes());
364
365 Ok(PathBuf::from(OsString::from_vec(path)))
366}
367
368#[cfg(all(test, unix))]
369fn socket_root_from_parts(rmux_tmpdir: Option<&OsStr>) -> io::Result<PathBuf> {
370 let rmux_tmpdir = rmux_tmpdir
371 .filter(|value| !value.is_empty())
372 .map(PathBuf::from);
373 let candidates = rmux_tmpdir
374 .into_iter()
375 .chain(std::iter::once(PathBuf::from(FALLBACK_SOCKET_ROOT)));
376
377 for candidate in candidates {
378 if let Ok(resolved) = fs::canonicalize(&candidate) {
379 return Ok(resolved);
380 }
381 }
382
383 Err(io::Error::new(
384 io::ErrorKind::NotFound,
385 "no suitable rmux socket directory",
386 ))
387}
388
389fn connect_or_absent_with_timeout_using<F>(
390 socket_path: &Path,
391 timeout: Duration,
392 connect_stream: F,
393) -> Result<ConnectResult, ClientError>
394where
395 F: FnOnce(&Path, Duration) -> io::Result<BlockingLocalStream>,
396{
397 match connect_stream(socket_path, timeout) {
398 Ok(stream) => Ok(ConnectResult::Connected(Connection::new(stream)?)),
399 Err(error) if is_absent_error(&error) => Ok(ConnectResult::Absent),
400 Err(error) => Err(ClientError::Io(error)),
401 }
402}
403
404fn connect_with_timeout_using<F>(
405 socket_path: &Path,
406 timeout: Duration,
407 connect_stream: F,
408) -> Result<Connection, ClientError>
409where
410 F: FnOnce(&Path, Duration) -> io::Result<BlockingLocalStream>,
411{
412 let stream = connect_stream(socket_path, timeout).map_err(ClientError::Io)?;
413 Connection::new(stream)
414}
415
416fn connect_stream_with_timeout(
417 socket_path: &Path,
418 timeout: Duration,
419) -> io::Result<BlockingLocalStream> {
420 connect_blocking(
421 &LocalEndpoint::from_path(socket_path.to_path_buf()),
422 timeout,
423 )
424}
425
426fn read_timeout(stream: &BlockingLocalStream) -> io::Result<Option<Duration>> {
427 stream.read_timeout()
428}
429
430fn set_read_timeout(stream: &BlockingLocalStream, timeout: Option<Duration>) -> io::Result<()> {
431 stream.set_read_timeout(timeout)
432}
433
434fn set_write_timeout(stream: &BlockingLocalStream, timeout: Option<Duration>) -> io::Result<()> {
435 stream.set_write_timeout(timeout)
436}
437
438fn is_absent_error(error: &io::Error) -> bool {
440 matches!(
441 error.kind(),
442 io::ErrorKind::NotFound | io::ErrorKind::ConnectionRefused
443 )
444}
445
446#[cfg(all(test, unix))]
447mod tests {
448 include!("connection/tests.rs");
449}