Added support for domain names as upstream

Author: MarcelCoding

zia-server/src/cfg.rs

@@ -7,7 +7,7 @@ use serde::Deserialize;
 #[derive(Deserialize)]
 pub(crate) struct ClientCfg {
   pub(crate) listen_addr: SocketAddr,
-  pub(crate) upstream: SocketAddr,
+  pub(crate) upstream: String,
   pub(crate) mode: Mode,
 }
 

zia-server/src/listener/mod.rs

@@ -1,5 +1,3 @@
-use std::net::SocketAddr;
-
 pub(crate) use self::tcp::*;
 pub(crate) use self::ws::*;
 
@@ -8,5 +6,5 @@ mod ws;
 
 #[async_trait::async_trait]
 pub(crate) trait Listener {
-  async fn listen(&self, upstream: SocketAddr) -> anyhow::Result<()>;
+  async fn listen(&self, upstream: &str) -> anyhow::Result<()>;
 }

zia-server/src/listener/tcp.rs

@@ -15,11 +15,12 @@ pub(crate) struct TcpListener {
 
 #[async_trait::async_trait]
 impl Listener for TcpListener {
-  async fn listen(&self, upstream: SocketAddr) -> anyhow::Result<()> {
+  async fn listen(&self, upstream: &str) -> anyhow::Result<()> {
     let listener = tokio::net::TcpListener::bind(self.addr).await?;
 
     loop {
       let (sock, _) = listener.accept().await?;
+      let upstream = upstream.to_string();
 
       tokio::spawn(async move {
         if let Err(err) = Self::handle(sock, upstream).await {
@@ -31,7 +32,7 @@ impl Listener for TcpListener {
 }
 
 impl TcpListener {
-  async fn handle(downstream: TcpStream, upstream_addr: SocketAddr) -> anyhow::Result<()> {
+  async fn handle(downstream: TcpStream, upstream_addr: String) -> anyhow::Result<()> {
     downstream.set_nodelay(true)?;
     let downstream_addr = downstream.peer_addr()?;
     info!("New downstream connection: {}", downstream_addr);

zia-server/src/listener/ws.rs

@@ -3,7 +3,7 @@ use std::time::Duration;
 
 use tokio::net::{TcpListener, TcpStream, UdpSocket};
 use tracing::{error, info};
-use zia_common::process_udp_over_ws;
+use zia_common::{process_udp_over_ws, Stream};
 use zia_common::Stream::Plain;
 
 use crate::listener::Listener;
@@ -14,12 +14,13 @@ pub(crate) struct WsListener {
 
 #[async_trait::async_trait]
 impl Listener for WsListener {
-  async fn listen(&self, upstream: SocketAddr) -> anyhow::Result<()> {
+  async fn listen(&self, upstream: &str) -> anyhow::Result<()> {
     let listener = TcpListener::bind(self.addr).await?;
     info!("Listening on ws://{}...", listener.local_addr()?);
 
     loop {
       let (sock, _) = listener.accept().await?;
+      let upstream = upstream.to_string();
 
       tokio::spawn(async move {
         if let Err(err) = Self::handle(sock, upstream).await {
@@ -31,7 +32,7 @@ impl Listener for WsListener {
 }
 
 impl WsListener {
-  async fn handle(downstream: TcpStream, upstream_addr: SocketAddr) -> anyhow::Result<()> {
+  async fn handle(downstream: TcpStream, upstream_addr: String) -> anyhow::Result<()> {
     downstream.set_nodelay(true)?;
     let downstream_addr = downstream.peer_addr()?;
     info!("New downstream connection: {}", downstream_addr);

zia-server/src/main.rs

@@ -1,7 +1,7 @@
 #[cfg(not(target_env = "msvc"))]
 use tikv_jemallocator::Jemalloc;
-use tokio::signal::ctrl_c;
 use tokio::select;
+use tokio::signal::ctrl_c;
 use tracing::info;
 
 use crate::cfg::{ClientCfg, Mode};
@@ -32,7 +32,7 @@ async fn main() -> anyhow::Result<()> {
   info!("Listening in {}://{}...", config.mode, config.listen_addr);
 
   select! {
-    result = listener.listen(config.upstream) => {
+    result = listener.listen(&config.upstream) => {
       result?;
       info!("Socket closed, quitting...");
     },