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

rmux_client/
control.rs
