use core::time::Duration;
use hierr::Result;
use hirun::net::{AioFd, Fd, SocketAddr, TcpClient};
use hirun::runtime;
fn main() {
let opts = hiopt::options!["addr:", "a:", "cache:", "threads:", "t:", "buffer:", "b:", "help", "h"];
let args = unsafe { hiopt::raw_args_from_i8(hictor::args()) };
let mut listening_addr = None;
let mut server_addr = None;
let mut cache = 0_usize;
let mut nth = 0_usize;
let mut buflen = 1024_usize;
for opt in opts.opt_iter(&args[1..]) {
let (opt_idx, opt_value) = opt.unwrap();
match opt_idx {
0..=1 => listening_addr = Some(SocketAddr::inet_from(opt_value.unwrap()).unwrap()),
2 => cache = opt_value.unwrap().parse().unwrap(),
3..=4 => nth = opt_value.unwrap().parse().unwrap(),
5..=6 => buflen = opt_value.unwrap().parse().unwrap(),
7..=8 => return help(),
_ => unreachable!(),
}
}
for addr in opts.noopt_iter(&args[1..]) {
server_addr = Some(SocketAddr::inet_from(addr).unwrap());
break;
}
if listening_addr.is_none() || server_addr.is_none() {
return help();
}
let _ = runtime::Builder::new().max_cache(cache).nth(nth).build().unwrap();
let _ = runtime::block_on_with(
listening(
listening_addr.as_ref().unwrap(),
server_addr.as_ref().unwrap(),
buflen,
),
runtime::Attr::new().priority(1),
);
}
fn help() {
println!(
"Usage: {} options server_addr",
hictor::program_invocation_name()
);
println!("server_addr: <ipv4:port> | <[ipv6]:port>");
println!("options:");
println!("--addr | -a: port | <ipv4:port> | <[ipv6]:port>");
println!("--cache: max task-cache size");
println!("--threads | -t: defaut is 0");
println!("--buffer | -b: buflen, defaut is 1024");
println!("--help | -h: print help message");
}
async fn do_proxy(down: Fd, up: TcpClient, buflen: usize) -> Result<()> {
let Ok((up, cookie)) = up.wait_connected().await else {
println!("connect failed: {}", hierr::Error::last());
return Ok(());
};
let mut buf = vec![0; buflen];
let mut down = AioFd::new(&down);
let mut up = AioFd::new_with(&up, cookie);
let _ = down.set_linger(Some(Duration::new(0, 0)));
let _ = up.set_linger(Some(Duration::new(0, 0)));
down.copy_bidirectional(&mut up, &mut buf, Duration::new(5, 0))
.await;
Ok(())
}
async fn listening(addr: &SocketAddr, remote: &SocketAddr, buflen: usize) -> Result<()> {
let fd = Fd::tcp_server(addr)?;
println!("listening {:?}", addr);
let mut aio = AioFd::new(&fd);
while let Ok((down, _)) = aio.accept().await {
let up = Fd::tcp_connect(None, remote)?;
let _ = runtime::task::spawn(do_proxy(down, up, buflen)).await;
}
Ok(())
}