use std::sync::Arc;
use anyhow::Context;
use tokio::io::{
AsyncReadExt,
AsyncWriteExt,
};
use tokio::net::UdpSocket as TokioUdpSocket;
use tokio::sync::Mutex;
use tokio_util::sync::CancellationToken;
use tracing::{
error,
info,
};
const BUFFER_SIZE: usize = 131072;
pub struct UdpForwarder;
impl UdpForwarder {
pub async fn bind_and_forward(
local_address: String, local_port: u16,
upstream_conn: impl tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send + 'static,
cancellation_token: CancellationToken,
) -> anyhow::Result<(u16, tokio::task::JoinHandle<()>)> {
let local_udp_addr = format!("{local_address}:{local_port}");
let local_udp_socket = Arc::new(
TokioUdpSocket::bind(&local_udp_addr)
.await
.context("Failed to bind local UDP socket")?,
);
let local_port = local_udp_socket.local_addr()?.port();
info!("Local UDP socket bound to {}", local_udp_addr);
let (tcp_read, tcp_write) = tokio::io::split(upstream_conn);
let tcp_read = Arc::new(Mutex::new(tcp_read));
let tcp_write = Arc::new(Mutex::new(tcp_write));
let local_udp_socket_read = local_udp_socket.clone();
let local_udp_socket_write = local_udp_socket;
let handle = tokio::spawn({
let tcp_read = tcp_read.clone();
let tcp_write = tcp_write.clone();
let cancel_token = cancellation_token.clone();
async move {
let mut udp_buffer = [0u8; BUFFER_SIZE];
let peer: Arc<Mutex<Option<std::net::SocketAddr>>> = Arc::new(Mutex::new(None));
loop {
tokio::select! {
result = local_udp_socket_read.recv_from(&mut udp_buffer) => {
match result {
Ok((len, src)) => {
*peer.lock().await = Some(src);
let mut writer = tcp_write.lock().await;
let packet_len = (len as u32).to_be_bytes();
if let Err(e) = writer.write_all(&packet_len).await {
error!("Failed to write packet length to TCP stream: {:?}", e);
break;
}
if let Err(e) = writer.write_all(&udp_buffer[..len]).await {
error!("Failed to write UDP packet to TCP stream: {:?}", e);
break;
}
if let Err(e) = writer.flush().await {
error!("Failed to flush TCP stream: {:?}", e);
break;
}
},
Err(e) => {
error!("Failed to receive from UDP socket: {:?}", e);
break;
}
}
},
result = async {
let mut reader = tcp_read.lock().await;
Self::read_tcp_length_and_packet(&mut *reader).await
} => {
match result {
Ok(Some(packet)) => {
let peer_addr = *peer.lock().await;
if let Some(peer_addr) = peer_addr {
if let Err(e) = local_udp_socket_write.send_to(&packet, &peer_addr).await {
error!("Failed to send UDP packet to peer: {:?}", e);
break;
}
} else {
error!("No UDP peer to send to");
break;
}
},
Ok(None) => break,
Err(e) => {
error!("Failed to read from TCP stream: {:?}", e);
break;
}
}
}
_ = cancel_token.cancelled() => {
info!("UDP forwarder cancelled, shutting down");
break;
}
}
}
if let Err(e) = tcp_write.lock().await.shutdown().await {
error!("Error shutting down TCP writer: {:?}", e);
}
}
});
Ok((local_port, handle))
}
async fn read_tcp_length_and_packet(
tcp_read: &mut (impl AsyncReadExt + Unpin),
) -> anyhow::Result<Option<Vec<u8>>> {
let mut len_bytes = [0u8; 4];
if tcp_read.read_exact(&mut len_bytes).await.is_err() {
return Ok(None);
}
let len = u32::from_be_bytes(len_bytes) as usize;
let mut packet = vec![0u8; len];
if tcp_read.read_exact(&mut packet).await.is_err() {
return Ok(None);
}
Ok(Some(packet))
}
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use tokio::io::duplex;
use tokio::net::UdpSocket;
use tokio_util::sync::CancellationToken;
use super::*;
#[tokio::test]
async fn test_read_tcp_length_and_packet() {
let (mut reader, mut writer) = duplex(1024);
let packet = b"hello world";
let len = packet.len() as u32;
let len_bytes = len.to_be_bytes();
writer.write_all(&len_bytes).await.unwrap();
writer.write_all(packet).await.unwrap();
writer.flush().await.unwrap();
let result = UdpForwarder::read_tcp_length_and_packet(&mut reader).await;
assert!(result.is_ok());
let data = result.unwrap();
assert!(data.is_some());
assert_eq!(data.unwrap(), packet);
}
#[tokio::test]
async fn test_read_tcp_length_and_packet_empty() {
let (mut reader, _) = duplex(0);
let result = UdpForwarder::read_tcp_length_and_packet(&mut reader).await;
assert!(result.is_ok());
let data = result.unwrap();
assert!(data.is_none());
}
#[tokio::test]
async fn test_bind_and_forward_basic() {
let (_, server_stream) = duplex(1024);
let cancellation_token = CancellationToken::new();
let result = UdpForwarder::bind_and_forward(
"127.0.0.1".to_string(),
0,
server_stream,
cancellation_token,
)
.await;
assert!(result.is_ok());
let (port, handle) = result.unwrap();
assert!(port > 0);
let client_socket = UdpSocket::bind("127.0.0.1:0").await.unwrap();
client_socket
.connect(format!("127.0.0.1:{port}"))
.await
.unwrap();
client_socket.send(b"hello").await.unwrap();
tokio::time::sleep(Duration::from_millis(100)).await;
handle.abort();
}
#[tokio::test]
async fn test_read_tcp_length_and_packet_partial() {
let (mut reader, mut writer) = duplex(1024);
let packet = b"hello world";
let len = packet.len() as u32;
let len_bytes = len.to_be_bytes();
writer.write_all(&len_bytes).await.unwrap();
writer.write_all(&packet[0..5]).await.unwrap();
writer.flush().await.unwrap();
let read_task =
tokio::spawn(
async move { UdpForwarder::read_tcp_length_and_packet(&mut reader).await },
);
tokio::time::sleep(Duration::from_millis(100)).await;
drop(writer);
let result = tokio::time::timeout(Duration::from_secs(1), read_task).await;
assert!(result.is_ok(), "Test timed out");
}
}