rws_proxy/udp_proxy/
mod.rs

1//! Layer-4 UDP proxy.
2//!
3//! [`UdpProxy`] receives UDP datagrams from clients and forwards each one to a
4//! backend server, then relays the backend's reply to the original sender.
5//! This request-reply model covers protocols such as DNS, NTP, and RADIUS.
6//!
7//! Each datagram is handled in its own thread, so the main `bind` loop is
8//! never blocked waiting for a backend reply.
9//!
10//! # Example
11//!
12//! ```rust,no_run
13//! # extern crate rws_proxy as rust_web_server;
14//! use rust_web_server::udp_proxy::UdpProxy;
15//!
16//! // Forward DNS queries round-robin across two resolvers.
17//! UdpProxy::new(["8.8.8.8:53", "8.8.4.4:53"])
18//!     .reply_timeout_ms(2000)
19//!     .bind("0.0.0.0:53")
20//!     .unwrap();
21//! ```
22
23use std::net::{ToSocketAddrs, UdpSocket};
24use std::sync::{
25    Arc,
26    atomic::{AtomicUsize, Ordering},
27};
28use std::time::Duration;
29
30/// Layer-4 (raw UDP) reverse proxy with round-robin load balancing.
31///
32/// Each client datagram is forwarded to one backend; the backend's reply is
33/// delivered back to the originating client address. No session state is kept
34/// between datagrams.
35///
36/// Call [`UdpProxy::bind`] to start. It blocks the calling thread indefinitely.
37pub struct UdpProxy {
38    backends: Vec<String>,
39    counter: Arc<AtomicUsize>,
40    reply_timeout: Duration,
41    buffer_size: usize,
42}
43
44impl UdpProxy {
45    /// Create a proxy that distributes datagrams across `backends` in
46    /// round-robin order. Each entry must be `"host:port"`.
47    pub fn new<I, S>(backends: I) -> Self
48    where
49        I: IntoIterator<Item = S>,
50        S: Into<String>,
51    {
52        UdpProxy {
53            backends: backends.into_iter().map(|b| b.into()).collect(),
54            counter: Arc::new(AtomicUsize::new(0)),
55            reply_timeout: Duration::from_secs(5),
56            buffer_size: 65536,
57        }
58    }
59
60    /// Override the timeout waiting for a backend reply (default: 5 s).
61    pub fn reply_timeout_ms(mut self, ms: u64) -> Self {
62        self.reply_timeout = Duration::from_millis(ms);
63        self
64    }
65
66    /// Override the per-datagram buffer size (default: 65 536 B).
67    pub fn buffer_size(mut self, bytes: usize) -> Self {
68        self.buffer_size = bytes;
69        self
70    }
71
72    /// Bind on `addr` and start forwarding datagrams. Blocks indefinitely.
73    pub fn bind(self, addr: &str) -> Result<(), String> {
74        if self.backends.is_empty() {
75            return Err("UdpProxy: no backends configured".to_string());
76        }
77        let socket = UdpSocket::bind(addr)
78            .map_err(|e| format!("UdpProxy: bind on {} failed: {}", addr, e))?;
79        println!("UdpProxy: listening on {}", addr);
80        let proxy = Arc::new(self);
81
82        loop {
83            let mut buf = vec![0u8; proxy.buffer_size];
84            let (n, client_addr) = match socket.recv_from(&mut buf) {
85                Ok(v) => v,
86                Err(e) => {
87                    eprintln!("UdpProxy: recv_from error: {}", e);
88                    continue;
89                }
90            };
91            let packet = buf[..n].to_vec();
92            let backend_addr = proxy.pick_backend().to_string();
93            let reply_socket = match socket.try_clone() {
94                Ok(s) => s,
95                Err(e) => {
96                    eprintln!("UdpProxy: socket clone error: {}", e);
97                    continue;
98                }
99            };
100            let timeout = proxy.reply_timeout;
101            let buf_size = proxy.buffer_size;
102
103            std::thread::spawn(move || {
104                let backend_sock_addr = match backend_addr.to_socket_addrs() {
105                    Ok(mut a) => match a.next() {
106                        Some(addr) => addr,
107                        None => {
108                            eprintln!("UdpProxy: no address for {}", backend_addr);
109                            return;
110                        }
111                    },
112                    Err(e) => {
113                        eprintln!("UdpProxy: DNS lookup for {} failed: {}", backend_addr, e);
114                        return;
115                    }
116                };
117
118                let backend = match UdpSocket::bind("0.0.0.0:0") {
119                    Ok(s) => s,
120                    Err(e) => {
121                        eprintln!("UdpProxy: ephemeral socket error: {}", e);
122                        return;
123                    }
124                };
125                let _ = backend.set_read_timeout(Some(timeout));
126
127                if let Err(e) = backend.send_to(&packet, backend_sock_addr) {
128                    eprintln!("UdpProxy: send to {} failed: {}", backend_addr, e);
129                    return;
130                }
131
132                let mut reply = vec![0u8; buf_size];
133                match backend.recv_from(&mut reply) {
134                    Ok((m, _)) => {
135                        let _ = reply_socket.send_to(&reply[..m], client_addr);
136                    }
137                    Err(e) if e.kind() != std::io::ErrorKind::WouldBlock
138                           && e.kind() != std::io::ErrorKind::TimedOut => {
139                        eprintln!("UdpProxy: backend reply error from {}: {}", backend_addr, e);
140                    }
141                    _ => {} // timeout — backend didn't reply in time, drop silently
142                }
143            });
144        }
145    }
146
147    fn pick_backend(&self) -> &str {
148        let i = self.counter.fetch_add(1, Ordering::Relaxed) % self.backends.len();
149        &self.backends[i]
150    }
151}
rws_proxy/udp_proxy/mod.rs

rws_proxy/udp_proxy/
mod.rs
