rmux_client/
control.rs

1//! Blocking tmux-compatible control-mode client transport.
2
3use std::io::{self, Read, Write};
4use std::sync::mpsc;
5use std::thread;
6
7use rmux_ipc::BlockingLocalStream;
8#[cfg(any(test, windows))]
9use rmux_proto::CONTROL_STDIN_EOF_MARKER;
10use rmux_proto::{
11    ClientTerminalContext, ControlMode, ControlModeRequest, Request, Response, CONTROL_CONTROL_END,
12    CONTROL_CONTROL_START, MAX_INITIAL_CONTROL_COMMANDS,
13};
14#[cfg(any(test, windows))]
15use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
16#[cfg(any(test, windows))]
17use tokio::sync::mpsc as tokio_mpsc;
18
19use crate::{
20    connection::{read_response_frame_exact, Connection, ControlModeUpgrade, ControlTransition},
21    ClientError,
22};
23
24#[cfg(unix)]
25#[path = "control/output.rs"]
26mod output;
27
28impl Connection {
29    /// Requests a control-mode upgrade and, on success, yields the raw local
30    /// stream for tmux-compatible text control traffic.
31    pub fn begin_control_mode(
32        self,
33        mode: ControlMode,
34        client_terminal: ClientTerminalContext,
35    ) -> Result<ControlTransition, ClientError> {
36        self.begin_control_mode_with_initial_commands(mode, client_terminal, &[])
37    }
38
39    /// Requests a control-mode upgrade and writes command-line commands across
40    /// the upgrade boundary so the server can frame them as tmux argv commands.
41    pub fn begin_control_mode_with_initial_commands(
42        mut self,
43        mode: ControlMode,
44        client_terminal: ClientTerminalContext,
45        initial_commands: &[String],
46    ) -> Result<ControlTransition, ClientError> {
47        if initial_commands.len() > MAX_INITIAL_CONTROL_COMMANDS {
48            return Err(ClientError::Protocol(rmux_proto::RmuxError::Server(
49                format!(
50                    "too many initial control-mode commands: {} (maximum {MAX_INITIAL_CONTROL_COMMANDS})",
51                    initial_commands.len()
52                ),
53            )));
54        }
55        let initial_command_count = u32::try_from(initial_commands.len()).map_err(|_| {
56            ClientError::Protocol(rmux_proto::RmuxError::Server(
57                "too many initial control-mode commands".to_owned(),
58            ))
59        })?;
60        self.write_request(&Request::ControlMode(ControlModeRequest {
61            mode,
62            client_terminal,
63            initial_command_count,
64        }))?;
65        write_initial_control_commands(self.stream_mut(), initial_commands)?;
66        let response = read_response_frame_exact(self.stream_mut())?;
67
68        match response {
69            Response::ControlMode(response) => Ok(ControlTransition::Upgraded(
70                self.into_control_upgrade(response)?,
71            )),
72            other => Ok(ControlTransition::Rejected(other)),
73        }
74    }
75}
76
77fn write_initial_control_commands<W>(
78    stream: &mut W,
79    initial_commands: &[String],
80) -> Result<(), ClientError>
81where
82    W: Write,
83{
84    for command in initial_commands {
85        stream
86            .write_all(command.as_bytes())
87            .map_err(ClientError::Io)?;
88        stream.write_all(b"\n").map_err(ClientError::Io)?;
89    }
90    Ok(())
91}
92
93/// Drives a control-mode session using the process stdio streams.
94pub fn drive_control_mode(
95    upgrade: ControlModeUpgrade,
96    initial_commands: &[String],
97) -> Result<(), ClientError> {
98    let stdin = io::stdin();
99    let stdout = io::stdout();
100    drive_control_mode_with_stdio(upgrade, initial_commands, stdin, stdout)
101}
102
103/// Drives a control-mode session using explicit input and output streams.
104pub fn drive_control_mode_with_stdio<R, W>(
105    upgrade: ControlModeUpgrade,
106    initial_commands: &[String],
107    input: R,
108    mut output: W,
109) -> Result<(), ClientError>
110where
111    R: Read + Send + 'static,
112    W: Write + Send,
113{
114    let mode = upgrade.mode();
115    if mode.is_control_control() {
116        output
117            .write_all(CONTROL_CONTROL_START.as_bytes())
118            .map_err(ClientError::Io)?;
119        output.flush().map_err(ClientError::Io)?;
120    }
121
122    let stream = upgrade.into_stream();
123    let copy_result = drive_control_stream(stream, initial_commands, input, &mut output);
124    if copy_result.is_ok() && output_needs_suffix(mode) {
125        output
126            .write_all(CONTROL_CONTROL_END.as_bytes())
127            .map_err(ClientError::Io)?;
128        output.flush().map_err(ClientError::Io)?;
129    }
130
131    copy_result
132}
133
134#[cfg(unix)]
135fn drive_control_stream<R, W>(
136    stream: BlockingLocalStream,
137    initial_commands: &[String],
138    mut input: R,
139    output: &mut W,
140) -> Result<(), ClientError>
141where
142    R: Read + Send + 'static,
143    W: Write + Send,
144{
145    write_initial_commands(&stream, initial_commands)?;
146    ensure_blocking(&stream).map_err(ClientError::Io)?;
147    let mut writer = stream.try_clone().map_err(ClientError::Io)?;
148    let (stdin_done_tx, stdin_done_rx) = mpsc::sync_channel(1);
149    let stdin_thread = thread::spawn(move || {
150        let result = io::copy(&mut input, &mut writer).map(|_| ());
151        let _ = shutdown_write(&writer);
152        let _ = stdin_done_tx.send(result);
153    });
154
155    let copy_result = output::copy_control_output(stream, output).map_err(ClientError::Io);
156    let stdin_result = poll_input_thread(&stdin_done_rx)?;
157    if stdin_result.is_some() {
158        stdin_thread
159            .join()
160            .map_err(|_| ClientError::Io(io::Error::other("control input thread panicked")))?;
161    }
162
163    copy_result?;
164    if let Some(stdin_result) = stdin_result {
165        finish_control_input_after_output_closed(stdin_result).map_err(ClientError::Io)?;
166    }
167    Ok(())
168}
169
170#[cfg(windows)]
171const CONTROL_STDIN_QUEUE_CAPACITY: usize = 256;
172#[cfg(windows)]
173const CONTROL_STDOUT_QUEUE_CAPACITY: usize = 256;
174
175#[cfg(windows)]
176fn drive_control_stream<R, W>(
177    stream: BlockingLocalStream,
178    initial_commands: &[String],
179    input: R,
180    output: &mut W,
181) -> Result<(), ClientError>
182where
183    R: Read + Send + 'static,
184    W: Write + Send,
185{
186    let (input_tx, input_rx) = tokio_mpsc::channel(CONTROL_STDIN_QUEUE_CAPACITY);
187    let (output_tx, output_rx) = tokio_mpsc::channel(CONTROL_STDOUT_QUEUE_CAPACITY);
188    let (stdin_done_tx, stdin_done_rx) = mpsc::sync_channel(1);
189    let stdin_thread = thread::spawn(move || {
190        let result = copy_control_input(input, input_tx);
191        let _ = stdin_done_tx.send(result);
192    });
193
194    let (pipe, runtime) = stream.into_async_parts();
195    let copy_result = thread::scope(|scope| {
196        let output_thread = scope.spawn(move || write_queued_control_output(output, output_rx));
197        let copy_result = runtime
198            .block_on(drive_async_control(
199                pipe,
200                initial_commands,
201                input_rx,
202                output_tx,
203            ))
204            .map_err(ClientError::Io);
205        let output_result = output_thread
206            .join()
207            .map_err(|_| ClientError::Io(io::Error::other("control output thread panicked")))?;
208
209        copy_result?;
210        output_result.map_err(ClientError::Io)
211    });
212    let stdin_result = poll_input_thread(&stdin_done_rx)?;
213
214    if stdin_result.is_some() {
215        stdin_thread
216            .join()
217            .map_err(|_| ClientError::Io(io::Error::other("control input thread panicked")))?;
218    }
219
220    copy_result?;
221    if let Some(stdin_result) = stdin_result {
222        finish_control_input_after_output_closed(stdin_result).map_err(ClientError::Io)?;
223    }
224    Ok(())
225}
226
227fn output_needs_suffix(mode: ControlMode) -> bool {
228    mode.is_control_control()
229}
230
231fn finish_control_input_after_output_closed(result: io::Result<()>) -> io::Result<()> {
232    match result {
233        Err(error)
234            if matches!(
235                error.kind(),
236                io::ErrorKind::BrokenPipe
237                    | io::ErrorKind::ConnectionReset
238                    | io::ErrorKind::NotConnected
239            ) =>
240        {
241            Ok(())
242        }
243        result => result,
244    }
245}
246
247fn poll_input_thread(
248    stdin_done_rx: &mpsc::Receiver<io::Result<()>>,
249) -> Result<Option<io::Result<()>>, ClientError> {
250    match stdin_done_rx.try_recv() {
251        Ok(result) => Ok(Some(result)),
252        Err(mpsc::TryRecvError::Empty) => Ok(None),
253        Err(mpsc::TryRecvError::Disconnected) => Err(ClientError::Io(io::Error::other(
254            "control input thread terminated unexpectedly",
255        ))),
256    }
257}
258
259#[cfg(unix)]
260fn write_initial_commands(
261    stream: &BlockingLocalStream,
262    initial_commands: &[String],
263) -> Result<(), ClientError> {
264    if initial_commands.is_empty() {
265        return Ok(());
266    }
267
268    let mut writer = stream.try_clone().map_err(ClientError::Io)?;
269    for command in initial_commands {
270        writer
271            .write_all(command.as_bytes())
272            .and_then(|()| writer.write_all(b"\n"))
273            .map_err(ClientError::Io)?;
274    }
275    Ok(())
276}
277
278#[cfg(unix)]
279fn ensure_blocking(stream: &BlockingLocalStream) -> io::Result<()> {
280    stream.set_nonblocking(false)
281}
282
283#[cfg(unix)]
284fn shutdown_write(stream: &BlockingLocalStream) -> io::Result<()> {
285    stream.shutdown(std::net::Shutdown::Write)
286}
287
288#[cfg(windows)]
289fn copy_control_input<R>(mut input: R, input_tx: tokio_mpsc::Sender<Vec<u8>>) -> io::Result<()>
290where
291    R: Read,
292{
293    let mut buffer = [0_u8; 8192];
294    loop {
295        let bytes_read = match input.read(&mut buffer) {
296            Ok(0) => return Ok(()),
297            Ok(bytes_read) => bytes_read,
298            Err(error) if error.kind() == io::ErrorKind::Interrupted => continue,
299            Err(error) => return Err(error),
300        };
301
302        if input_tx
303            .blocking_send(buffer[..bytes_read].to_vec())
304            .is_err()
305        {
306            return Ok(());
307        }
308    }
309}
310
311#[cfg(any(test, windows))]
312async fn drive_async_control<Stream>(
313    stream: Stream,
314    initial_commands: &[String],
315    mut input_rx: tokio_mpsc::Receiver<Vec<u8>>,
316    output_tx: tokio_mpsc::Sender<Vec<u8>>,
317) -> io::Result<()>
318where
319    Stream: AsyncRead + AsyncWrite + Unpin,
320{
321    let mut input_closed = false;
322    let (mut reader, mut writer) = tokio::io::split(stream);
323    write_async_initial_commands(&mut writer, initial_commands).await?;
324    let mut buffer = [0_u8; 8192];
325
326    loop {
327        tokio::select! {
328            input = input_rx.recv(), if !input_closed => {
329                match input {
330                    Some(bytes) => {
331                        writer.write_all(&bytes).await?;
332                    }
333                    None => {
334                        writer.write_all(CONTROL_STDIN_EOF_MARKER.as_bytes()).await?;
335                        writer.write_all(b"\n").await?;
336                        writer.flush().await?;
337                        writer.shutdown().await?;
338                        input_closed = true;
339                    }
340                }
341            }
342            bytes_read = reader.read(&mut buffer) => {
343                let bytes_read = match bytes_read {
344                    Ok(bytes_read) => bytes_read,
345                    Err(error) if error.kind() == io::ErrorKind::BrokenPipe => return Ok(()),
346                    Err(error) => return Err(error),
347                };
348                if bytes_read == 0 {
349                    return Ok(());
350                }
351                send_control_output(&output_tx, &buffer[..bytes_read]).await?;
352            }
353        }
354    }
355}
356
357#[cfg(windows)]
358fn write_queued_control_output<W>(
359    output: &mut W,
360    mut output_rx: tokio_mpsc::Receiver<Vec<u8>>,
361) -> io::Result<()>
362where
363    W: Write,
364{
365    while let Some(bytes) = output_rx.blocking_recv() {
366        output.write_all(&bytes)?;
367        output.flush()?;
368    }
369    Ok(())
370}
371
372#[cfg(any(test, windows))]
373async fn send_control_output(
374    output_tx: &tokio_mpsc::Sender<Vec<u8>>,
375    bytes: &[u8],
376) -> io::Result<()> {
377    output_tx
378        .send(bytes.to_vec())
379        .await
380        .map_err(|_| io::Error::new(io::ErrorKind::BrokenPipe, "control output writer stopped"))
381}
382
383#[cfg(any(test, windows))]
384async fn write_async_initial_commands<Writer>(
385    writer: &mut Writer,
386    initial_commands: &[String],
387) -> io::Result<()>
388where
389    Writer: AsyncWrite + Unpin,
390{
391    for command in initial_commands {
392        writer.write_all(command.as_bytes()).await?;
393        writer.write_all(b"\n").await?;
394    }
395    writer.flush().await?;
396    Ok(())
397}
398
399#[cfg(all(test, unix))]
400mod tests {
401    use std::io::{self, Cursor, Read, Write};
402    use std::sync::mpsc;
403    use std::time::Duration;
404
405    use rmux_proto::{
406        ClientTerminalContext, ControlMode, ControlModeResponse, MAX_INITIAL_CONTROL_COMMANDS,
407    };
408
409    use super::{drive_control_mode_with_stdio, finish_control_input_after_output_closed};
410    use crate::connection::{Connection, ControlModeUpgrade};
411
412    #[test]
413    fn excessive_initial_commands_are_rejected_before_any_stream_write() {
414        let (client, mut server) = std::os::unix::net::UnixStream::pair().expect("socket pair");
415        let connection = Connection::new(client).expect("client connection");
416        let commands = vec![String::new(); MAX_INITIAL_CONTROL_COMMANDS + 1];
417
418        let error = connection
419            .begin_control_mode_with_initial_commands(
420                ControlMode::Plain,
421                ClientTerminalContext::default(),
422                &commands,
423            )
424            .expect_err("oversized command batch must fail locally");
425
426        assert!(
427            error
428                .to_string()
429                .contains("too many initial control-mode commands"),
430            "unexpected error: {error}"
431        );
432        let mut byte = [0_u8; 1];
433        assert_eq!(
434            server.read(&mut byte).expect("read closed client stream"),
435            0,
436            "client must not write a partial upgrade before rejecting the batch"
437        );
438    }
439
440    #[test]
441    fn closed_control_transport_supersedes_input_side_connection_errors() {
442        for kind in [
443            io::ErrorKind::BrokenPipe,
444            io::ErrorKind::ConnectionReset,
445            io::ErrorKind::NotConnected,
446        ] {
447            finish_control_input_after_output_closed(Err(io::Error::from(kind)))
448                .expect("closed server transport makes further control input irrelevant");
449        }
450
451        let error = finish_control_input_after_output_closed(Err(io::Error::from(
452            io::ErrorKind::InvalidData,
453        )))
454        .expect_err("unrelated input failures remain visible");
455        assert_eq!(error.kind(), io::ErrorKind::InvalidData);
456    }
457
458    #[test]
459    fn control_control_mode_wraps_output_with_dcs_sequences() {
460        let (left, right) = std::os::unix::net::UnixStream::pair().expect("socket pair");
461        let writer = std::thread::spawn(move || {
462            let mut right = right;
463            right.write_all(b"%exit\n").expect("write output");
464        });
465
466        let mut output = Vec::new();
467        drive_control_mode_with_stdio(
468            ControlModeUpgrade {
469                response: ControlModeResponse {
470                    mode: ControlMode::ControlControl,
471                },
472                stream: left,
473            },
474            &[],
475            Cursor::new(Vec::<u8>::new()),
476            &mut output,
477        )
478        .expect("control mode succeeds");
479        writer.join().expect("writer thread");
480
481        let rendered = String::from_utf8(output).expect("utf8");
482        assert!(rendered.starts_with(rmux_proto::CONTROL_CONTROL_START));
483        assert!(rendered.contains("%exit\n"));
484        assert!(rendered.ends_with(rmux_proto::CONTROL_CONTROL_END));
485    }
486
487    #[test]
488    fn control_mode_returns_after_server_exit_without_waiting_for_input_eof() {
489        let (left, right) = std::os::unix::net::UnixStream::pair().expect("socket pair");
490        let (input_reader, input_writer) =
491            std::os::unix::net::UnixStream::pair().expect("input socket pair");
492        let server = std::thread::spawn(move || {
493            let mut right = right;
494            right.write_all(b"%exit\n").expect("write exit");
495        });
496        let (done_tx, done_rx) = mpsc::channel();
497        let worker = std::thread::spawn(move || {
498            let mut output = Vec::new();
499            let result = drive_control_mode_with_stdio(
500                ControlModeUpgrade {
501                    response: ControlModeResponse {
502                        mode: ControlMode::Plain,
503                    },
504                    stream: left,
505                },
506                &[],
507                input_reader,
508                &mut output,
509            );
510            done_tx
511                .send((result, output))
512                .expect("report control mode result");
513        });
514
515        let done = done_rx.recv_timeout(Duration::from_secs(1));
516        drop(input_writer);
517        worker.join().expect("worker thread");
518        server.join().expect("server thread");
519
520        let (result, output) = done.expect("control mode should exit promptly");
521        result.expect("control mode succeeds");
522        assert_eq!(String::from_utf8(output).expect("utf8"), "%exit\n");
523    }
524}
525
526#[cfg(test)]
527#[path = "control/windows_tests.rs"]
528mod windows_tests;
rmux_client/control.rs

rmux_client/
control.rs
