yprox/
lib.rs

1use std::{
2    net::{SocketAddr, TcpListener},
3    sync::{mpsc, Arc},
4    thread,
5};
6
7use broadcaster::broadcaster;
8use client::client;
9pub use error::{Error, Result};
10use server::server;
11pub use server::Message;
12use target::Target;
13
14mod broadcaster;
15mod cli;
16mod client;
17mod error;
18mod server;
19mod target;
20mod utils;
21
22/// Starts a TCP server that forwards incoming connections to multiple destinations.
23///
24/// # Arguments
25///
26/// * `bind_addr` - The address to bind the server to.
27/// * `targets` - A vector of names and destination addresses to forward incoming connections to.
28///
29/// # Example
30///
31/// ```
32/// #[tokio::main]
33///
34/// async fn main() {
35///     let bind_addr = SocketAddr::parse("127.0.0.1:8080");
36///     let targets = vec![
37///         ("server1".to_string(), SocketAddr::new("127.0.0.1:8081")),
38///         ("server2".to_string(), SocketAddr::new("127.0.0.1:8082"))
39///     ];
40///     start(bind_addr, targets).await;
41/// }
42/// ```
43pub fn start_proxy(bind_addr: SocketAddr, targets: Vec<(String, SocketAddr)>) -> Result<()> {
44    let listener = TcpListener::bind(bind_addr)?;
45
46    // used to send messages to the server
47    let (send_message, receive_message) = mpsc::channel();
48
49    // used to send broadcasts to all targets
50    let (send_broadcast, receive_broadcast) = mpsc::channel();
51
52    // spawn the server thread (handles server -> client and server -> broadcast)
53    // handles messages between client and server, and sends broadcasts
54    thread::spawn(|| server(receive_message, send_broadcast));
55
56    // spawn the broadcasting thread (handles server -> targets and targets -> server)
57    // the breadcaster receives broadcast requests and sends them to all targets
58    // it also receives the send_message handle so that each target can send individual
59    // responses to the server
60    let send_message_clone = send_message.clone();
61    thread::spawn(|| {
62        broadcaster(targets, receive_broadcast, send_message_clone)
63            .map_err(|err| eprintln!("{:?}", err))
64    });
65
66    // spawn the client threads (handle client -> server)
67    for stream in listener.incoming() {
68        match stream {
69            Ok(stream) => {
70                let stream = Arc::new(stream);
71                let send_message = send_message.clone();
72                thread::spawn(|| client(stream, send_message));
73            }
74            Err(err) => {
75                eprintln!("Error accepting connection: {}", err);
76            }
77        }
78    }
79
80    Ok(())
81}