From ac4e92dfc4e18b648dd267cd2c92bb4608135fdd Mon Sep 17 00:00:00 2001 From: Marcin Kulik Date: Wed, 18 Jun 2025 17:58:03 +0200 Subject: [PATCH] Refactor tty module --- src/cmd/session.rs | 2 +- src/player.rs | 3 +- src/session.rs | 8 +- src/tty.rs | 570 +++++++++++++++++++++++++++++---------------- 4 files changed, 375 insertions(+), 208 deletions(-) diff --git a/src/cmd/session.rs b/src/cmd/session.rs index 37b4221..88a3e6d 100644 --- a/src/cmd/session.rs +++ b/src/cmd/session.rs @@ -189,7 +189,7 @@ impl cli::Session { } async fn get_term_info(&self) -> Result { - let tty = self.get_tty(false).await?; + let mut tty = self.get_tty(false).await?; Ok(TermInfo { type_: env::var("TERM").ok(), diff --git a/src/player.rs b/src/player.rs index 3efe61a..a2d7fe9 100644 --- a/src/player.rs +++ b/src/player.rs @@ -1,10 +1,11 @@ use anyhow::Result; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::sync::mpsc; use tokio::time::{self, Duration, Instant}; use crate::asciicast::{self, Event, EventData}; use crate::config::Key; -use crate::tty::{DevTty, Tty}; +use crate::tty::DevTty; pub struct KeyBindings { pub quit: Key, diff --git a/src/session.rs b/src/session.rs index 95be153..d1d3564 100644 --- a/src/session.rs +++ b/src/session.rs @@ -8,6 +8,7 @@ use nix::sys::wait::{WaitPidFlag, WaitStatus}; use signal_hook::consts::signal::*; use signal_hook_tokio::Signals; use tokio::io; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::sync::mpsc; use tokio::time::Instant; use tracing::error; @@ -138,6 +139,7 @@ impl Session { let mut input: Vec = Vec::with_capacity(BUF_SIZE); let mut output: Vec = Vec::with_capacity(BUF_SIZE); let mut wait_status = None; + let (mut tty_reader, mut tty_writer) = tty.split(); loop { tokio::select! { @@ -157,7 +159,7 @@ impl Session { input.drain(..n); } - result = tty.read(&mut input_buf) => { + result = tty_reader.read(&mut input_buf) => { let n = result?; if n > 0 { @@ -169,7 +171,7 @@ impl Session { } } - result = tty.write(&output), if !output.is_empty() => { + result = tty_writer.write(&output), if !output.is_empty() => { let n = result?; output.drain(..n); } @@ -203,7 +205,7 @@ impl Session { if !output.is_empty() { self.handle_output(&output).await; - let _ = tty.write_all(&output).await; + let _ = tty_writer.write_all(&output).await; } let wait_status = match wait_status { diff --git a/src/tty.rs b/src/tty.rs index a14b06b..6db2ccd 100644 --- a/src/tty.rs +++ b/src/tty.rs @@ -1,8 +1,9 @@ use std::fs::File; -use std::future::pending; use std::io::{Read, Write}; use std::os::fd::{AsFd, AsRawFd}; use std::os::unix::fs::OpenOptionsExt; +use std::pin::Pin; +use std::task::{ready, Context, Poll}; use async_trait::async_trait; use nix::libc; @@ -10,26 +11,13 @@ use nix::pty::Winsize; use nix::sys::termios::{self, SetArg, Termios}; use rgb::RGB8; use tokio::io::unix::AsyncFd; -use tokio::io::{self, Interest}; +use tokio::io::{self, AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, ReadBuf}; use tokio::time::{self, Duration}; const QUERY_READ_TIMEOUT: u64 = 500; const COLORS_QUERY: &str = "\x1b]10;?\x07\x1b]11;?\x07\x1b]4;0;?\x07\x1b]4;1;?\x07\x1b]4;2;?\x07\x1b]4;3;?\x07\x1b]4;4;?\x07\x1b]4;5;?\x07\x1b]4;6;?\x07\x1b]4;7;?\x07\x1b]4;8;?\x07\x1b]4;9;?\x07\x1b]4;10;?\x07\x1b]4;11;?\x07\x1b]4;12;?\x07\x1b]4;13;?\x07\x1b]4;14;?\x07\x1b]4;15;?\x07"; const XTVERSION_QUERY: &str = "\x1b[>0q"; -pub struct DevTty { - file: AsyncFd, - settings: libc::termios, -} - -pub struct NullTty; - -pub struct FixedSizeTty { - inner: T, - cols: Option, - rows: Option, -} - #[derive(Clone, Copy, Debug, PartialEq)] pub struct TtySize(pub u16, pub u16); @@ -40,14 +28,50 @@ pub struct TtyTheme { pub palette: Vec, } +pub struct DevTty { + file: AsyncFd, + settings: libc::termios, +} + +pub struct DevTtyReadHalf<'a> { + tty: &'a DevTty, +} + +pub struct DevTtyWriteHalf<'a> { + tty: &'a DevTty, +} + +pub struct NullTty; + +pub struct NullTtyReadHalf; + +pub struct NullTtyWriteHalf; + +pub struct FixedSizeTty { + inner: T, + cols: Option, + rows: Option, +} + +pub struct FixedSizeTtyReadHalf<'a> { + inner: Box, +} + +pub struct FixedSizeTtyWriteHalf<'a> { + inner: Box, +} + #[async_trait] pub trait Tty { fn get_size(&self) -> Winsize; - async fn get_theme(&self) -> Option; - async fn get_version(&self) -> Option; - async fn read<'e>(&self, buffer: &'e mut [u8]) -> io::Result; - async fn write<'e>(&self, buffer: &'e [u8]) -> io::Result; - async fn write_all<'e>(&self, buffer: &'e [u8]) -> io::Result<()>; + async fn get_theme(&mut self) -> Option; + async fn get_version(&mut self) -> Option; + fn split( + &self, + ) -> ( + Box, + Box, + ); } impl Default for TtySize { @@ -99,26 +123,27 @@ impl DevTty { Ok(Self { file, settings }) } - async fn query(&self, query: &str) -> anyhow::Result> { + async fn query(&mut self, query: &str) -> anyhow::Result> { let mut query = query.to_string().into_bytes(); query.extend_from_slice(b"\x1b[c"); let mut query = &query[..]; let mut response = Vec::new(); let mut buf = [0u8; 1024]; + let (mut reader, mut writer) = self.split(); loop { tokio::select! { - result = self.read(&mut buf) => { + result = reader.read(&mut buf) => { let n = result?; response.extend_from_slice(&buf[..n]); - if let Some(len) = self.complete_response_len(&response) { + if let Some(len) = complete_da_response_len(&response) { response.truncate(len); break; } } - result = self.write(query), if !query.is_empty() => { + result = writer.write(query), if !query.is_empty() => { let n = result?; query = &query[n..]; } @@ -132,35 +157,6 @@ impl DevTty { Ok(response) } - fn complete_response_len(&self, response: &[u8]) -> Option { - let mut reversed = response.iter().rev(); - let mut includes_da_response = false; - let mut da_response_len = 0; - - if let Some(b'c') = reversed.next() { - da_response_len += 1; - - for b in reversed { - if *b == b'[' { - includes_da_response = true; - break; - } - - if *b != b';' && *b != b'?' && !b.is_ascii_digit() { - break; - } - - da_response_len += 1; - } - } - - if includes_da_response { - Some(response.len() - da_response_len - 2) - } else { - None - } - } - pub async fn resize(&mut self, size: TtySize) -> io::Result<()> { let xtwinops_seq = format!("\x1b[8;{};{}t", size.1, size.0); self.write_all(xtwinops_seq.as_bytes()).await?; @@ -169,6 +165,296 @@ impl DevTty { } } +impl AsyncRead for DevTty { + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + loop { + let mut guard = ready!(self.file.poll_read_ready(cx))?; + let unfilled = buf.initialize_unfilled(); + + match guard.try_io(|inner| inner.get_ref().read(unfilled)) { + Ok(Ok(len)) => { + buf.advance(len); + return Poll::Ready(Ok(())); + } + + Ok(Err(err)) => return Poll::Ready(Err(err)), + Err(_would_block) => continue, + } + } + } +} + +impl AsyncWrite for DevTty { + fn poll_write( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + loop { + let mut guard = ready!(self.file.poll_write_ready(cx))?; + + match guard.try_io(|inner| inner.get_ref().write(buf)) { + Ok(result) => return Poll::Ready(result), + Err(_would_block) => continue, + } + } + } + + fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } +} + +impl Drop for DevTty { + fn drop(&mut self) { + let termios = Termios::from(self.settings); + let _ = termios::tcsetattr(self.file.as_fd(), SetArg::TCSANOW, &termios); + } +} + +#[async_trait] +impl Tty for DevTty { + fn get_size(&self) -> Winsize { + let mut winsize = Winsize { + ws_row: 24, + ws_col: 80, + ws_xpixel: 0, + ws_ypixel: 0, + }; + + unsafe { libc::ioctl(self.file.as_raw_fd(), libc::TIOCGWINSZ, &mut winsize) }; + + winsize + } + + async fn get_theme(&mut self) -> Option { + let response = self.query(COLORS_QUERY).await.ok()?; + let response = String::from_utf8_lossy(response.as_slice()); + let mut colors = response.match_indices("rgb:"); + let (idx, _) = colors.next()?; + let fg = parse_color(&response[idx + 4..])?; + let (idx, _) = colors.next()?; + let bg = parse_color(&response[idx + 4..])?; + let mut palette = Vec::new(); + + for _ in 0..16 { + let (idx, _) = colors.next()?; + let color = parse_color(&response[idx + 4..])?; + palette.push(color); + } + + Some(TtyTheme { fg, bg, palette }) + } + + async fn get_version(&mut self) -> Option { + let response = self.query(XTVERSION_QUERY).await.ok()?; + + if let [b'\x1b', b'P', b'>', b'|', version @ .., b'\x1b', b'\\'] = &response[..] { + Some(String::from_utf8_lossy(version).to_string()) + } else { + None + } + } + + fn split( + &self, + ) -> ( + Box, + Box, + ) { + let reader = DevTtyReadHalf { tty: self }; + let writer = DevTtyWriteHalf { tty: self }; + + (Box::new(reader), Box::new(writer)) + } +} + +impl AsyncRead for DevTtyReadHalf<'_> { + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + loop { + let mut guard = ready!(self.tty.file.poll_read_ready(cx))?; + let unfilled = buf.initialize_unfilled(); + + match guard.try_io(|inner| inner.get_ref().read(unfilled)) { + Ok(Ok(len)) => { + buf.advance(len); + return Poll::Ready(Ok(())); + } + + Ok(Err(err)) => return Poll::Ready(Err(err)), + Err(_would_block) => continue, + } + } + } +} + +impl AsyncWrite for DevTtyWriteHalf<'_> { + fn poll_write( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + loop { + let mut guard = ready!(self.tty.file.poll_write_ready(cx))?; + + match guard.try_io(|inner| inner.get_ref().write(buf)) { + Ok(result) => return Poll::Ready(result), + Err(_would_block) => continue, + } + } + } + + fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } +} + +impl FixedSizeTty { + pub fn new(inner: T, cols: Option, rows: Option) -> Self { + Self { inner, cols, rows } + } +} + +#[async_trait] +impl Tty for NullTty { + fn get_size(&self) -> Winsize { + Winsize { + ws_row: 24, + ws_col: 80, + ws_xpixel: 0, + ws_ypixel: 0, + } + } + + async fn get_theme(&mut self) -> Option { + None + } + + async fn get_version(&mut self) -> Option { + None + } + + fn split( + &self, + ) -> ( + Box, + Box, + ) { + (Box::new(NullTtyReadHalf), Box::new(NullTtyWriteHalf)) + } +} + +impl AsyncRead for NullTtyReadHalf { + fn poll_read( + self: Pin<&mut Self>, + _cx: &mut Context<'_>, + _buf: &mut ReadBuf<'_>, + ) -> Poll> { + Poll::Pending + } +} + +impl AsyncWrite for NullTtyWriteHalf { + fn poll_write( + self: Pin<&mut Self>, + _cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + Poll::Ready(Ok(buf.len())) + } + + fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } +} + +#[async_trait] +impl Tty for FixedSizeTty { + fn get_size(&self) -> Winsize { + let mut winsize = self.inner.get_size(); + + if let Some(cols) = self.cols { + winsize.ws_col = cols; + } + + if let Some(rows) = self.rows { + winsize.ws_row = rows; + } + + winsize + } + + async fn get_theme(&mut self) -> Option { + self.inner.get_theme().await + } + + async fn get_version(&mut self) -> Option { + self.inner.get_version().await + } + + fn split( + &self, + ) -> ( + Box, + Box, + ) { + let (reader, writer) = self.inner.split(); + + ( + Box::new(FixedSizeTtyReadHalf { inner: reader }), + Box::new(FixedSizeTtyWriteHalf { inner: writer }), + ) + } +} + +impl AsyncRead for FixedSizeTtyReadHalf<'_> { + fn poll_read( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + Pin::new(&mut self.inner).poll_read(cx, buf) + } +} + +impl AsyncWrite for FixedSizeTtyWriteHalf<'_> { + fn poll_write( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + Pin::new(&mut self.inner).poll_write(cx, buf) + } + + fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.inner).poll_flush(cx) + } + + fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.inner).poll_shutdown(cx) + } +} + fn make_raw(fd: F) -> anyhow::Result { let termios = termios::tcgetattr(fd.as_fd())?; let mut raw_termios = termios.clone(); @@ -178,10 +464,32 @@ fn make_raw(fd: F) -> anyhow::Result { Ok(termios.into()) } -impl Drop for DevTty { - fn drop(&mut self) { - let termios = Termios::from(self.settings); - let _ = termios::tcsetattr(self.file.as_fd(), SetArg::TCSANOW, &termios); +fn complete_da_response_len(response: &[u8]) -> Option { + let mut reversed = response.iter().rev(); + let mut includes_da_response = false; + let mut da_response_len = 0; + + if let Some(b'c') = reversed.next() { + da_response_len += 1; + + for b in reversed { + if *b == b'[' { + includes_da_response = true; + break; + } + + if *b != b';' && *b != b'?' && !b.is_ascii_digit() { + break; + } + + da_response_len += 1; + } + } + + if includes_da_response { + Some(response.len() - da_response_len - 2) + } else { + None } } @@ -202,154 +510,10 @@ fn parse_color(rgb: &str) -> Option { Some(RGB8::new(r, g, b)) } -#[async_trait] -impl Tty for DevTty { - fn get_size(&self) -> Winsize { - let mut winsize = Winsize { - ws_row: 24, - ws_col: 80, - ws_xpixel: 0, - ws_ypixel: 0, - }; - - unsafe { libc::ioctl(self.file.as_raw_fd(), libc::TIOCGWINSZ, &mut winsize) }; - - winsize - } - - async fn get_theme(&self) -> Option { - let response = self.query(COLORS_QUERY).await.ok()?; - let response = String::from_utf8_lossy(response.as_slice()); - let mut colors = response.match_indices("rgb:"); - let (idx, _) = colors.next()?; - let fg = parse_color(&response[idx + 4..])?; - let (idx, _) = colors.next()?; - let bg = parse_color(&response[idx + 4..])?; - let mut palette = Vec::new(); - - for _ in 0..16 { - let (idx, _) = colors.next()?; - let color = parse_color(&response[idx + 4..])?; - palette.push(color); - } - - Some(TtyTheme { fg, bg, palette }) - } - - async fn get_version(&self) -> Option { - let response = self.query(XTVERSION_QUERY).await.ok()?; - - if let [b'\x1b', b'P', b'>', b'|', version @ .., b'\x1b', b'\\'] = &response[..] { - Some(String::from_utf8_lossy(version).to_string()) - } else { - None - } - } - - async fn read<'e>(&self, buffer: &'e mut [u8]) -> io::Result { - self.file - .async_io(Interest::READABLE, |mut file| file.read(buffer)) - .await - } - - async fn write<'e>(&self, buffer: &'e [u8]) -> io::Result { - self.file - .async_io(Interest::WRITABLE, |mut file| file.write(buffer)) - .await - } - - async fn write_all<'e>(&self, buffer: &'e [u8]) -> io::Result<()> { - let mut buffer = buffer; - - while !buffer.is_empty() { - let n = self.write(buffer).await?; - buffer = &buffer[n..]; - } - - Ok(()) - } -} - -#[async_trait] -impl Tty for NullTty { - fn get_size(&self) -> Winsize { - Winsize { - ws_row: 24, - ws_col: 80, - ws_xpixel: 0, - ws_ypixel: 0, - } - } - - async fn get_theme(&self) -> Option { - None - } - - async fn get_version(&self) -> Option { - None - } - - async fn read<'e>(&self, _buffer: &'e mut [u8]) -> io::Result { - pending::<()>().await; - unreachable!() - } - - async fn write<'e>(&self, buffer: &'e [u8]) -> io::Result { - Ok(buffer.len()) - } - - async fn write_all<'e>(&self, _buffer: &'e [u8]) -> io::Result<()> { - Ok(()) - } -} - -impl FixedSizeTty { - pub fn new(inner: T, cols: Option, rows: Option) -> Self { - Self { inner, cols, rows } - } -} - -#[async_trait] -impl Tty for FixedSizeTty { - fn get_size(&self) -> Winsize { - let mut winsize = self.inner.get_size(); - - if let Some(cols) = self.cols { - winsize.ws_col = cols; - } - - if let Some(rows) = self.rows { - winsize.ws_row = rows; - } - - winsize - } - - async fn get_theme(&self) -> Option { - self.inner.get_theme().await - } - - async fn get_version(&self) -> Option { - self.inner.get_version().await - } - - async fn read<'e>(&self, buffer: &'e mut [u8]) -> io::Result { - self.inner.read(buffer).await - } - - async fn write<'e>(&self, buffer: &'e [u8]) -> io::Result { - self.inner.write(buffer).await - } - - async fn write_all<'e>(&self, buffer: &'e [u8]) -> io::Result<()> { - self.inner.write_all(buffer).await - } -} - #[cfg(test)] mod tests { - use super::{FixedSizeTty, Tty}; - use crate::tty::NullTty; + use super::{FixedSizeTty, NullTty, Tty}; + use rgb::RGB8; #[test]