wstunnel/src/stdio.rs

#[cfg(unix)]
pub mod server {
    use std::pin::Pin;
    use std::process;
    use std::task::{Context, Poll};
    use tokio::io::{AsyncRead, ReadBuf};
    use tokio_fd::AsyncFd;
    use tracing::info;

    pub struct AbortOnDropStdin {
        stdin: AsyncFd,
    }

    // Wrapper around stdin is needed in order to properly abort the process in case the tunnel drop.
    // As we are going to launch the tunnel in a threadpool, we cant know when the tunnel is dropped.
    impl Drop for AbortOnDropStdin {
        fn drop(&mut self) {
            // Hackish !
            process::exit(0);
        }
    }

    impl AsyncRead for AbortOnDropStdin {
        fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<std::io::Result<()>> {
            unsafe { self.map_unchecked_mut(|s| &mut s.stdin) }.poll_read(cx, buf)
        }
    }

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

        let stdin = AsyncFd::try_from(nix::libc::STDIN_FILENO)?;
        let stdout = AsyncFd::try_from(nix::libc::STDOUT_FILENO)?;

        Ok((AbortOnDropStdin { stdin }, stdout))
    }
}

#[cfg(not(unix))]
pub mod server {
    use bytes::BytesMut;
    use log::error;
    use scopeguard::guard;
    use std::io::{Read, Write};
    use std::{io, process, 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> {
        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();
                process::exit(0);
            });
            let stdin = stdin;
            let mut stdin = stdin.lock();
            let mut buf = [0u8; 65536];

            loop {
                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]))) {
                    error!("Failed send inout: {:?}", err);
                    break;
                }
            }
        });
        let stdin = StreamReader::new(UnboundedReceiverStream::new(recv));

        let (stdout, mut recv) = tokio::io::duplex(65536);
        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 {
                    let Ok(n) = recv.read(&mut buf).await else {
                        break;
                    };

                    if n == 0 {
                        break;
                    }

                    if let Err(err) = stdout.write_all(&buf[..n]) {
                        error!("Failed to write to stdout: {:?}", err);
                        break;
                    };
                    let _ = stdout.flush();
                }
            };

            let local = LocalSet::new();
            local.spawn_local(task);

            rt.block_on(local);
        });

        Ok((stdin, stdout))
    }
}