Skip to content

Commit

Permalink
Improve stdio tunnel on windows
Browse files Browse the repository at this point in the history
- Handle CTRL+C to exit properly
- Restore terminal mode at exit
- Use logger to stderr
  • Loading branch information
erebe committed May 18, 2024
1 parent 0595e23 commit 9b82006
Show file tree
Hide file tree
Showing 3 changed files with 41 additions and 28 deletions.
43 changes: 22 additions & 21 deletions src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -724,27 +724,28 @@ async fn main() {
let args = Wstunnel::parse();

// Setup logging
match &args.commands {
// Disable logging if there is a stdio tunnel
Commands::Client(args)
if args
.local_to_remote
.iter()
.filter(|x| x.local_protocol == LocalProtocol::Stdio)
.count()
> 0 => {}
_ => {
let mut env_filter = EnvFilter::builder().parse(&args.log_lvl).expect("Invalid log level");
if !(args.log_lvl.contains("h2::") || args.log_lvl.contains("h2=")) {
env_filter =
env_filter.add_directive(Directive::from_str("h2::codec=off").expect("Invalid log directive"));
}
tracing_subscriber::fmt()
.with_ansi(args.no_color.is_none())
.with_env_filter(env_filter)
.init();
}
let mut env_filter = EnvFilter::builder().parse(&args.log_lvl).expect("Invalid log level");
if !(args.log_lvl.contains("h2::") || args.log_lvl.contains("h2=")) {
env_filter = env_filter.add_directive(Directive::from_str("h2::codec=off").expect("Invalid log directive"));
}
let logger = tracing_subscriber::fmt()
.with_ansi(args.no_color.is_none())
.with_env_filter(env_filter);

// stdio tunnel capture stdio, so need to log into stderr
if let Commands::Client(args) = &args.commands {
if args
.local_to_remote
.iter()
.filter(|x| x.local_protocol == LocalProtocol::Stdio)
.count()
> 0
{
logger.with_writer(io::stderr).init();
}
} else {
logger.init();
};

match args.commands {
Commands::Client(args) => {
Expand Down Expand Up @@ -1018,7 +1019,7 @@ async fn main() {
});
}
#[cfg(not(unix))]
LocalProtocol::Unix { path } => {
LocalProtocol::Unix { .. } => {
panic!("Unix socket is not available for non Unix platform")
}
LocalProtocol::Stdio
Expand Down
24 changes: 18 additions & 6 deletions src/stdio.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,9 @@
pub mod server {

use tokio_fd::AsyncFd;
use tracing::info;
pub async fn run_server() -> Result<(AsyncFd, AsyncFd), anyhow::Error> {
eprintln!("Starting STDIO server");
info!("Starting STDIO server");

let stdin = AsyncFd::try_from(nix::libc::STDIN_FILENO)?;
let stdout = AsyncFd::try_from(nix::libc::STDOUT_FILENO)?;
Expand All @@ -15,31 +16,39 @@ pub mod server {
#[cfg(not(unix))]
pub mod server {
use bytes::BytesMut;
use log::error;
use scopeguard::guard;
use std::io::{Read, Write};
use std::{io, thread};
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite};
use tokio::task::LocalSet;
use tokio_stream::wrappers::UnboundedReceiverStream;
use tokio_util::io::StreamReader;
use tracing::info;

pub async fn run_server() -> Result<(impl AsyncRead, impl AsyncWrite), anyhow::Error> {
eprintln!("Starting STDIO server");
info!("Starting STDIO server. Press ctrl+c twice to exit");

crossterm::terminal::enable_raw_mode()?;

let stdin = io::stdin();
let (send, recv) = tokio::sync::mpsc::unbounded_channel();
thread::spawn(move || {
let _restore_terminal = guard((), move |_| {
let _ = crossterm::terminal::disable_raw_mode();
});
let stdin = stdin;
let mut stdin = stdin.lock();
let mut buf = [0u8; 65536];

loop {
let n = stdin.read(&mut buf).unwrap();
if n == 0 {
let n = stdin.read(&mut buf).unwrap_or(0);
if n == 0 || (n == 1 && buf[0] == 3) {
// ctrl+c send char 3
break;
}
if let Err(err) = send.send(Result::<_, io::Error>::Ok(BytesMut::from(&buf[..n]))) {
eprintln!("Failed send inout: {:?}", err);
error!("Failed send inout: {:?}", err);
break;
}
}
Expand All @@ -50,6 +59,9 @@ pub mod server {
let rt = tokio::runtime::Handle::current();
thread::spawn(move || {
let task = async move {
let _restore_terminal = guard((), move |_| {
let _ = crossterm::terminal::disable_raw_mode();
});
let mut stdout = io::stdout().lock();
let mut buf = [0u8; 65536];
loop {
Expand All @@ -62,7 +74,7 @@ pub mod server {
}

if let Err(err) = stdout.write_all(&buf[..n]) {
eprintln!("Failed to write to stdout: {:?}", err);
error!("Failed to write to stdout: {:?}", err);
break;
};
let _ = stdout.flush();
Expand Down
2 changes: 1 addition & 1 deletion src/tunnel/server.rs
Original file line number Diff line number Diff line change
Expand Up @@ -167,7 +167,7 @@ async fn run_tunnel(
Ok((remote, Box::pin(local_rx), Box::pin(local_tx)))
}
#[cfg(not(unix))]
LocalProtocol::ReverseUnix { ref path } => {
LocalProtocol::ReverseUnix { .. } => {
error!("Received an unsupported target protocol {:?}", remote);
Err(anyhow::anyhow!("Invalid upgrade request"))
}
Expand Down

0 comments on commit 9b82006

Please sign in to comment.