Apply TcpStream fix to UdpSocket as well

This commit is contained in:
Aram 🍐 2022-01-10 00:35:14 -05:00
parent 2b15e581f2
commit 782f5e74bf

View file

@ -59,17 +59,33 @@ pub async fn udp_proxy_server(
} }
} }
event = endpoint.recv() => { event = endpoint.recv() => {
if let Event::RemoteData(port, data) = event { if let Event::RemoteData(virtual_port, data) = event {
if let Some(peer) = port_pool.get_peer_addr(port).await { if let Some(peer) = port_pool.get_peer_addr(virtual_port).await {
if let Err(e) = socket.send_to(&data, peer).await { // Have remote data to send to the local client
error!( if let Err(e) = socket.writable().await {
"[{}] Failed to send UDP datagram to real client ({}): {:?}", error!("[{}] Failed to check if writable: {:?}", virtual_port, e);
port,
peer,
e,
);
} }
port_pool.update_last_transmit(port).await; let expected = data.len();
let mut sent = 0;
loop {
if sent >= expected {
break;
}
match socket.send_to(&data[sent..expected], peer).await {
Ok(written) => {
debug!("[{}] Sent {} (expected {}) bytes to local client", virtual_port, written, expected);
sent += written;
if sent < expected {
debug!("[{}] Will try to resend remaining {} bytes to local client", virtual_port, (expected - written));
}
},
Err(e) => {
error!("[{}] Failed to send {} bytes to local client: {:?}", virtual_port, expected, e);
break;
}
}
}
port_pool.update_last_transmit(virtual_port).await;
} }
} }
} }