diff --git a/src/tcp_relay_client/mod.rs b/src/tcp_relay_client/mod.rs index 8010952..6ca6424 100644 --- a/src/tcp_relay_client/mod.rs +++ b/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().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); - panic!(); + Err(std::io::Error::new( + ErrorKind::Other, + format!("Failed to fetch relay configuration from server! {}", e), + ))?; + unreachable!(); } }; diff --git a/src/test/dummy_tcp_sockets.rs b/src/test/dummy_tcp_sockets.rs index 348a5cd..03262f5 100644 --- a/src/test/dummy_tcp_sockets.rs +++ b/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 { diff --git a/src/test/invalid_with_token_auth.rs b/src/test/invalid_with_token_auth.rs new file mode 100644 index 0000000..75ed953 --- /dev/null +++ b/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; +} diff --git a/src/test/mod.rs b/src/test/mod.rs index cc0f23b..fd8d92c 100644 --- a/src/test/mod.rs +++ b/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; diff --git a/src/test/valid_with_token_auth.rs b/src/test/valid_with_token_auth.rs index da3aba6..bbdb7c9 100644 --- a/src/test/valid_with_token_auth.rs +++ b/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, - DummyTCPServer, + dummy_tcp_client_square_root_requests, wait_for_port, 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(); + // Start internal service + let local_server = DummyTCPServer::start(port(1)).await; tokio::spawn(async move { - // Start internal service - let local_server = DummyTCPServer::start(port(1)).await; - local_server.next_conn_square_operations().await; + 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; }