diff --git a/Cargo.lock b/Cargo.lock index 2000500..07caeb1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -126,6 +126,17 @@ dependencies = [ "yaml_serde", ] +[[package]] +name = "async-http-proxy" +version = "1.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29faa5d4d308266048bd7505ba55484315a890102f9345b9ff4b87de64201592" +dependencies = [ + "httparse", + "thiserror 1.0.69", + "tokio", +] + [[package]] name = "async-lock" version = "3.4.1" @@ -444,6 +455,7 @@ name = "communicator" version = "0.1.0" dependencies = [ "anyhow", + "async-http-proxy", "bytes", "futures", "futures-util", @@ -455,7 +467,9 @@ dependencies = [ "serde_json", "tokio", "tokio-rustls", + "tokio-socks", "tokio-util", + "url", "webpki-roots", ] @@ -2715,6 +2729,18 @@ dependencies = [ "tokio", ] +[[package]] +name = "tokio-socks" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d4770b8024672c1101b3f6733eab95b18007dbe0847a8afe341fcf79e06043f" +dependencies = [ + "either", + "futures-util", + "thiserror 1.0.69", + "tokio", +] + [[package]] name = "tokio-util" version = "0.7.18" diff --git a/Cargo.toml b/Cargo.toml index 4ec32bb..5c659ca 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -34,4 +34,8 @@ lazy_static = "1.5.0" structopt = "0.3.26" tokio-util = "0.7.18" futures-util = "0.3.32" -twox-hash = "2.1.2" \ No newline at end of file +twox-hash = "2.1.2" +futures = "0.3.32" +tokio-socks = "0.5.2" +url = "2.5.8" +async-http-proxy = "1.2.5" \ No newline at end of file diff --git a/communicator/Cargo.toml b/communicator/Cargo.toml index 6322837..c1e4110 100644 --- a/communicator/Cargo.toml +++ b/communicator/Cargo.toml @@ -17,4 +17,7 @@ webpki-roots = { workspace = true } anyhow = { workspace = true } tokio-util = { workspace = true } futures-util = { workspace = true } -futures = "0.3.32" +futures = { workspace = true } +tokio-socks = { workspace = true } +url = {workspace = true} +async-http-proxy = {workspace = true, features = ["tokio", "runtime-tokio"]} diff --git a/communicator/src/config.rs b/communicator/src/config.rs index 12f0f25..0ad5654 100644 --- a/communicator/src/config.rs +++ b/communicator/src/config.rs @@ -34,6 +34,7 @@ pub struct TcpClientTunnelConfig { pub host: Option, pub url: String, pub token: String, + pub proxy: Option, } impl TcpClientTunnelConfig { @@ -43,6 +44,7 @@ impl TcpClientTunnelConfig { host: self.host, url: self.url, token: self.token, + proxy: self.proxy, } } -} \ No newline at end of file +} diff --git a/communicator/src/stream.rs b/communicator/src/stream.rs index c5e6bcf..40c4d88 100644 --- a/communicator/src/stream.rs +++ b/communicator/src/stream.rs @@ -1,10 +1,11 @@ use anyhow::anyhow; +use async_http_proxy::http_connect_tokio; use bytes::Bytes; use h2::{RecvStream, client, server}; use http::Request; use log::{debug, error, info}; use serde::{Deserialize, Serialize}; -use std::collections::{HashMap, HashSet}; +use std::collections::HashMap; use std::error::Error; use std::fs::File; use std::io::BufReader; @@ -17,6 +18,8 @@ use tokio::time::{Duration, sleep, timeout}; use tokio_rustls::rustls; use tokio_rustls::rustls::pki_types::{CertificateDer, ServerName}; use tokio_rustls::{TlsAcceptor, TlsConnector}; +use tokio_socks::tcp::Socks5Stream; +use url::Url; #[derive(Debug, Copy, Clone, PartialEq, Eq, Serialize, Deserialize)] pub enum Identity { @@ -39,6 +42,7 @@ pub struct ClientTunnelConfig { pub url: String, pub token: String, pub identity: Identity, + pub proxy: Option, } impl Default for ClientTunnelConfig { @@ -48,6 +52,7 @@ impl Default for ClientTunnelConfig { url: "127.0.0.1:3333".to_string(), token: "super_secret_magic_token".to_string(), identity: Identity::Client, + proxy: None, } } } @@ -575,6 +580,47 @@ pub async fn connect_tunnel(config: ClientTunnelConfig) -> Result anyhow::Result { + let Some(proxy) = &config.proxy else { + return Ok(TcpStream::connect(&config.url).await?); + }; + + let parsed_proxy = Url::parse(proxy)?; + + match parsed_proxy.scheme() { + "socks5" => { + let host = parsed_proxy.host_str().unwrap(); + let port = parsed_proxy.port().unwrap_or(1080); + let stream = Socks5Stream::connect((host, port), config.url.clone()) + .await? + .into_inner(); + Ok(stream) + } + "http" | "https" => { + let proxy_addr = format!( + "{}:{}", + parsed_proxy.host_str().unwrap(), + parsed_proxy.port().unwrap_or(80) + ); + let mut stream = TcpStream::connect(proxy_addr).await?; + + let target_url = Url::parse(&format!("tcp://{}", config.url))?; + http_connect_tokio( + &mut stream, + target_url.host_str().unwrap(), + target_url.port().unwrap_or(80), + ) + .await?; + + Ok(stream) + } + _ => Err(anyhow::anyhow!( + "Unsupported proxy scheme: {}", + parsed_proxy.scheme() + )), + } +} + async fn do_client_reconnect( config: &ClientTunnelConfig, current_sid: &mut u64, @@ -596,7 +642,7 @@ async fn do_client_reconnect( *current_sid = sid; resume_tunnel_client(config, sid).await } else { - let mut stream = TcpStream::connect(&config.url).await?; + let mut stream = connect_with_auto_proxy(config).await?; perform_client_handshake(&mut stream, &config.token, config.identity).await?; let raw = upgrade_to_h2_raw(stream, config.identity).await?; *current_sid = 0; @@ -609,7 +655,7 @@ async fn bootstrap_tls_and_get_sid( host: &str, ) -> anyhow::Result<(u64, ())> { let connector = build_client_tls_connector(); - let tcp = TcpStream::connect(&config.url).await?; + let tcp = connect_with_auto_proxy(config).await?; let server_name = ServerName::try_from(host.to_string()) .map_err(|_| anyhow!("Invalid TLS host: {}", host))? .to_owned(); @@ -637,7 +683,7 @@ async fn resume_tunnel_client( config: &ClientTunnelConfig, session_id: u64, ) -> anyhow::Result { - let mut plain_stream = TcpStream::connect(&config.url).await?; + let mut plain_stream = connect_with_auto_proxy(config).await?; plain_stream.write_all(RESUME_MAGIC).await?; plain_stream.write_all(&session_id.to_be_bytes()).await?; upgrade_to_h2_raw(plain_stream, config.identity).await