diff --git a/Cargo.lock b/Cargo.lock index 55cb2e9..d1cd9b8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -418,9 +418,9 @@ checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9" [[package]] name = "hyper" -version = "1.5.1" +version = "1.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "97818827ef4f364230e16705d4706e2897df2bb60617d6ca15d598025a3c481f" +checksum = "256fb8d4bd6413123cc9d91832d78325c48ff41677595be797d90f42969beae0" dependencies = [ "bytes", "futures-channel", @@ -744,7 +744,7 @@ dependencies = [ [[package]] name = "proxer-cli" -version = "0.3.1" +version = "0.3.2" dependencies = [ "base64", "bytes", @@ -754,6 +754,7 @@ dependencies = [ "hyper", "hyper-util", "json5", + "rlimit", "serde", "serde_json", "sha2", @@ -826,6 +827,15 @@ version = "0.8.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2b15c43186be67a4fd63bee50d0303afffcef381492ebe2c5d87f324e1b8815c" +[[package]] +name = "rlimit" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7043b63bd0cd1aaa628e476b80e6d4023a3b50eb32789f2728908107bd0c793a" +dependencies = [ + "libc", +] + [[package]] name = "rustc-demangle" version = "0.1.24" @@ -891,18 +901,18 @@ dependencies = [ [[package]] name = "serde" -version = "1.0.215" +version = "1.0.216" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6513c1ad0b11a9376da888e3e0baa0077f1aed55c17f50e7b2397136129fb88f" +checksum = "0b9781016e935a97e8beecf0c933758c97a5520d32930e460142b4cd80c6338e" dependencies = [ "serde_derive", ] [[package]] name = "serde_derive" -version = "1.0.215" +version = "1.0.216" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ad1e866f866923f252f05c889987993144fb74e722403468a4ebd70c3cd756c0" +checksum = "46f859dbbf73865c6627ed570e78961cd3ac92407a2d117204c49232485da55e" dependencies = [ "proc-macro2", "quote", diff --git a/Cargo.toml b/Cargo.toml index e40551d..7b672f8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "proxer-cli" -version = "0.3.1" +version = "0.3.2" edition = "2021" authors = ["doroved"] description = "Proxy TCP traffic on macOS with domain filtering." @@ -13,10 +13,10 @@ categories = ["network-programming", "command-line-utilities"] # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html [dependencies] -hyper = { version = "1.5.1", features = ["http1", "http2", "server", "client"] } +hyper = { version = "1.5.2", features = ["http1", "http2", "server", "client"] } hyper-util = { version = "0.1.10", features = ["tokio"] } -serde = { version = "1.0.215", features = ["derive"] } +serde = { version = "1.0.216", features = ["derive"] } serde_json = "1.0.133" tokio = { version = "1.42.0", features = ["full"] } @@ -34,6 +34,7 @@ http-body-util = "0.1.2" tracing = "0.1.41" tracing-subscriber = { version = "0.3.19", features = ["env-filter"] } json5 = "0.4.1" +rlimit = "0.10.2" [profile.release] diff --git a/src/server/mod.rs b/src/server/mod.rs index 8142d80..4dc6aa3 100644 --- a/src/server/mod.rs +++ b/src/server/mod.rs @@ -42,6 +42,25 @@ pub struct ProxyConfig { } pub async fn run() -> Result<(), Box> { + tracing::info!( + "Default open connection limit: {:?}", + rlimit::Resource::NOFILE.get_soft()? + ); + + let connection_limit = match rlimit::Resource::NOFILE.get() { + Ok(limit) if limit.0 < 1024 * 10 => { + tracing::info!("Setting open connection limit to {}", limit.1); + limit.1 + } + Ok(limit) => limit.0, + Err(err) => { + tracing::error!("Failed to get open connection limit: {}", err); + return Err(Box::new(err)); + } + }; + + let _ = rlimit::setrlimit(rlimit::Resource::NOFILE, connection_limit, rlimit::INFINITY); + // Close all proxer-cli processes terminate_proxer(); @@ -82,7 +101,6 @@ pub async fn run() -> Result<(), Box> { loop { let (stream, _) = listener.accept().await?; - let io = TokioIo::new(stream); let proxy_config = Arc::clone(&proxy_config); let service = service_fn(move |req: Request| { @@ -94,7 +112,7 @@ pub async fn run() -> Result<(), Box> { if let Err(err) = http1::Builder::new() .preserve_header_case(true) .title_case_headers(true) - .serve_connection(io, service) + .serve_connection(TokioIo::new(stream), service) .with_upgrades() .await { diff --git a/src/server/proxy.rs b/src/server/proxy.rs index 7d67efe..29ae4f1 100644 --- a/src/server/proxy.rs +++ b/src/server/proxy.rs @@ -58,12 +58,10 @@ pub async fn handle_request( match TcpStream::connect((host, port)).await { Ok(stream) => { - let io = TokioIo::new(stream); - let (mut sender, conn) = Builder::new() .preserve_header_case(true) .title_case_headers(true) - .handshake(io) + .handshake(TokioIo::new(stream)) .await?; tokio::spawn(async move { diff --git a/src/server/tunnel.rs b/src/server/tunnel.rs index 2385332..0fe2bcc 100644 --- a/src/server/tunnel.rs +++ b/src/server/tunnel.rs @@ -29,8 +29,13 @@ pub async fn tunnel_direct( ) -> Result<(), Box> { tracing::info!("{addr} → DIRECT connection"); - let mut server = TcpStream::connect(&addr).await?; - let _ = tokio::io::copy_bidirectional(&mut TokioIo::new(upgraded), &mut server).await?; + let mut server = timeout(Duration::from_secs(30), TcpStream::connect(&addr)).await??; + + timeout( + Duration::from_secs(30), + tokio::io::copy_bidirectional(&mut TokioIo::new(upgraded), &mut server), + ) + .await??; Ok(()) } @@ -52,14 +57,13 @@ pub async fn tunnel_via_proxy( let proxy_addr = format!("{}:{}", &proxy.host, &proxy.port); let proxy_host = proxy_addr.split(':').next().unwrap(); - let tcp = TcpStream::connect(&proxy_addr).await?; - let mut upgraded = TokioIo::new(upgraded); + let tcp = timeout(Duration::from_secs(30), TcpStream::connect(&proxy_addr)).await??; let mut stream: Pin> = if proxy.scheme.eq_ignore_ascii_case("http") { Box::pin(tcp) } else { - let tls = TlsConnector::from(native_tls::TlsConnector::new().unwrap()); - Box::pin(tls.connect(proxy_host, tcp).await?) + let tls = TlsConnector::from(native_tls::TlsConnector::new()?); + Box::pin(timeout(Duration::from_secs(30), tls.connect(proxy_host, tcp)).await??) }; let mut connect_req = format!( @@ -100,8 +104,11 @@ pub async fn tunnel_via_proxy( return Err(format!("{:?}", String::from_utf8_lossy(&response[..n])).into()); } - let (from_client, from_server) = - tokio::io::copy_bidirectional(&mut upgraded, &mut stream).await?; + let (from_client, from_server) = timeout( + Duration::from_secs(30), + tokio::io::copy_bidirectional(&mut TokioIo::new(upgraded), &mut stream), + ) + .await??; tracing::info!( "{addr} → Client wrote {:.2} KB and received {:.2} KB", @@ -109,5 +116,7 @@ pub async fn tunnel_via_proxy( from_server as f64 / 1024.0 ); + stream.shutdown().await?; + Ok(()) }