First commit

src/servers/mod.rs

@@ -0,0 +1,208 @@
+use futures::future::try_join;
+use log::{debug, error, info, warn};
+use std::collections::{HashMap, HashSet};
+use std::net::SocketAddr;
+use std::sync::Arc;
+use tokio::io;
+use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
+use tokio::net::{TcpListener, TcpStream};
+use tokio::task::JoinHandle;
+
+mod tls;
+use self::tls::get_sni;
+use crate::config::BaseConfig;
+
+#[derive(Debug)]
+pub struct Server {
+    pub proxies: Vec<Arc<Proxy>>,
+    pub config: BaseConfig,
+}
+
+#[derive(Debug, Clone)]
+pub struct Proxy {
+    pub name: String,
+    pub listen: SocketAddr,
+    pub tls: bool,
+    pub sni: Option<HashMap<String, String>>,
+    pub default: String,
+    pub upstream: HashMap<String, String>,
+}
+
+impl Server {
+    pub fn new(config: BaseConfig) -> Self {
+        let mut new_server = Server {
+            proxies: Vec::new(),
+            config: config.clone(),
+        };
+
+        for (name, proxy) in config.servers.iter() {
+            let tls = proxy.tls.unwrap_or(false);
+            let sni = proxy.sni.clone();
+            let default = proxy.default.clone().unwrap_or_else(|| "ban".to_string());
+            let upstream = config.upstream.clone();
+            let mut upstream_set: HashSet<String> = HashSet::new();
+            for (key, _) in &upstream {
+                if key.eq("ban") || key.eq("echo") {
+                    continue;
+                }
+                upstream_set.insert(key.clone());
+            }
+            for listen in proxy.listen.clone() {
+                println!("{:?}", listen);
+                let listen_addr: SocketAddr = match listen.parse() {
+                    Ok(addr) => addr,
+                    Err(_) => {
+                        error!("Invalid listen address: {}", listen);
+                        continue;
+                    }
+                };
+                let proxy = Proxy {
+                    name: name.clone(),
+                    listen: listen_addr,
+                    tls,
+                    sni: sni.clone(),
+                    default: default.clone(),
+                    upstream: upstream.clone(),
+                };
+                new_server.proxies.push(Arc::new(proxy));
+            }
+        }
+
+        new_server
+    }
+
+    #[tokio::main]
+    pub async fn run(&mut self) -> Result<(), Box<dyn std::error::Error>> {
+        let proxies = self.proxies.clone();
+        let mut handles: Vec<JoinHandle<()>> = Vec::new();
+
+        for config in proxies {
+            info!("Starting server {} on {}", config.name, config.listen);
+            let handle = tokio::spawn(async move {
+                let _ = proxy(config).await;
+            });
+            handles.push(handle);
+        }
+
+        for handle in handles {
+            handle.await?;
+        }
+        Ok(())
+    }
+}
+
+async fn proxy(config: Arc<Proxy>) -> Result<(), Box<dyn std::error::Error>> {
+    let listener = TcpListener::bind(config.listen).await?;
+    let config = config.clone();
+
+    loop {
+        let thread_proxy = config.clone();
+        match listener.accept().await {
+            Err(err) => {
+                error!("Failed to accept connection: {}", err);
+                return Err(Box::new(err));
+            }
+            Ok((stream, _)) => {
+                tokio::spawn(async move {
+                    match accept(stream, thread_proxy).await {
+                        Ok(_) => {}
+                        Err(err) => {
+                            error!("Relay thread returned an error: {}", err);
+                        }
+                    };
+                });
+            }
+        }
+    }
+}
+
+async fn accept(inbound: TcpStream, proxy: Arc<Proxy>) -> Result<(), Box<dyn std::error::Error>> {
+    debug!("New connection from {:?}", inbound.peer_addr()?);
+
+    let upstream_name = match proxy.tls {
+        false => proxy.default.clone(),
+        true => {
+            let mut hello_buf = [0u8; 1024];
+            inbound.peek(&mut hello_buf).await?;
+            let snis = get_sni(&hello_buf);
+            if snis.is_empty() {
+                proxy.default.clone()
+            } else {
+                match proxy.sni.clone() {
+                    Some(sni_map) => {
+                        let mut upstream = proxy.default.clone();
+                        for sni in snis {
+                            let m = sni_map.get(&sni);
+                            if m.is_some() {
+                                upstream = m.unwrap().clone();
+                                break;
+                            }
+                        }
+                        upstream
+                    }
+                    None => proxy.default.clone(),
+                }
+            }
+        }
+    };
+
+    debug!("Upstream: {}", upstream_name);
+
+    let upstream = match proxy.upstream.get(&upstream_name) {
+        Some(upstream) => upstream,
+        None => {
+            warn!(
+                "No upstream named {:?} on server {:?}",
+                proxy.default, proxy.name
+            );
+            return process(inbound, &proxy.default).await;
+        }
+    };
+    return process(inbound, upstream).await;
+}
+
+async fn process(mut inbound: TcpStream, upstream: &str) -> Result<(), Box<dyn std::error::Error>> {
+    if upstream == "ban" {
+        let _ = inbound.shutdown();
+        return Ok(());
+    } else if upstream == "echo" {
+        loop {
+            let mut buf = [0u8; 1];
+            let b = inbound.read(&mut buf).await?;
+            if b == 0 {
+                break;
+            } else {
+                inbound.write(&buf).await?;
+            }
+        }
+        return Ok(());
+    }
+
+    let outbound = TcpStream::connect(upstream).await?;
+
+    let (mut ri, mut wi) = io::split(inbound);
+    let (mut ro, mut wo) = io::split(outbound);
+
+    let inbound_to_outbound = copy(&mut ri, &mut wo);
+    let outbound_to_inbound = copy(&mut ro, &mut wi);
+
+    let (bytes_tx, bytes_rx) = try_join(inbound_to_outbound, outbound_to_inbound).await?;
+
+    debug!("Bytes read: {:?} write: {:?}", bytes_tx, bytes_rx);
+
+    Ok(())
+}
+
+async fn copy<'a, R, W>(reader: &'a mut R, writer: &'a mut W) -> io::Result<u64>
+where
+    R: AsyncRead + Unpin + ?Sized,
+    W: AsyncWrite + Unpin + ?Sized,
+{
+    match io::copy(reader, writer).await {
+        Ok(u64) => {
+            writer.shutdown().await?;
+            Ok(u64)
+        }
+        Err(_) => Ok(0),
+    }
+}