use anyctx::AnyCtx;
use futures_concurrency::future::TryJoin;
use futures_util::AsyncReadExt;
use nursery_macro::nursery;
use sillad::listener::Listener;
use smol::future::FutureExt;
use crate::{Config, litecopy::litecopy, session::open_conn};
pub async fn port_forward(ctx: &AnyCtx<Config>) -> anyhow::Result<()> {
let tasks = ctx
.init()
.port_forward
.iter()
.map(|fwd| async {
let mut listener = sillad::tcp::TcpListener::bind(fwd.listen).await?;
nursery!({
loop {
let listened = listener.accept().await?;
let connect = fwd.connect.clone();
let ctx = &ctx;
tracing::debug!(
listen = display(&fwd.listen),
connect = display(&fwd.connect),
"accepted port forward"
);
let fwd = fwd.clone();
spawn!(async move {
let connected = open_conn(ctx, "tcp", &connect).await?;
tracing::debug!(
listen = display(&fwd.listen),
connect = display(&fwd.connect),
"started port forward"
);
let (read_listened, write_listened) = listened.split();
let (read_connected, write_connected) = connected.split();
litecopy(read_listened, write_connected)
.race(litecopy(read_connected, write_listened))
.await?;
anyhow::Ok(())
})
.detach();
}
})
})
.collect::<Vec<_>>();
if tasks.is_empty() {
smol::future::pending().await
} else {
tasks.try_join().await.map(|_: Vec<()>| ())
}
}