fixing handshake handle logic and stricting the memory orders

This commit is contained in:
Saber Haj Rabiee 2022-12-25 06:47:03 -08:00
parent 1040b3dec9
commit ac07749134
2 changed files with 31 additions and 10 deletions

View File

@ -268,7 +268,9 @@ async fn main() -> io::Result<()> {
}
};
let mut should_receive_handshake_packet = false;
if let Some(ref p) = *handshake_packet {
should_receive_handshake_packet = true;
if tcp_sock.send(p).await.is_none() {
error!(
"Failed to send handshake packet to remote, closing connection."
@ -316,7 +318,13 @@ async fn main() -> io::Result<()> {
debug!("Received EOF from {addr}, closing connection {sock_index}");
break;
}
let udp_sock_index = udp_sock_index.fetch_add(1, Ordering::Relaxed) % udp_socks_amount;
// discard handshake packet since it is not related to
// underlying logic
if should_receive_handshake_packet {
should_receive_handshake_packet = false;
continue;
}
let udp_sock_index = udp_sock_index.fetch_add(1, Ordering::SeqCst) % udp_socks_amount;
let udp_sock = udp_socks[udp_sock_index].clone();
if let Some(ref enc) = *encryption {
enc.decrypt(&mut buf_tcp[..size]);
@ -367,7 +375,7 @@ async fn main() -> io::Result<()> {
debug!("Zero-sized data are not supported, discarding received data from {addr}");
continue;
}
let tcp_sock_index = tcp_sock_index.fetch_add(1, Ordering::Relaxed) % tcp_socks_amount;
let tcp_sock_index = tcp_sock_index.fetch_add(1, Ordering::SeqCst) % tcp_socks_amount;
let tcp_sock = tcp_socks[tcp_sock_index].clone();
if let Some(ref enc) = *encryption {
enc.encrypt(&mut buf_udp[..size]);

View File

@ -181,6 +181,7 @@ async fn main() -> io::Result<()> {
.get_one::<String>("handshake_packet")
.map(fs::read)
.transpose()?;
let handshake_packet = Arc::new(handshake_packet);
let tun = TunBuilder::new()
.name(tun_name) // if name is empty, then it is set by kernel.
@ -207,14 +208,6 @@ async fn main() -> io::Result<()> {
'main_loop: loop {
let tcp_sock = Arc::new(stack.accept().await);
info!("New connection: {}", tcp_sock);
if let Some(ref p) = handshake_packet {
if tcp_sock.send(p).await.is_none() {
error!("Failed to send handshake packet to remote, closing connection.");
continue;
}
debug!("Sent handshake packet to: {}", tcp_sock);
}
let udp_sock = UdpSocket::bind(if remote_addr.is_ipv4() {
"0.0.0.0:0"
@ -301,6 +294,11 @@ async fn main() -> io::Result<()> {
let encryption = encryption.clone();
let packet_received = packet_received.clone();
let cancellation = cancellation.clone();
let handshake_packet = handshake_packet.clone();
let mut should_receive_handshake_packet = false;
if handshake_packet.is_some() {
should_receive_handshake_packet = true;
}
tokio::spawn(async move {
let mut buf_tcp = [0u8; MAX_PACKET_LEN];
let mut udp_sock_index = 0;
@ -319,6 +317,21 @@ async fn main() -> io::Result<()> {
debug!("Received EOF from {local_addr}, closing connection");
break;
}
// discard handshake packet since it is not related to
// underlying logic
if should_receive_handshake_packet {
should_receive_handshake_packet = false;
if let Some(ref p) = *handshake_packet {
if tcp_sock.send(p).await.is_none() {
error!("Failed to send handshake packet to remote, closing connection.");
break;
}
debug!("Sent handshake packet to: {}", tcp_sock);
continue;
}
}
udp_sock_index = (udp_sock_index + 1) % udp_socks_amount;
let udp_sock = udp_socks[udp_sock_index].clone();
if let Some(ref enc) = *encryption {