Moved upstreams to their own dedicated namespace

Signed-off-by: Jacob Kiers <code@kiers.eu>

src/upstreams/mod.rs

@@ -0,0 +1,51 @@
+mod proxy_to_upstream;
+
+use log::debug;
+use serde::Deserialize;
+use std::error::Error;
+use tokio::io;
+use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt};
+use tokio::net::TcpStream;
+
+pub use crate::upstreams::proxy_to_upstream::ProxyToUpstream;
+
+#[derive(Debug, Clone, Deserialize)]
+pub enum Upstream {
+    Ban,
+    Echo,
+    Proxy(ProxyToUpstream),
+}
+
+impl Upstream {
+    pub(crate) async fn process(&self, mut inbound: TcpStream) -> Result<(), Box<dyn Error>> {
+        match self {
+            Upstream::Ban => {
+                inbound.shutdown().await?;
+            }
+            Upstream::Echo => {
+                let (mut ri, mut wi) = io::split(inbound);
+                let inbound_to_inbound = copy(&mut ri, &mut wi);
+                let bytes_tx = inbound_to_inbound.await;
+                debug!("Bytes read: {:?}", bytes_tx);
+            }
+            Upstream::Proxy(config) => {
+                config.proxy(inbound).await?;
+            }
+        };
+        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) => {
+            let _ = writer.shutdown().await;
+            Ok(u64)
+        }
+        Err(_) => Ok(0),
+    }
+}

src/upstreams/proxy_to_upstream.rs

@@ -0,0 +1,68 @@
+use crate::servers::upstream_address::UpstreamAddress;
+
+use crate::upstreams::copy;
+use futures::future::try_join;
+use log::{debug, error};
+use serde::Deserialize;
+use std::net::SocketAddr;
+use tokio::io;
+use tokio::net::TcpStream;
+use tokio::sync::Mutex;
+
+#[derive(Debug, Default)]
+struct Addr(Mutex<UpstreamAddress>);
+
+impl Clone for Addr {
+    fn clone(&self) -> Self {
+        tokio::task::block_in_place(|| Self(Mutex::new(self.0.blocking_lock().clone())))
+    }
+}
+
+#[derive(Debug, Clone, Deserialize, Default)]
+pub struct ProxyToUpstream {
+    pub addr: String,
+    pub protocol: String,
+    #[serde(skip_deserializing)]
+    addresses: Addr,
+}
+
+impl ProxyToUpstream {
+    pub async fn resolve_addresses(&self) -> std::io::Result<Vec<SocketAddr>> {
+        let mut addr = self.addresses.0.lock().await;
+        addr.resolve((*self.protocol).into()).await
+    }
+
+    pub fn new(address: String, protocol: String) -> Self {
+        Self {
+            addr: address.clone(),
+            protocol,
+            addresses: Addr(Mutex::new(UpstreamAddress::new(address))),
+        }
+    }
+
+    pub(crate) async fn proxy(&self, inbound: TcpStream) -> Result<(), Box<dyn std::error::Error>> {
+        let outbound = match self.protocol.as_ref() {
+            "tcp4" | "tcp6" | "tcp" => {
+                TcpStream::connect(self.resolve_addresses().await?.as_slice()).await?
+            }
+            _ => {
+                error!("Reached unknown protocol: {:?}", self.protocol);
+                return Err("Reached unknown protocol".into());
+            }
+        };
+
+        debug!("Connected to {:?}", outbound.peer_addr().unwrap());
+
+        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(())
+    }
+}