featherbit/stream/
udp.rs

1//! UDP stream proxy: relay datagrams between clients and a load-balanced
2//! upstream, tracking one session per client address.
3//!
4//! UDP has no connections, so a single owning task demultiplexes the shared
5//! listener socket by client address into a session map (no lock needed — the
6//! map is touched only by that task). Each session owns an ephemeral socket
7//! `connect`ed to the chosen upstream and a reader task that pumps upstream
8//! replies back to the client. A session is torn down after `idle` with **no
9//! traffic in either direction** (tracked by a shared `last_seen`), reaped by a
10//! 1-second prune tick.
11
12use std::collections::HashMap;
13use std::io;
14use std::net::SocketAddr;
15use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
16use std::sync::Arc;
17use std::time::{Duration, Instant};
18
19use tokio::net::UdpSocket;
20use tokio::sync::watch;
21use tokio::time::timeout;
22use tracing::{debug, warn};
23
24use crate::balancer::{Balancer, OwnedConnGuard};
25use crate::config::StreamListenerConfig;
26
27/// Max UDP payload plus headroom.
28const BUF_SIZE: usize = 65_535;
29
30/// One client's proxy session: the ephemeral upstream socket, shared activity
31/// clock, liveness flag, and the in-flight guard held for the session's life.
32struct Session {
33    upstream: Arc<UdpSocket>,
34    last_seen: Arc<AtomicU64>,
35    alive: Arc<AtomicBool>,
36    _guard: OwnedConnGuard,
37}
38
39fn now_ms(epoch: Instant) -> u64 {
40    epoch.elapsed().as_millis() as u64
41}
42
43/// Binds the UDP listener (fail-fast) and spawns its receive loop, returning the
44/// bound address (the OS-assigned port when `cfg.port == 0`).
45pub async fn spawn(
46    cfg: &StreamListenerConfig,
47    balancer: Arc<Balancer>,
48    idle: Duration,
49    mut shutdown_rx: watch::Receiver<bool>,
50) -> io::Result<SocketAddr> {
51    let ip = cfg
52        .bind
53        .parse()
54        .map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, format!("invalid bind: {}", e)))?;
55    let addr = SocketAddr::new(ip, cfg.port);
56    let sock = Arc::new(UdpSocket::bind(addr).await?);
57    let local = sock.local_addr()?;
58
59    tokio::spawn(async move {
60        let epoch = Instant::now();
61        let mut sessions: HashMap<SocketAddr, Session> = HashMap::new();
62        let mut buf = [0u8; BUF_SIZE];
63        let mut prune = tokio::time::interval(Duration::from_secs(1));
64
65        loop {
66            tokio::select! {
67                _ = shutdown_rx.changed() => break,
68                _ = prune.tick() => {
69                    sessions.retain(|_, s| s.alive.load(Ordering::Relaxed));
70                }
71                recv = sock.recv_from(&mut buf) => {
72                    let (n, client) = match recv {
73                        Ok(v) => v,
74                        Err(e) => { warn!("udp stream recv error: {}", e); continue; }
75                    };
76                    let now = now_ms(epoch);
77
78                    // Drop a dead session so it gets recreated below.
79                    if sessions.get(&client).is_some_and(|s| !s.alive.load(Ordering::Relaxed)) {
80                        sessions.remove(&client);
81                    }
82                    if let std::collections::hash_map::Entry::Vacant(e) = sessions.entry(client) {
83                        match create_session(&balancer, &sock, client, idle, epoch).await {
84                            Ok(session) => { e.insert(session); }
85                            Err(e) => { warn!("udp stream session setup failed: {}", e); continue; }
86                        }
87                    }
88                    if let Some(session) = sessions.get(&client) {
89                        session.last_seen.store(now, Ordering::Relaxed);
90                        if let Err(e) = session.upstream.send(&buf[..n]).await {
91                            debug!("udp stream send to upstream failed: {}", e);
92                        }
93                    }
94                }
95            }
96        }
97    });
98
99    Ok(local)
100}
101
102/// Creates a session for `client`: an ephemeral socket connected to a selected
103/// upstream, plus a spawned reader task pumping replies back to the client.
104async fn create_session(
105    balancer: &Arc<Balancer>,
106    listener: &Arc<UdpSocket>,
107    client: SocketAddr,
108    idle: Duration,
109    epoch: Instant,
110) -> io::Result<Session> {
111    let idx = balancer.select(&client.to_string());
112    let guard = balancer.owned_acquire(idx);
113    let target = balancer.target(idx);
114
115    let upstream = Arc::new(UdpSocket::bind(("0.0.0.0", 0)).await?);
116    upstream
117        .connect((target.host.as_str(), target.port))
118        .await?;
119
120    let last_seen = Arc::new(AtomicU64::new(now_ms(epoch)));
121    let alive = Arc::new(AtomicBool::new(true));
122
123    tokio::spawn(reader(
124        listener.clone(),
125        upstream.clone(),
126        client,
127        alive.clone(),
128        last_seen.clone(),
129        idle,
130        epoch,
131    ));
132
133    Ok(Session {
134        upstream,
135        last_seen,
136        alive,
137        _guard: guard,
138    })
139}
140
141/// Pumps upstream replies back to the client, exiting when the session has been
142/// idle in **both** directions for `idle` (so an active client with a briefly
143/// quiet upstream is not torn down).
144async fn reader(
145    listener: Arc<UdpSocket>,
146    upstream: Arc<UdpSocket>,
147    client: SocketAddr,
148    alive: Arc<AtomicBool>,
149    last_seen: Arc<AtomicU64>,
150    idle: Duration,
151    epoch: Instant,
152) {
153    let mut buf = [0u8; BUF_SIZE];
154    let idle_ms = idle.as_millis() as u64;
155    loop {
156        match timeout(idle, upstream.recv(&mut buf)).await {
157            Ok(Ok(n)) => {
158                if let Err(e) = listener.send_to(&buf[..n], client).await {
159                    debug!("udp stream send to client failed: {}", e);
160                    break;
161                }
162                last_seen.store(now_ms(epoch), Ordering::Relaxed);
163            }
164            Ok(Err(e)) => {
165                debug!("udp stream upstream recv error: {}", e);
166                break;
167            }
168            Err(_) => {
169                // recv idle elapsed — only give up if the client side is idle too.
170                if now_ms(epoch).saturating_sub(last_seen.load(Ordering::Relaxed)) >= idle_ms {
171                    break;
172                }
173            }
174        }
175    }
176    alive.store(false, Ordering::Relaxed);
177}
featherbit/stream/udp.rs

featherbit/stream/
udp.rs
