feat: work in process

This commit is contained in:
2022-08-21 11:06:11 +08:00
parent 21c857bdf4
commit 4723298aec
7 changed files with 269 additions and 19 deletions

View File

@@ -6,12 +6,12 @@ use s2n_quic::client::Connect;
use crate::config::ListenConfig;
async fn run(listen_config: &ListenConfig) -> XResult<()> {
pub async fn run(listen_config: &ListenConfig) -> XResult<()> {
let cert_pem = opt_result!(fs::read_to_string(&listen_config.cert_pem_file),
"Read cert pem file: {}, failed: {}", &listen_config.key_pem_file);
"Read cert pem file: {}, failed: {}", &listen_config.cert_pem_file);
let client = Client::builder()
.with_tls(cert_pem.as_str())?
.with_io("0.0.0.0:0")?
.with_io(listen_config.listen.as_str())?
.start()?;
let addr: SocketAddr = "127.0.0.1:4433".parse()?;
let connect = Connect::new(addr).with_server_name("localhost");

View File

@@ -1,5 +1,12 @@
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct ListenConfig {
// client: 0.0.0.0:0
// server: 127.0.0.1:4433
pub listen: String,
pub key_pem_file: String,
pub cert_pem_file: String,
pub key_pem_file: Option<String>,
pub proxy_address: String,
pub proxy_server_name: Option<String>,
}

163
src/io_util.rs Normal file
View File

@@ -0,0 +1,163 @@
use std::io::{Error, ErrorKind};
use std::time::Duration;
use futures::future::try_join;
use s2n_quic::stream::BidirectionalStream;
use tokio::{select, time};
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
use tokio::net::TcpStream;
enum StreamDirection {
Up,
Down,
}
pub async fn transfer_for_server_to_remote(inbound: BidirectionalStream, proxy_addr: String, conn_count: String) -> Result<(), String> {
let mut outbound = match TcpStream::connect(&proxy_addr).await {
Ok(outbound) => outbound,
Err(e) => {
return Err(format!("[conn {}] Failed to connect to: {}, err: {}", &conn_count, &proxy_addr, e));
}
};
if let (Ok(ref in_peer_addr), Ok(ref in_local_addr), Ok(ref out_local_addr), Ok(ref out_peer_addr))
= (inbound.connection().remote_addr(), inbound.connection().local_addr(), outbound.local_addr(), outbound.peer_addr()) {
let peer = format!("{} -> [{} * {}] -> {}", in_peer_addr, in_local_addr, out_local_addr, out_peer_addr);
information!("[conn {}] New server-remote tcp connection: {}", &conn_count, peer);
}
let (mut ri, mut wi) = inbound.split();
let (mut ro, mut wo) = outbound.split();
inner_transfer(&mut ri, &mut wi, &mut ro, &mut wo, conn_count).await
}
pub async fn transfer_for_client_to_server(mut inbound: TcpStream, outbound: BidirectionalStream, conn_count: String) -> Result<(), String> {
if let (Ok(ref in_peer_addr), Ok(ref in_local_addr), Ok(ref out_local_addr), Ok(ref out_peer_addr))
= (inbound.peer_addr(), inbound.local_addr(), outbound.connection().local_addr(), outbound.connection().remote_addr()) {
let peer = format!("{} -> [{} * {}] -> {}", in_peer_addr, in_local_addr, out_local_addr, out_peer_addr);
information!("[conn {}] New client-server tcp connection: {}", &conn_count, peer);
}
let (mut ri, mut wi) = inbound.split();
let (mut ro, mut wo) = outbound.split();
inner_transfer(&mut ri, &mut wi, &mut ro, &mut wo, conn_count).await
}
async fn inner_transfer<'a, R1, W1, R2, W2>(mut ri: &'a mut R1, mut wi: &'a mut W1, mut ro: &'a mut R2, mut wo: &'a mut W2, conn_count: String) -> Result<(), String>
where R1: AsyncRead + Unpin + ?Sized, W1: AsyncWrite + Unpin + ?Sized,
R2: AsyncRead + Unpin + ?Sized, W2: AsyncWrite + Unpin + ?Sized {
// IO copy timeout 6 HOURS
let tcp_io_copy_timeout = Duration::from_secs(6 * 3600);
let shutdown_tcp_timeout = Duration::from_secs(60);
let (client_to_server_tx, client_to_server_rx) = tokio::sync::oneshot::channel::<bool>();
let (server_to_client_tx, server_to_client_rx) = tokio::sync::oneshot::channel::<bool>();
let client_to_server = async {
let r = select! {
_timeout = time::sleep(tcp_io_copy_timeout) => {
failure!("[conn {}] TCP copy client -> server timeout", &conn_count);
Err(Error::new(ErrorKind::TimedOut, "timeout"))
}
_tcp_break = client_to_server_rx => {
failure!("[conn {}] TCP copy client -> server shutdown", &conn_count);
Err(Error::new(ErrorKind::BrokenPipe, "shutdown"))
}
data_copy_result = copy_data(&mut ri, &mut wo, StreamDirection::Up, &conn_count) => {
match data_copy_result {
Err(e) => {
failure!("[conn {}] TCP copy client -> server error: {}", &conn_count, e);
Err(e)
},
Ok(r) => {
information!("[conn {}] TCP copy client -> server success: {} byte(s)", &conn_count, r);
Ok(r)
},
}
}
};
information!("[conn {}] Close client to server connection", &conn_count);
match time::timeout(shutdown_tcp_timeout, wo.shutdown()).await {
Err(e) => warning!("[conn {}] TCP close client -> server timeout: {}", &conn_count, e),
Ok(Err(e)) => warning!("[conn {}] TCP close client -> server error: {}", &conn_count, e),
_ => {}
}
time::sleep(Duration::from_secs(2)).await;
let _ = server_to_client_tx.send(true);
r
};
let server_to_client = async {
let r = select! {
_timeout = time::sleep(tcp_io_copy_timeout) => {
failure!("[conn {}] TCP copy server -> client timeout", &conn_count);
Err(Error::new(ErrorKind::TimedOut, "timeout"))
}
_tcp_break = server_to_client_rx => {
failure!("[conn {}] TCP copy server -> client shutdown", &conn_count);
Err(Error::new(ErrorKind::BrokenPipe, "shutdown"))
}
data_copy_result = copy_data(&mut ro, &mut wi, StreamDirection::Down, &conn_count) => {
match data_copy_result {
Err(e) => {
failure!("[conn {}] TCP copy server -> client error: {}", &conn_count, e);
Err(e)
},
Ok(r) => {
information!("[conn {}] TCP copy server -> client success: {} byte(s)", &conn_count, r);
Ok(r)
},
}
}
};
information!("[conn {}] Close server to client connection", &conn_count);
match time::timeout(shutdown_tcp_timeout, wi.shutdown()).await {
Err(e) => warning!("[conn {}] TCP close server -> client timeout: {}", &conn_count, e),
Ok(Err(e)) => warning!("[conn {}] TCP close server -> client error: {}", &conn_count, e),
_ => {}
}
time::sleep(Duration::from_secs(2)).await;
let _ = client_to_server_tx.send(true);
r
};
let r = match try_join(client_to_server, server_to_client).await {
Err(e) => Err(format!("[conn {}] Failed try_join: {}", &conn_count, e)),
Ok((upstream_bytes, downstream_bytes)) => {
information!("[conn {}] Finished, proxy-in: {} bytes, proxy-out: {} bytes",
&conn_count, upstream_bytes, downstream_bytes
);
Ok(())
}
};
r
}
// fn is_in_peer_addr_matches(inbound: &TcpStream, allow_ips: &[IpAddressMask]) -> bool {
// if let Ok(ref in_peer_addr) = inbound.peer_addr() {
// if let Some(ip_filter) = &*GLOBAL_TEMP_IP_FILTER.read().unwrap() {
// if let Some(ip) = ip_filter.is_ip_address_matches(in_peer_addr) {
// return true;
// }
// }
// if !allow_ips.is_empty() {
// // ONLY if allow ips is not config, returns true
// // If default deny IPs, should config like: "allow_ips": ["127.0.0.1"]
// return allow_ips.iter().any(|ip| ip.is_matches(in_peer_addr));
// }
// }
// true
// }
async fn copy_data<'a, R, W>(
reader: &'a mut R, writer: &'a mut W, _direction: StreamDirection, _conn_count: &str) -> tokio::io::Result<u64>
where R: AsyncRead + Unpin + ?Sized,
W: AsyncWrite + Unpin + ?Sized {
let mut total_copied_bytes = 0_u64;
let mut buff = vec![0; 1024 * 4];
loop {
match reader.read(&mut buff).await {
Ok(0) => return Ok(total_copied_bytes),
Ok(n) => match writer.write_all(&buff[..n]).await {
Ok(_) => {
total_copied_bytes += n as u64;
}
Err(e) => return Err(e),
}
Err(e) => return Err(e),
}
}
}

