116 lines
3.7 KiB
Rust
116 lines
3.7 KiB
Rust
#[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))
|
|
}
|
|
}
|