rws_proxy/tcp_proxy/
mod.rs

1//! Layer-4 TCP proxy.
2//!
3//! [`TcpProxy`] accepts raw TCP connections and forwards them to backend servers,
4//! bidirectionally tunneling bytes with one thread per direction.
5//!
6//! Unlike [`crate::proxy::ReverseProxy`] (which operates at the HTTP layer),
7//! `TcpProxy` is protocol-agnostic: any TCP-based protocol (database wire formats,
8//! custom binary protocols, raw TLS passthrough) is forwarded unchanged.
9//!
10//! # Example
11//!
12//! ```rust,no_run
13//! # extern crate rws_proxy as rust_web_server;
14//! use rust_web_server::tcp_proxy::TcpProxy;
15//!
16//! // Proxy raw TCP on port 5432 across two PostgreSQL backends.
17//! TcpProxy::new(["backend-1:5432", "backend-2:5432"])
18//!     .connect_timeout_ms(3000)
19//!     .bind("0.0.0.0:5432")
20//!     .unwrap();
21//! ```
22
23use std::io;
24use std::net::{TcpListener, TcpStream, ToSocketAddrs};
25use std::sync::{
26    Arc,
27    atomic::{AtomicUsize, Ordering},
28};
29use std::time::Duration;
30
31/// Layer-4 (raw TCP) reverse proxy with round-robin load balancing.
32///
33/// Call [`TcpProxy::bind`] to start accepting connections. Each connection is
34/// handled in its own thread pair (one thread per direction), so `bind` blocks
35/// the calling thread indefinitely.
36pub struct TcpProxy {
37    backends: Vec<String>,
38    counter: Arc<AtomicUsize>,
39    connect_timeout: Duration,
40}
41
42impl TcpProxy {
43    /// Create a proxy that distributes connections across `backends` in
44    /// round-robin order. Each entry must be `"host:port"`.
45    pub fn new<I, S>(backends: I) -> Self
46    where
47        I: IntoIterator<Item = S>,
48        S: Into<String>,
49    {
50        TcpProxy {
51            backends: backends.into_iter().map(|b| b.into()).collect(),
52            counter: Arc::new(AtomicUsize::new(0)),
53            connect_timeout: Duration::from_secs(5),
54        }
55    }
56
57    /// Override the TCP connect timeout to each backend (default: 5 s).
58    pub fn connect_timeout_ms(mut self, ms: u64) -> Self {
59        self.connect_timeout = Duration::from_millis(ms);
60        self
61    }
62
63    /// Bind on `addr` and start proxying. Blocks until the listener is closed.
64    pub fn bind(self, addr: &str) -> Result<(), String> {
65        if self.backends.is_empty() {
66            return Err("TcpProxy: no backends configured".to_string());
67        }
68        let listener = TcpListener::bind(addr)
69            .map_err(|e| format!("TcpProxy: bind on {} failed: {}", addr, e))?;
70        println!("TcpProxy: listening on {}", addr);
71        let proxy = Arc::new(self);
72        for incoming in listener.incoming() {
73            let client = match incoming {
74                Ok(s) => s,
75                Err(e) => {
76                    eprintln!("TcpProxy: accept error: {}", e);
77                    continue;
78                }
79            };
80            let p = Arc::clone(&proxy);
81            std::thread::spawn(move || {
82                if let Err(e) = p.relay(client) {
83                    eprintln!("TcpProxy: relay error: {}", e);
84                }
85            });
86        }
87        Ok(())
88    }
89
90    fn pick_backend(&self) -> &str {
91        let i = self.counter.fetch_add(1, Ordering::Relaxed) % self.backends.len();
92        &self.backends[i]
93    }
94
95    fn relay(&self, client: TcpStream) -> Result<(), String> {
96        let addr_str = self.pick_backend().to_string();
97        let sock_addr = addr_str
98            .to_socket_addrs()
99            .map_err(|e| format!("DNS lookup for {} failed: {}", addr_str, e))?
100            .next()
101            .ok_or_else(|| format!("no address resolved for {}", addr_str))?;
102
103        let backend = TcpStream::connect_timeout(&sock_addr, self.connect_timeout)
104            .map_err(|e| format!("TcpProxy: connect to {} failed: {}", addr_str, e))?;
105
106        let mut client_r = client.try_clone().map_err(|e| e.to_string())?;
107        let mut backend_r = backend.try_clone().map_err(|e| e.to_string())?;
108        let mut client_w = client;
109        let mut backend_w = backend;
110
111        let t1 = std::thread::spawn(move || {
112            io::copy(&mut client_r, &mut backend_w).ok();
113            let _ = backend_w.shutdown(std::net::Shutdown::Write);
114        });
115        let t2 = std::thread::spawn(move || {
116            io::copy(&mut backend_r, &mut client_w).ok();
117            let _ = client_w.shutdown(std::net::Shutdown::Write);
118        });
119
120        let _ = t1.join();
121        let _ = t2.join();
122        Ok(())
123    }
124}
rws_proxy/tcp_proxy/mod.rs

rws_proxy/tcp_proxy/
mod.rs