View File

@@ -2,11 +2,13 @@
extern crate rust_util;
mod config;
mod io_util;
mod client;
mod server;
use clap::{App, Arg};
use rust_util::XResult;
use crate::config::ListenConfig;
#[tokio::main]
async fn main() -> XResult<()> {
@@ -35,8 +37,11 @@ async fn main() -> XResult<()> {
failure_and_exit!("Must run in server on client mode")
}
let listen_config: ListenConfig = deser_hjson::from_str("").unwrap();
// TODO
Ok(())
if client_mode {
client::run(&listen_config).await
} else {
server::run(&listen_config).await
}
}

View File

@@ -4,12 +4,15 @@ use rust_util::XResult;
use s2n_quic::Server;
use crate::config::ListenConfig;
use crate::io_util;
async fn run(listen_config: &ListenConfig) -> XResult<()> {
pub async fn run(listen_config: &ListenConfig) -> XResult<()> {
let cert_pem = opt_result!(fs::read_to_string(&listen_config.cert_pem_file),
"Read cert pem file: {}, failed: {}", &listen_config.key_pem_file);
let key_pem = opt_result!(fs::read_to_string(&listen_config.key_pem_file),
"Read key pem file :{}, failed: {}", &listen_config.key_pem_file);
"Read cert pem file: {}, failed: {}", &listen_config.cert_pem_file);
let key_pem_file = opt_value_result!(&listen_config.key_pem_file,
"key_pem_file in config is required in server mode");
let key_pem = opt_result!(fs::read_to_string(&key_pem_file),
"Read key pem file :{}, failed: {}", &key_pem_file);
let mut server = Server::builder()
.with_tls((cert_pem.as_str(), key_pem.as_str()))?
@@ -17,18 +20,18 @@ async fn run(listen_config: &ListenConfig) -> XResult<()> {
.start()?;
while let Some(mut connection) = server.accept().await {
// spawn a new task for the connection
let proxy_address = listen_config.proxy_address.clone();
tokio::spawn(async move {
information!("Connection accepted from {:?}", connection.remote_addr());
while let Ok(Some(mut stream)) = connection.accept_bidirectional_stream().await {
while let Ok(Some(stream)) = connection.accept_bidirectional_stream().await {
// spawn a new task for the stream
let proxy_address = proxy_address.clone();
tokio::spawn(async move {
information!("Stream opened from {:?}", stream.connection().remote_addr());
// echo any data back to the stream
// while let Ok(Some(data)) = stream.receive().await {
// stream.send(data).await.expect("stream should be open");
// }
// println!("Stream closed from {:?}", stream.connection().remote_addr());
// TODO ...
let conn_count = format!("{}", 1);
if let Err(e) = io_util::transfer_for_server_to_remote(stream, proxy_address, conn_count).await {
failure!("Server - Client error: {}", e);
}
});
}
println!("Connection closed from {:?}", connection.remote_addr());
@@ -36,4 +39,4 @@ async fn run(listen_config: &ListenConfig) -> XResult<()> {
}
Ok(())
}
}