1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
use anyhow::{Context, Result};
use tokio::net::UdpSocket;
use std::net::SocketAddr;
use std::sync::Arc;
use tokio::sync::RwLock;
use std::collections::HashMap;
use crate::config::PortMapping;
/// UDP connection statistics
#[derive(Debug, Default, Clone)]
pub struct UdpStats {
pub packets_sent: u64,
pub packets_received: u64,
pub bytes_sent: u64,
pub bytes_received: u64,
}
/// UDP forwarder
pub struct UdpForwarder {
stats: Arc<RwLock<UdpStats>>,
}
impl UdpForwarder {
pub fn new() -> Self {
Self {
stats: Arc::new(RwLock::new(UdpStats::default())),
}
}
/// Start forwarding for a port mapping
pub async fn start_forwarding(&self, mapping: &PortMapping) -> Result<()> {
let listen_addr = format!("0.0.0.0:{}", mapping.remote_port);
let socket = Arc::new(UdpSocket::bind(&listen_addr).await
.context(format!("Failed to bind UDP port {}", mapping.remote_port))?);
tracing::info!(
"UDP forwarder listening on :{} -> localhost:{}",
mapping.remote_port,
mapping.local_port
);
let local_addr = format!("127.0.0.1:{}", mapping.local_port);
let stats = self.stats.clone();
// Map of client addresses to their local socket
let clients: Arc<RwLock<HashMap<SocketAddr, Arc<UdpSocket>>>> =
Arc::new(RwLock::new(HashMap::new()));
tokio::spawn(async move {
let mut buf = vec![0u8; 65535];
loop {
match socket.recv_from(&mut buf).await {
Ok((len, client_addr)) => {
let data = buf[..len].to_vec();
// Get or create socket for this client
let client_socket = {
let clients_read = clients.read().await;
if let Some(sock) = clients_read.get(&client_addr) {
sock.clone()
} else {
drop(clients_read);
// Create new socket outside of lock
let sock = Arc::new(UdpSocket::bind("0.0.0.0:0").await.unwrap());
// Add to map
{
let mut clients_write = clients.write().await;
// Check again in case another task added it
if let Some(existing) = clients_write.get(&client_addr) {
existing.clone()
} else {
clients_write.insert(client_addr, sock.clone());
// Start receive loop for this client
let remote_socket = socket.clone();
let local_socket = sock.clone();
let stats = stats.clone();
tokio::spawn(async move {
let mut buf = vec![0u8; 65535];
loop {
match local_socket.recv_from(&mut buf).await {
Ok((len, _)) => {
if remote_socket.send_to(&buf[..len], client_addr).await.is_err() {
break;
}
let mut s = stats.write().await;
s.packets_sent += 1;
s.bytes_sent += len as u64;
}
Err(_) => break,
}
}
});
sock
}
}
}
};
// Forward to local service
if client_socket.send_to(&data, &local_addr).await.is_ok() {
let mut s = stats.write().await;
s.packets_received += 1;
s.bytes_received += len as u64;
}
}
Err(e) => {
tracing::error!("UDP recv error: {}", e);
}
}
}
});
Ok(())
}
/// Get current statistics
pub async fn stats(&self) -> UdpStats {
self.stats.read().await.clone()
}
}
impl Default for UdpForwarder {
fn default() -> Self {
Self::new()
}
}