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