diff --git a/Cargo.lock b/Cargo.lock index 29e4046..b394216 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -29,6 +29,12 @@ version = "1.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" +[[package]] +name = "bitflags" +version = "2.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed570934406eb16438a4e976b1b4500774099c13b8cb96eec99f620f05090ddf" + [[package]] name = "block-buffer" version = "0.10.4" @@ -91,6 +97,27 @@ dependencies = [ "crypto-common", ] +[[package]] +name = "dirs-next" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b98cf8ebf19c3d1b223e151f99a4f9f0690dca41414773390fc824184ac833e1" +dependencies = [ + "cfg-if", + "dirs-sys-next", +] + +[[package]] +name = "dirs-sys-next" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ebda144c4fe02d1f7ea1a7d9641b6fc6b580adcfa024ae48797ecdeb6825b4d" +dependencies = [ + "libc", + "redox_users", + "winapi", +] + [[package]] name = "downcast-rs" version = "1.2.0" @@ -242,6 +269,15 @@ dependencies = [ "libc", ] +[[package]] +name = "home" +version = "0.5.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3d1354bf6b7235cb4a0576c2619fd4ed18183f689b12b006a0ee7329eeff9a5" +dependencies = [ + "windows-sys 0.52.0", +] + [[package]] name = "http" version = "0.2.9" @@ -296,6 +332,17 @@ version = "0.2.142" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6a987beff54b60ffa6d51982e1aa1146bc42f19bd26be28b0586f252fccf5317" +[[package]] +name = "libredox" +version = "0.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85c833ca1e66078851dba29046874e38f08b2c883700aa29a03ddd3b23814ee8" +dependencies = [ + "bitflags 2.4.2", + "libc", + "redox_syscall", +] + [[package]] name = "log" version = "0.4.17" @@ -329,7 +376,7 @@ dependencies = [ "libc", "log", "wasi", - "windows-sys", + "windows-sys 0.45.0", ] [[package]] @@ -339,7 +386,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f346ff70e7dbfd675fe90590b92d59ef2de15a8779ae305ebcbfd3f0caf59be4" dependencies = [ "autocfg", - "bitflags", + "bitflags 1.3.2", "cfg-if", "libc", "memoffset", @@ -381,7 +428,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "806ee80c2a03dbe1a9fb9534f8d19e4c0546b790cde8fd1fea9d6390644cb0be" dependencies = [ "anyhow", - "bitflags", + "bitflags 1.3.2", "downcast-rs", "filedescriptor", "lazy_static", @@ -449,6 +496,26 @@ dependencies = [ "getrandom", ] +[[package]] +name = "redox_syscall" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4722d768eff46b75989dd134e5c353f0d6296e5aaa3132e776cbdb56be7731aa" +dependencies = [ + "bitflags 1.3.2", +] + +[[package]] +name = "redox_users" +version = "0.4.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a18479200779601e498ada4e8c1e1f50e3ee19deb0259c25825a98b5603b2cb4" +dependencies = [ + "getrandom", + "libredox", + "thiserror", +] + [[package]] name = "regex" version = "1.8.1" @@ -466,6 +533,12 @@ version = "0.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a5996294f19bd3aae0453a862ad728f60e6600695733dd5df01da90c54363a3c" +[[package]] +name = "rustversion" +version = "1.0.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7ffc183a10b4478d04cbbbfc96d0873219d962dd5accaff2ffbd4ceb7df837f4" + [[package]] name = "serial" version = "0.4.0" @@ -535,6 +608,15 @@ version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "24188a676b6ae68c3b2cb3a01be17fbf7240ce009799bb56d5b1409051e78fde" +[[package]] +name = "signal-hook-registry" +version = "1.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d8229b473baa5980ac72ef434c4415e70c4b5e71b423043adb4ba059f89c99a1" +dependencies = [ + "libc", +] + [[package]] name = "slab" version = "0.4.8" @@ -565,6 +647,26 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "term" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c59df8ac95d96ff9bede18eb7300b0fda5e5d8d90960e76f8e14ae765eedbf1f" +dependencies = [ + "dirs-next", + "rustversion", + "winapi", +] + +[[package]] +name = "termini" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2ad441d87dd98bc5eeb31cf2fb7e4839968763006b478efb38668a3bf9da0d59" +dependencies = [ + "home", +] + [[package]] name = "termios" version = "0.2.2" @@ -602,6 +704,8 @@ dependencies = [ "libc", "portable-pty", "regex", + "term", + "termini", "thiserror", "tokio", "tokio-tungstenite", @@ -635,9 +739,10 @@ dependencies = [ "mio", "num_cpus", "pin-project-lite", + "signal-hook-registry", "socket2", "tokio-macros", - "windows-sys", + "windows-sys 0.45.0", ] [[package]] @@ -766,7 +871,16 @@ version = "0.45.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "75283be5efb2831d37ea142365f009c02ec203cd29a3ebecbc093d52315b66d0" dependencies = [ - "windows-targets", + "windows-targets 0.42.2", +] + +[[package]] +name = "windows-sys" +version = "0.52.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d" +dependencies = [ + "windows-targets 0.52.0", ] [[package]] @@ -775,13 +889,28 @@ version = "0.42.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8e5180c00cd44c9b1c88adb3693291f1cd93605ded80c250a75d472756b4d071" dependencies = [ - "windows_aarch64_gnullvm", - "windows_aarch64_msvc", - "windows_i686_gnu", - "windows_i686_msvc", - "windows_x86_64_gnu", - "windows_x86_64_gnullvm", - "windows_x86_64_msvc", + "windows_aarch64_gnullvm 0.42.2", + "windows_aarch64_msvc 0.42.2", + "windows_i686_gnu 0.42.2", + "windows_i686_msvc 0.42.2", + "windows_x86_64_gnu 0.42.2", + "windows_x86_64_gnullvm 0.42.2", + "windows_x86_64_msvc 0.42.2", +] + +[[package]] +name = "windows-targets" +version = "0.52.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8a18201040b24831fbb9e4eb208f8892e1f50a37feb53cc7ff887feb8f50e7cd" +dependencies = [ + "windows_aarch64_gnullvm 0.52.0", + "windows_aarch64_msvc 0.52.0", + "windows_i686_gnu 0.52.0", + "windows_i686_msvc 0.52.0", + "windows_x86_64_gnu 0.52.0", + "windows_x86_64_gnullvm 0.52.0", + "windows_x86_64_msvc 0.52.0", ] [[package]] @@ -790,42 +919,84 @@ version = "0.42.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "597a5118570b68bc08d8d59125332c54f1ba9d9adeedeef5b99b02ba2b0698f8" +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.52.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb7764e35d4db8a7921e09562a0304bf2f93e0a51bfccee0bd0bb0b666b015ea" + [[package]] name = "windows_aarch64_msvc" version = "0.42.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e08e8864a60f06ef0d0ff4ba04124db8b0fb3be5776a5cd47641e942e58c4d43" +[[package]] +name = "windows_aarch64_msvc" +version = "0.52.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbaa0368d4f1d2aaefc55b6fcfee13f41544ddf36801e793edbbfd7d7df075ef" + [[package]] name = "windows_i686_gnu" version = "0.42.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c61d927d8da41da96a81f029489353e68739737d3beca43145c8afec9a31a84f" +[[package]] +name = "windows_i686_gnu" +version = "0.52.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a28637cb1fa3560a16915793afb20081aba2c92ee8af57b4d5f28e4b3e7df313" + [[package]] name = "windows_i686_msvc" version = "0.42.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "44d840b6ec649f480a41c8d80f9c65108b92d89345dd94027bfe06ac444d1060" +[[package]] +name = "windows_i686_msvc" +version = "0.52.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ffe5e8e31046ce6230cc7215707b816e339ff4d4d67c65dffa206fd0f7aa7b9a" + [[package]] name = "windows_x86_64_gnu" version = "0.42.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8de912b8b8feb55c064867cf047dda097f92d51efad5b491dfb98f6bbb70cb36" +[[package]] +name = "windows_x86_64_gnu" +version = "0.52.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d6fa32db2bc4a2f5abeacf2b69f7992cd09dca97498da74a151a3132c26befd" + [[package]] name = "windows_x86_64_gnullvm" version = "0.42.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "26d41b46a36d453748aedef1486d5c7a85db22e56aff34643984ea85514e94a3" +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.52.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1a657e1e9d3f514745a572a6846d3c7aa7dbe1658c056ed9c3344c4109a6949e" + [[package]] name = "windows_x86_64_msvc" version = "0.42.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9aec5da331524158c6d1a4ac0ab1541149c0b9505fde06423b02f5ef0106b9f0" +[[package]] +name = "windows_x86_64_msvc" +version = "0.52.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dff9641d1cd4be8d1a070daf9e3773c5f67e78b4d9d42263020c057706765c04" + [[package]] name = "winreg" version = "0.10.1" diff --git a/Cargo.toml b/Cargo.toml index cc2fce3..c271040 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -7,8 +7,10 @@ edition = "2021" futures = "0.3.28" libc = "0.2" regex = "1" +term = "0.7" +termini = "1" thiserror = "1" -tokio = { version = "1", features = ["io-std", "io-util", "macros", "rt", "rt-multi-thread", "sync", "time"] } +tokio = { version = "1", features = ["io-std", "io-util", "macros", "rt", "rt-multi-thread", "signal", "sync", "time"] } tokio-tungstenite = "0.20.1" [target.'cfg(target_os = "windows")'.dependencies] diff --git a/src/lib.rs b/src/lib.rs index 968442c..65e9355 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -12,12 +12,19 @@ pub use crate::input::EscapeSequence; pub use crate::raw::RawModeGuard; use crate::input::{stdin_read_task, stdin_relay_task}; + #[cfg(target_family = "unix")] use std::os::fd::AsRawFd as AsRawFdHandle; #[cfg(target_family = "windows")] use std::os::windows::io::AsRawHandle as AsRawFdHandle; -use futures::{SinkExt, StreamExt}; +use futures::stream::FuturesUnordered; +#[cfg(target_family = "unix")] +use tokio::signal::unix::{signal, SignalKind}; +#[cfg(target_family = "windows")] +use tokio::signal::windows::*; + +use futures::{FutureExt, SinkExt, StreamExt}; use thiserror::Error; use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}; use tokio::sync::mpsc; @@ -37,6 +44,8 @@ pub enum Error { StdoutWrite(#[from] std::io::Error), #[error("Server error: {0}")] ServerError(String), + #[error("Terminated by signal: {0}")] + Signal(&'static str), } /// A simple abstraction over a TTY's async I/O streams. @@ -141,7 +150,34 @@ impl Console { &mut self, upgraded: impl AsyncRead + AsyncWrite + Unpin, ) -> Result<(), Error> { + // need Signal structs to live at least as long as their futures + let mut signal_storage = Vec::new(); + + let mut signaled = FuturesUnordered::new(); + + #[cfg(target_family = "unix")] + { + signal_storage.push((signal(SignalKind::hangup())?, "HUP")); + signal_storage.push((signal(SignalKind::interrupt())?, "INT")); + signal_storage.push((signal(SignalKind::pipe())?, "PIPE")); + signal_storage.push((signal(SignalKind::quit())?, "QUIT")); + signal_storage.push((signal(SignalKind::terminate())?, "TERM")); + } + #[cfg(target_family = "windows")] + { + // no ctrl_c(), we're already in VT100 mode, and raw mode in that + signal_storage.push((WinCtrlSignal::CBreak(ctrl_break()?), "CTRL-BREAK")); + signal_storage.push((WinCtrlSignal::CClose(ctrl_close()?), "CTRL-CLOSE")); + signal_storage.push((WinCtrlSignal::CLogoff(ctrl_logoff()?), "CTRL-LOGOFF")); + signal_storage.push((WinCtrlSignal::CShutdown(ctrl_shutdown()?), "CTRL-SHUTDOWN")); + } + + for (s_fut, s_name) in &mut signal_storage { + signaled.push(s_fut.recv().then(|opt| async move { opt.map(|_| s_name) })); + } + let mut ws_stream = WebSocketStream::from_raw_socket(upgraded, Role::Client, None).await; + loop { tokio::select! { in_buf = self.read_stdin() => { @@ -182,6 +218,10 @@ impl Console { _ => continue, } } + Some(Some(signal_name)) = signaled.next() => { + eprint!("\r\nExiting on signal.\r\n"); + return Err(Error::Signal(signal_name)); + } } } // let _: the connection may have already been dropped at this point. @@ -190,6 +230,28 @@ impl Console { } } +// unfortunately tokio::signal makes these all separate types... +#[cfg(target_family = "windows")] +enum WinCtrlSignal { + CC(CtrlC), + CBreak(CtrlBreak), + CClose(CtrlClose), + CLogoff(CtrlLogoff), + CShutdown(CtrlShutdown), +} +#[cfg(target_family = "windows")] +impl WinCtrlSignal { + async fn recv(&mut self) -> Option<()> { + match self { + Self::CC(c) => c.recv().await, + Self::CBreak(c) => c.recv().await, + Self::CClose(c) => c.recv().await, + Self::CLogoff(c) => c.recv().await, + Self::CShutdown(c) => c.recv().await, + } + } +} + impl Drop for Console { fn drop(&mut self) { self.relay_handle.abort(); @@ -258,4 +320,39 @@ mod tests { // ...and end the event loop. timeout(ONE_SEC, join_handle).await.unwrap().unwrap(); } + + #[cfg(target_family = "unix")] + #[tokio::test] + async fn test_cleanup_on_signal() { + let (_in_testdrv, in_console) = tokio::io::duplex(16); + let (mut out_testdrv, out_console) = tokio::io::duplex(16); + let (ws_testdrv, ws_console) = tokio::io::duplex(16); + + let mut ws = WebSocketStream::from_raw_socket(ws_testdrv, Role::Server, None).await; + let mut console = Console::new_inner(in_console, out_console, None, None); + + let join_handle = + tokio::spawn(async move { console.attach_to_websocket(ws_console).await }); + + ws.send(Message::Binary(vec![1, 2, 3, 4, 5, 6])) + .await + .unwrap(); + + let mut read_buf = [0u8; 6]; + const ONE_SEC: Duration = Duration::from_secs(1); + timeout(ONE_SEC, out_testdrv.read_exact(&mut read_buf)) + .await + .unwrap() + .unwrap(); + assert_eq!(read_buf, [1, 2, 3, 4, 5, 6]); + + let syscall_return = unsafe { libc::kill(std::process::id() as libc::c_int, libc::SIGINT) }; + assert_eq!(syscall_return, 0); + + let Err(super::Error::Signal("INT")) = + timeout(ONE_SEC, join_handle).await.unwrap().unwrap() + else { + panic!("Expected SIGINT!") + }; + } } diff --git a/src/raw.rs b/src/raw.rs index 7f2d799..33c6fbb 100644 --- a/src/raw.rs +++ b/src/raw.rs @@ -90,11 +90,17 @@ mod platform_impl { unsafe { let r = SetConsoleMode(self.in_handle, self.in_mode); if r == 0 { - Err::<(), _>(std::io::Error::last_os_error()).unwrap(); + panic!( + "\r\n{}\r\n", + Error::SetConsoleMode(std::io::Error::last_os_error()) + ); } let r = SetConsoleMode(self.out_handle, self.out_mode); if r == 0 { - Err::<(), _>(std::io::Error::last_os_error()).unwrap(); + panic!( + "\r\n{}\r\n", + Error::SetConsoleMode(std::io::Error::last_os_error()) + ); } } } @@ -103,22 +109,29 @@ mod platform_impl { #[cfg(target_family = "unix")] mod platform_impl { - use std::os::fd::RawFd; + use std::os::fd::{FromRawFd, IntoRawFd, RawFd}; use thiserror::Error; #[derive(Error, Debug)] pub enum Error { #[error("tcgetattr(stdout, termios) call failed: {0}")] TcGetAttr(std::io::Error), - #[error("tcsetattr(stdout, TCSAFLUSH, termios) call failed: {0}")] - TcSetAttr(std::io::Error), + #[error("tcsetattr(stdout, {1}, termios) call failed: {0}")] + TcSetAttr(std::io::Error, &'static str), } /// Guard object that will set the terminal to raw mode and restore it - /// to its previous state when it's dropped. + /// to its previous state when it's dropped. If it is unable to restore + /// the previous termcap state, [Drop::drop] will panic. + /// + /// Additionally, if the terminfo database for the current terminal is + /// available, the reset sequences from it are emitted to return it from + /// any unknown state. Failures in finding these in terminfo are ignored, + /// but failure to output them to the terminal will also panic. pub struct RawModeGuard(libc::c_int, libc::termios); impl RawModeGuard { + /// Attach to the terminal whose stdout is the given fd. pub(crate) fn new(fd: RawFd) -> Result { let termios = unsafe { let mut curr_termios = std::mem::zeroed(); @@ -134,7 +147,10 @@ mod platform_impl { libc::cfmakeraw(&mut raw_termios); let r = libc::tcsetattr(fd, libc::TCSAFLUSH, &raw_termios); if r == -1 { - return Err(Error::TcSetAttr(std::io::Error::last_os_error())); + return Err(Error::TcSetAttr( + std::io::Error::last_os_error(), + "TCSAFLUSH", + )); } } Ok(guard) @@ -143,9 +159,52 @@ mod platform_impl { impl Drop for RawModeGuard { fn drop(&mut self) { + // reset the termcaps to what they were before we were constructed. + // (similar to `stty sane` if the tty was that way to begin with) let r = unsafe { libc::tcsetattr(self.0, libc::TCSADRAIN, &self.1) }; if r == -1 { - Err::<(), _>(std::io::Error::last_os_error()).unwrap(); + // some \r\n because we might still be in a raw mode... + panic!( + "\r\n{}\r\n", + Error::TcSetAttr(std::io::Error::last_os_error(), "TCSADRAIN") + ); + } + // if we have a terminfo database and this terminal is in it, + // try to emit ANSI strings to reset from an indeterminate state. + if let Ok(info) = termini::TermInfo::from_env() { + use std::io::Write; + use termini::StringCapability::*; + let mut stdout = unsafe { std::fs::File::from_raw_fd(self.0) }; + for cap in [ExitAlternativeMode, CursorNormal, ExitAttributeMode] { + if let Some(s) = info.raw_string_cap(cap) { + stdout.write_all(s).ok(); + } + } + // disable mouse mode, which takes a parameter (see section + // "Parameterized Strings" in terminfo(5) - needs variable + // expansion support. in practice this will always result + // in b"\x1b[1006;1000l" or similar, but let's be thorough) + if let Some(val) = info.extended_cap("XM") { + let xm = match val { + termini::Value::RawString(raw) => raw, + termini::Value::Utf8String(s) => s.as_bytes(), + _ => &[], // do nothing + }; + if let Ok(exp) = term::terminfo::parm::expand( + xm, + &[term::terminfo::parm::Param::Number(0)], + &mut term::terminfo::parm::Variables::new(), + ) { + stdout.write_all(&exp).ok(); + } + } else { + // despite it all, some terminfo entries are missing "XM" + // this despite being likely to be used in conjunction with + // mouse functionality -- tmux and screen, for example + stdout.write_all(b"\x1b[1006;1000l").ok(); + } + // relinquish ownership again - we may not want to close on drop + stdout.into_raw_fd(); } } }