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