Disabled timeout for server (#81)

- server does not require connection mode in UPP_CASE anymore
- client does not crash any more if upstream offline
- server does no longer just abandon TCP connections after 60 seconds of inactivity

Author: TheCataliasTNT2k

zia-client/src/handler.rs

@@ -38,7 +38,13 @@ impl<U: Upstream + Send> UdpHandler<U> {
         Entry::Vacant(vacant) => {
           info!("New socket at {}/udp, opening upstream connection...", addr);
 
-          let (sink, mut stream) = self.upstream.open().await?;
+          let (sink, mut stream) = match self.upstream.open().await {
+            Ok(conn) => conn,
+            Err(err) => {
+              warn!("Error while opening upstream connection: {err}");
+              continue;
+            }
+          };
 
           let known = self.known.clone();
           let socket = socket.clone();

zia-client/src/upstream/mod.rs

@@ -13,7 +13,7 @@ pub(crate) async fn transmit(
   upstream: &Url,
   proxy: &Option<Url>,
 ) -> anyhow::Result<()> {
-  match upstream.scheme() {
+  match upstream.scheme().to_lowercase().as_str() {
     "tcp" | "tcps" => tcp::transmit(socket, upstream, proxy).await,
     "ws" | "wss" => ws::transmit(socket, upstream, proxy).await,
     _ => Err(anyhow!("Unsupported upstream scheme {}", upstream.scheme())),

zia-common/src/forwarding/tcp.rs

@@ -3,7 +3,6 @@
 use std::convert::{Infallible, TryFrom};
 use std::mem;
 use std::sync::Arc;
-use std::time::Duration;
 
 use anyhow::Context;
 use futures_util::future::select;
@@ -12,7 +11,7 @@ use tokio::io::{split, AsyncReadExt, AsyncWriteExt, ReadHalf, WriteHalf};
 use tokio::net::{TcpStream, UdpSocket};
 use tracing::error;
 
-use crate::{maybe_timeout, Stream};
+use crate::{Stream};
 
 /// A UDP datagram header has a 16 bit field containing an unsigned integer
 /// describing the length of the datagram (including the header itself).
@@ -29,14 +28,13 @@ const HEADER_LEN: usize = mem::size_of::<u16>();
 pub async fn process_udp_over_tcp(
   udp_socket: UdpSocket,
   tcp_stream: Stream<TcpStream>,
-  tcp_recv_timeout: Option<Duration>,
 ) {
   let udp_in = Arc::new(udp_socket);
   let udp_out = udp_in.clone();
   let (tcp_in, tcp_out) = split(tcp_stream);
 
   let tcp2udp = tokio::spawn(async move {
-    if let Err(error) = process_tcp2udp(tcp_in, udp_out, tcp_recv_timeout).await {
+    if let Err(error) = process_tcp2udp(tcp_in, udp_out).await {
       error!("Error: {}", error);
     }
   });
@@ -58,16 +56,14 @@ pub async fn process_udp_over_tcp(
 /// Returns if the TCP socket is closed, or an IO error happens on either socket.
 async fn process_tcp2udp(
   mut tcp_in: ReadHalf<Stream<TcpStream>>,
-  udp_out: Arc<UdpSocket>,
-  tcp_recv_timeout: Option<Duration>,
+  udp_out: Arc<UdpSocket>
 ) -> anyhow::Result<()> {
   let mut buffer = datagram_buffer();
   // `buffer` has unprocessed data from the TCP socket up until this index.
   let mut unprocessed_i = 0;
   loop {
-    let tcp_read_len = maybe_timeout(tcp_recv_timeout, tcp_in.read(&mut buffer[unprocessed_i..]))
+    let tcp_read_len = tcp_in.read(&mut buffer[unprocessed_i..])
       .await
-      .context("Timeout while reading from TCP")?
       .context("Failed reading from TCP")?;
     if tcp_read_len == 0 {
       break;

zia-common/src/forwarding/ws.rs

@@ -3,7 +3,6 @@
 use std::convert::Infallible;
 use std::mem;
 use std::sync::Arc;
-use std::time::Duration;
 
 use anyhow::Context;
 use futures_util::future::select;
@@ -17,7 +16,7 @@ use tokio_tungstenite::tungstenite::Message;
 use tokio_tungstenite::WebSocketStream;
 use tracing::error;
 
-use crate::{maybe_timeout, Stream};
+use crate::{Stream};
 
 /// A UDP datagram header has a 16 bit field containing an unsigned integer
 /// describing the length of the datagram (including the header itself).
@@ -32,7 +31,6 @@ const MAX_DATAGRAM_SIZE: usize = u16::MAX as usize - mem::size_of::<u16>();
 pub async fn process_udp_over_ws(
   udp_socket: UdpSocket,
   ws_stream: WebSocketStream<Stream<TcpStream>>,
-  ws_recv_timeout: Option<Duration>,
 ) {
   let (mut ws_out, mut ws_in) = ws_stream.split();
 
@@ -41,7 +39,7 @@ pub async fn process_udp_over_ws(
     let udp_out = udp_in.clone();
 
     let ws2udp = async {
-      if let Err(error) = process_ws2udp(&mut ws_in, udp_out, ws_recv_timeout).await {
+      if let Err(error) = process_ws2udp(&mut ws_in, udp_out).await {
         error!("Error: {}", error);
       }
     };
@@ -77,11 +75,9 @@ pub async fn process_udp_over_ws(
 async fn process_ws2udp(
   ws_in: &mut SplitStream<WebSocketStream<Stream<TcpStream>>>,
   udp_out: Arc<UdpSocket>,
-  ws_recv_timeout: Option<Duration>,
 ) -> anyhow::Result<()> {
-  while let Some(message) = maybe_timeout(ws_recv_timeout, ws_in.next())
+  while let Some(message) = ws_in.next()
     .await
-    .context("Timeout while reading from WS")?
     .transpose()
     .context("Failed reading from WS")?
   {

zia-common/src/lib.rs

@@ -1,21 +1,5 @@
-use std::future::Future;
-use std::time::Duration;
-
-use tokio::time::error::Elapsed;
-use tokio::time::timeout;
-
 pub use crate::forwarding::*;
 pub use crate::stream::*;
 
 mod forwarding;
 mod stream;
-
-pub async fn maybe_timeout<F: Future>(
-  duration: Option<Duration>,
-  future: F,
-) -> Result<F::Output, Elapsed> {
-  match duration {
-    Some(duration) => timeout(duration, future).await,
-    None => Ok(future.await),
-  }
-}

zia-server/src/cfg.rs

@@ -10,15 +10,13 @@ pub(crate) struct ClientCfg {
   pub(crate) listen_addr: SocketAddr,
   #[arg(short, long, env = "ZIA_UPSTREAM")]
   pub(crate) upstream: String,
-  #[arg(short, long, env = "ZIA_MODE", default_value = "WS", value_enum)]
+  #[arg(short, long, env = "ZIA_MODE", default_value = "WS", value_enum, ignore_case(true))]
   pub(crate) mode: Mode,
 }
 
 #[derive(ValueEnum, Clone)]
 pub(crate) enum Mode {
-  #[value(rename_all = "SCREAMING_SNAKE_CASE")]
   Ws,
-  #[value(rename_all = "SCREAMING_SNAKE_CASE")]
   Tcp,
 }
 

zia-server/src/listener/tcp.rs

@@ -1,5 +1,4 @@
 use std::net::SocketAddr;
-use std::time::Duration;
 
 use tokio::net::{TcpStream, UdpSocket};
 use tracing::{error, info};
@@ -48,7 +47,7 @@ impl TcpListener {
       downstream_addr
     );
 
-    process_udp_over_tcp(upstream, Plain(downstream), Some(Duration::from_secs(60))).await;
+    process_udp_over_tcp(upstream, Plain(downstream)).await;
 
     info!("Connection with downstream {} closed...", downstream_addr);
 

zia-server/src/listener/ws.rs

@@ -1,5 +1,4 @@
 use std::net::SocketAddr;
-use std::time::Duration;
 
 use tokio::net::{TcpListener, TcpStream, UdpSocket};
 use tracing::{error, info};
@@ -49,7 +48,7 @@ impl WsListener {
       downstream_addr
     );
 
-    process_udp_over_ws(upstream, downstream, Some(Duration::from_secs(60))).await;
+    process_udp_over_ws(upstream, downstream).await;
 
     info!("Connection with downstream {} closed...", downstream_addr);