src/tcp_relay_client/mod.rs

@@ -1,6 +1,7 @@
 extern crate core;
 
 use std::error::Error;
+use std::io::ErrorKind;
 use std::sync::Arc;
 
 use futures::future::join_all;
@@ -39,11 +40,13 @@ async fn get_server_config(conf: &ClientConfig) -> Result<RemoteConfig, Box<dyn
         .send()
         .await?;
     if req.status().as_u16() != 200 {
-        log::error!(
+        Err(std::io::Error::new(
+            ErrorKind::Other,
+            format!(
                 "Could not retrieve configuration! (got status {})",
                 req.status()
-        );
+            ),
-        std::process::exit(2);
+        ))?;
     }
 
     Ok(req.json::<RemoteConfig>().await?)
@@ -70,8 +73,11 @@ pub async fn run_app(mut args: ClientConfig) -> std::io::Result<()> {
     let remote_conf = match get_server_config(&args).await {
         Ok(c) => c,
         Err(e) => {
-            log::error!("Failed to fetch relay configuration from server! {}", e);
+            Err(std::io::Error::new(
-            panic!();
+                ErrorKind::Other,
+                format!("Failed to fetch relay configuration from server! {}", e),
+            ))?;
+            unreachable!();
         }
     };
 

src/test/dummy_tcp_sockets.rs

@@ -85,6 +85,12 @@ impl DummyTCPServer {
                 .unwrap();
         }
     }
+
+    pub async fn loop_conn_square_operations(&self) {
+        loop {
+            self.next_conn_square_operations().await
+        }
+    }
 }
 
 pub async fn dummy_tcp_client_read_conn(port: u16) -> Vec<u8> {

src/test/invalid_with_token_auth.rs

@@ -0,0 +1,57 @@
+use tokio::task;
+
+use crate::tcp_relay_client::client_config::ClientConfig;
+use crate::tcp_relay_server::server_config::ServerConfig;
+use crate::test::dummy_tcp_sockets::{wait_for_port, DummyTCPServer};
+use crate::test::{get_port_number, PortsAllocation, LOCALHOST};
+
+const VALID_TOKEN: &str = "AvalidTOKEN";
+const INVALID_TOKEN: &str = "AnInvalidTOKEN";
+
+fn port(index: u16) -> u16 {
+    get_port_number(PortsAllocation::InvalidWithTokenAuth, index)
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 5)]
+async fn invalid_with_token_auth() {
+    let _ = env_logger::builder().is_test(true).try_init();
+
+    tokio::spawn(async move {
+        // Start internal service
+        let local_server = DummyTCPServer::start(port(1)).await;
+        local_server.loop_conn_square_operations().await;
+    });
+
+    let local_set = task::LocalSet::new();
+    local_set
+        .run_until(async move {
+            wait_for_port(port(1)).await;
+
+            // Start server relay
+            task::spawn_local(crate::tcp_relay_server::run_app(ServerConfig {
+                tokens: vec![VALID_TOKEN.to_string()],
+                tokens_file: None,
+                ports: vec![port(1)],
+                upstream_server: "127.0.0.1".to_string(),
+                listen_address: format!("127.0.0.1:{}", port(0)),
+                increment_ports: 1,
+                tls_cert: None,
+                tls_key: None,
+                tls_client_auth_root_cert: None,
+                tls_revocation_list: None,
+            }));
+            wait_for_port(port(0)).await;
+
+            // Start client relay
+            crate::tcp_relay_client::run_app(ClientConfig {
+                token: Some(INVALID_TOKEN.to_string()),
+                relay_url: format!("http://{}:{}", LOCALHOST, port(0)),
+                listen_address: LOCALHOST.to_string(),
+                root_certificate: None,
+                ..Default::default()
+            })
+            .await
+            .unwrap_err();
+        })
+        .await;
+}

src/test/mod.rs

@@ -2,6 +2,7 @@
 enum PortsAllocation {
     DummyTCPServer,
     ValidWithTokenAuth,
+    InvalidWithTokenAuth,
 }
 
 fn get_port_number(alloc: PortsAllocation, index: u16) -> u16 {
@@ -12,4 +13,5 @@ const LOCALHOST: &str = "127.0.0.1";
 
 mod dummy_tcp_sockets;
 
+mod invalid_with_token_auth;
 mod valid_with_token_auth;

src/test/valid_with_token_auth.rs

@@ -3,16 +3,12 @@ use tokio::task;
 use crate::tcp_relay_client::client_config::ClientConfig;
 use crate::tcp_relay_server::server_config::ServerConfig;
 use crate::test::dummy_tcp_sockets::{
-    dummy_tcp_client_square_root_requests, dummy_tcp_client_write_then_read_conn, wait_for_port,
+    dummy_tcp_client_square_root_requests, wait_for_port, DummyTCPServer,
-    DummyTCPServer,
 };
 use crate::test::{get_port_number, PortsAllocation, LOCALHOST};
 
 const VALID_TOKEN: &str = "AvalidTOKEN";
 
-const DATA_1: &[u8] = "DATA1".as_bytes();
-const DATA_2: &[u8] = "DATA2".as_bytes();
-
 fn port(index: u16) -> u16 {
     get_port_number(PortsAllocation::ValidWithTokenAuth, index)
 }
@@ -21,10 +17,10 @@ fn port(index: u16) -> u16 {
 async fn valid_with_token_auth() {
     let _ = env_logger::builder().is_test(true).try_init();
 
-    tokio::spawn(async move {
     // Start internal service
     let local_server = DummyTCPServer::start(port(1)).await;
-        local_server.next_conn_square_operations().await;
+    tokio::spawn(async move {
+        local_server.loop_conn_square_operations().await;
     });
 
     let local_set = task::LocalSet::new();
@@ -56,6 +52,8 @@ async fn valid_with_token_auth() {
                 ..Default::default()
             }));
             wait_for_port(port(2)).await;
+
+            dummy_tcp_client_square_root_requests(port(2), 10).await;
         })
         .await;
 }