diff --git a/Cargo.lock b/Cargo.lock index f1be79d..a8bc3ce 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -540,6 +540,21 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "futures" +version = "0.3.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f21eda599937fba36daeb58a22e8f5cee2d14c4a17b5b7739c7c8e5e3b8230c" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + [[package]] name = "futures-channel" version = "0.3.24" @@ -547,6 +562,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "30bdd20c28fadd505d0fd6712cdfcb0d4b5648baf45faef7f852afb2399bb050" dependencies = [ "futures-core", + "futures-sink", ] [[package]] @@ -555,6 +571,34 @@ version = "0.3.24" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4e5aa3de05362c3fb88de6531e6296e85cde7739cccad4b9dfeeb7f6ebce56bf" +[[package]] +name = "futures-executor" +version = "0.3.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ff63c23854bee61b6e9cd331d523909f238fc7636290b96826e9cfa5faa00ab" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-io" +version = "0.3.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbf4d2a7a308fd4578637c0b17c7e1c7ba127b8f6ba00b29f717e9655d85eb68" + +[[package]] +name = "futures-macro" +version = "0.3.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42cd15d1c7456c04dbdf7e88bcd69760d74f3a798d6444e16974b505b0e62f17" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "futures-sink" version = "0.3.24" @@ -573,8 +617,13 @@ version = "0.3.24" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "44fb6cb1be61cc1d2e43b262516aafcf63b241cffdb1d3fa115f91d9c7b09c90" dependencies = [ + "futures-channel", "futures-core", + "futures-io", + "futures-macro", + "futures-sink", "futures-task", + "memchr", "pin-project-lite", "pin-utils", "slab", @@ -1341,6 +1390,7 @@ dependencies = [ "base", "clap", "env_logger", + "futures", "log", "reqwest", "tokio", diff --git a/tcp_relay_client/Cargo.toml b/tcp_relay_client/Cargo.toml index eaa079e..76d4d01 100644 --- a/tcp_relay_client/Cargo.toml +++ b/tcp_relay_client/Cargo.toml @@ -9,4 +9,5 @@ clap = { version = "3.2.18", features = ["derive", "env"] } log = "0.4.17" env_logger = "0.9.0" reqwest = { version = "0.11", features = ["json"] } -tokio = { version = "1", features = ["full"] } \ No newline at end of file +tokio = { version = "1", features = ["full"] } +futures = "0.3.24" \ No newline at end of file diff --git a/tcp_relay_client/src/main.rs b/tcp_relay_client/src/main.rs index e963693..9045fa0 100644 --- a/tcp_relay_client/src/main.rs +++ b/tcp_relay_client/src/main.rs @@ -1,6 +1,7 @@ use std::sync::Arc; use clap::Parser; +use futures::future::join_all; use base::RemoteConfig; use tcp_relay_client::relay_client::relay_client; @@ -48,17 +49,16 @@ async fn main() -> Result<(), Box> { for port in conf { let listen_address = format!("{}:{}", args.listen_address, port.port); - let args_clone = args.clone(); - let handle = std::thread::spawn(move || { - relay_client(&args_clone.token, port.id, - &args_clone.listen_address, &listen_address); - }); - handles.push(handle); + let h = tokio::spawn(relay_client( + args.token.clone(), + port.id, + args.listen_address.clone(), + listen_address, + )); + handles.push(h); } - for h in handles { - h.join().unwrap(); - } + join_all(handles).await; Ok(()) } \ No newline at end of file diff --git a/tcp_relay_client/src/relay_client.rs b/tcp_relay_client/src/relay_client.rs index 94581c9..3c225c8 100644 --- a/tcp_relay_client/src/relay_client.rs +++ b/tcp_relay_client/src/relay_client.rs @@ -1,9 +1,11 @@ use std::future::Future; -use std::net::TcpListener; +use std::time::Duration; -pub fn relay_client(token: &str, port_id: usize, server: &str, listen_address: &str) { +use tokio::net::TcpListener; + +pub async fn relay_client(token: String, port_id: usize, server: String, listen_address: String) { log::info!("({}) Start to listen on {}", port_id, listen_address); - let listener = match TcpListener::bind(&listen_address) { + let listener = match TcpListener::bind(&listen_address).await { Ok(l) => l, Err(e) => { log::error!("Failed to start to listen on {}!", listen_address); @@ -11,4 +13,14 @@ pub fn relay_client(token: &str, port_id: usize, server: &str, listen_address: & } }; + loop { + let (mut socket, _) = listener.accept().await + .expect("Failed to accept new connection!"); + + let token = token.clone(); + let listen_address = listen_address.clone(); + tokio::spawn(async move { + log::info!("New connection to {}", &listen_address); + }); + } } \ No newline at end of file