diff --git a/slinger-mitm/Cargo.toml b/slinger-mitm/Cargo.toml index 5212b63..b7f95f9 100644 --- a/slinger-mitm/Cargo.toml +++ b/slinger-mitm/Cargo.toml @@ -10,7 +10,7 @@ license.workspace = true [dependencies] # Core slinger library -slinger = { path = "../slinger", features = ["tls", "rustls", "http2"] } +slinger = { version = "0.2.10", features = ["tls", "rustls", "http2"] } # Workspace dependencies thiserror.workspace = true http.workspace = true diff --git a/slinger-mitm/src/interceptor.rs b/slinger-mitm/src/interceptor.rs index c03a7a0..e60df7e 100644 --- a/slinger-mitm/src/interceptor.rs +++ b/slinger-mitm/src/interceptor.rs @@ -92,9 +92,9 @@ pub struct LoggingInterceptor; #[async_trait::async_trait] impl RequestInterceptor for LoggingInterceptor { async fn intercept_request(&self, request: Request) -> Result> { - println!("[MITM] Request: {} {}", request.method(), request.uri()); + tracing::info!("[MITM] Request: {} {}", request.method(), request.uri()); for (name, value) in request.headers() { - println!(" {}: {:?}", name, value); + tracing::info!(" {}: {:?}", name, value); } Ok(Some(request)) } @@ -103,11 +103,11 @@ impl RequestInterceptor for LoggingInterceptor { #[async_trait::async_trait] impl ResponseInterceptor for LoggingInterceptor { async fn intercept_response(&self, response: Response) -> Result> { - println!("[MITM] Response: {}", response.status_code()); + tracing::info!("[MITM] Response: {}", response.status_code()); for (name, value) in response.headers() { - println!(" {}: {:?}", name, value); + tracing::info!(" {}: {:?}", name, value); } - // println!("[MITM] Response Body: {:?}", response.body()); + // tracing::debug!("[MITM] Response Body: {:?}", response.body()); Ok(Some(response)) } } diff --git a/slinger-mitm/src/lib.rs b/slinger-mitm/src/lib.rs index c251aa6..668acd7 100644 --- a/slinger-mitm/src/lib.rs +++ b/slinger-mitm/src/lib.rs @@ -35,7 +35,7 @@ pub use ca::{CertificateAuthority, CertificateManager}; pub use error::{Error, Result}; pub use interceptor::{Interceptor, InterceptorHandler, RequestInterceptor, ResponseInterceptor}; pub use proxy::{MitmConfig, MitmProxy}; -pub use server::ProxyServer; +pub use server::{ProxyServer, ProxyServerBuilder}; pub use socks5::{Socks5Server, TargetAddr}; #[cfg(test)] diff --git a/slinger-mitm/src/server.rs b/slinger-mitm/src/server.rs index f19cb80..68aff7f 100644 --- a/slinger-mitm/src/server.rs +++ b/slinger-mitm/src/server.rs @@ -6,7 +6,7 @@ use crate::interceptor::InterceptorHandler; use crate::proxy::MitmConfig; use bytes::Bytes; use http::Version; -use slinger::{Client, Request, Response}; +use slinger::{Client, ClientBuilder, Request, Response}; use std::sync::Arc; use std::time::Duration; use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader}; @@ -24,6 +24,133 @@ pub struct ProxyServer { client: Client, } +/// Builder for `ProxyServer`. +/// +/// Allows configuring the server and the inner `slinger::Client`. +#[derive(Default)] +pub struct ProxyServerBuilder { + config: Option, + cert_manager: Option>, + interceptor_handler: Option>>, + client: Option, + // Optional client configurator: takes a ClientBuilder and returns a configured ClientBuilder + client_config: Option ClientBuilder + Send + Sync>>, +} + +impl ProxyServerBuilder { + /// Start building from an existing `ProxyServer` configuration. + pub fn from_server(server: &ProxyServer) -> Self { + Self { + config: Some(server.config.clone()), + cert_manager: Some(server.cert_manager.clone()), + interceptor_handler: Some(server.interceptor_handler.clone()), + client: Some(server.client.clone()), + client_config: None, + } + } + + /// Set the `MitmConfig` to use. + pub fn config(mut self, config: MitmConfig) -> Self { + self.config = Some(config); + self + } + + /// Set the `CertificateManager` to use. + pub fn cert_manager(mut self, cert_manager: Arc) -> Self { + self.cert_manager = Some(cert_manager); + self + } + + /// Set the `InterceptorHandler` to use. + pub fn interceptor_handler(mut self, handler: Arc>) -> Self { + self.interceptor_handler = Some(handler); + self + } + + /// Provide a fully constructed `slinger::Client` to use. + pub fn client(mut self, client: Client) -> Self { + self.client = Some(client); + self + } + + /// Configure the internal `slinger::Client` using a closure that accepts a + /// `ClientBuilder` and returns a configured `ClientBuilder`. + pub fn configure_client(mut self, f: F) -> Self + where + F: Fn(ClientBuilder) -> ClientBuilder + Send + Sync + 'static, + { + self.client_config = Some(Box::new(f)); + self + } + + /// Build the `ProxyServer`. + /// + /// Priority for creating the inner `Client`: + /// 1. If `client` is provided, use it. + /// 2. Else if `client_config` is provided, call it with `Client::builder()`. + /// 3. Else fall back to default behavior: honor `config.upstream_proxy` if present + /// and otherwise create a default client similar to `ProxyServer::new`. + pub fn build(self) -> Result { + // Resolve config + let config = self.config.unwrap_or_default(); + + // For synchronous build we require a pre-created CertificateManager because + // creation is async. Callers who don't have one should use `build_async()`. + let cert_manager = match self.cert_manager { + Some(c) => c, + None => { + return Err(Error::proxy_error( + "CertificateManager not provided; use ProxyServer::builder().build_async().await to create one automatically".to_string(), + )) + } + }; + + // Resolve interceptor handler + let interceptor_handler = self + .interceptor_handler + .unwrap_or_else(|| Arc::new(RwLock::new(InterceptorHandler::new()))); + + // Resolve client + let client = if let Some(client) = self.client { + client + } else if let Some(cfg_fn) = self.client_config { + let builder = Client::builder(); + let configured = cfg_fn(builder); + configured + .build() + .map_err(|e| Error::proxy_error(format!("Failed to build client: {}", e)))? + } else { + // Fallback: honor upstream_proxy in config similar to ProxyServer::new + if let Some(proxy) = &config.upstream_proxy { + Client::builder() + .timeout(Some(Duration::from_secs(60))) + .keepalive(true) + .proxy(proxy.clone()) + .build() + .map_err(|e| { + Error::proxy_error(format!( + "Failed to build client with proxy {}: {}", + proxy.uri(), + e + )) + })? + } else { + Client::builder() + .keepalive(true) + .build() + .map_err(|e| Error::proxy_error(format!("Failed to build default client: {}", e)))? + } + }; + + Ok(ProxyServer { + config, + cert_manager, + interceptor_handler, + client, + }) + } +} + impl ProxyServer { /// Create a new proxy server pub fn new( @@ -77,12 +204,12 @@ impl ProxyServer { if let Err(e) = Self::handle_connection(stream, config, cert_manager, interceptor, client).await { - eprintln!("[MITM] Error handling connection: {}", e); + tracing::error!("[MITM] Error handling connection: {}", e); } }); } Err(e) => { - eprintln!("[MITM] Failed to accept connection: {}", e); + tracing::error!("[MITM] Failed to accept connection: {}", e); } } } @@ -140,9 +267,7 @@ impl ProxyServer { } } } - Err(e) => { - Err(e) - } + Err(e) => Err(e), } } else { let mut request_line = vec![first_byte[0]]; @@ -203,7 +328,7 @@ impl ProxyServer { /// Handle HTTPS CONNECT with MITM interception async fn handle_https_connect( - mut client_stream: TcpStream, + client_stream: TcpStream, uri: &str, cert_manager: Arc, interceptor: Arc>, @@ -211,16 +336,78 @@ impl ProxyServer { ) -> Result<()> { // Parse domain and port let (domain, port) = Self::parse_host_port(uri)?; - // Send 200 Connection Established to client - client_stream - .write_all(b"HTTP/1.1 200 Connection Established\r\n\r\n") - .await?; - client_stream - .flush() // ← 关键：确保 HTTP 响应完全发送 - .await - .map_err(Error::Io)?; + // Perform TLS accept + MITM handling, send HTTP 200 before TLS handshake + Self::accept_tls_and_handle( + client_stream, + &domain, + port, + true, + cert_manager, + interceptor, + slinger_client, + ) + .await + } + + /// Handle HTTPS tunnel without interception (transparent proxy) + async fn handle_https_tunnel(client_stream: TcpStream, uri: &str) -> Result<()> { + Self::tcp_tunnel(client_stream, uri, true).await + } + + /// Handle TCP tunnel without interception (for SOCKS5) + /// This function doesn't send any response - the SOCKS5 handshake already sent the reply + async fn handle_tcp_tunnel(client_stream: TcpStream, target_addr: &str) -> Result<()> { + Self::tcp_tunnel(client_stream, target_addr, false).await + } + + /// Handle HTTPS CONNECT with MITM interception for SOCKS5 + /// This function doesn't send HTTP response - the SOCKS5 handshake already sent the reply + async fn handle_https_connect_socks5( + client_stream: TcpStream, + uri: &str, + cert_manager: Arc, + interceptor: Arc>, + slinger_client: Client, + ) -> Result<()> { + // Parse domain and port + let (domain, port) = Self::parse_host_port(uri)?; + + // Perform TLS accept + MITM handling, do NOT send HTTP 200 (SOCKS5 already responded) + Self::accept_tls_and_handle( + client_stream, + &domain, + port, + false, + cert_manager, + interceptor, + slinger_client, + ) + .await + } + + /// Accept TLS on an incoming stream (using certs from CertificateManager) and handle HTTPS requests over it. + /// If `send_response` is true, send an HTTP/1.1 200 Connection Established before performing the TLS handshake + async fn accept_tls_and_handle( + mut client_stream: TcpStream, + domain: &str, + port: u16, + send_response: bool, + cert_manager: Arc, + interceptor: Arc>, + slinger_client: Client, + ) -> Result<()> { + if send_response { + client_stream + .write_all(b"HTTP/1.1 200 Connection Established\r\n\r\n") + .await?; + client_stream + .flush() // ensure response fully sent + .await + .map_err(Error::Io)?; + } + // Generate server certificate for this domain - let (cert_chain, key) = cert_manager.get_server_cert(&domain).await?; + let (cert_chain, key) = cert_manager.get_server_cert(domain).await?; // Create TLS acceptor let tls_config = Self::create_tls_server_config(cert_chain, key)?; let acceptor = TlsAcceptor::from(Arc::new(tls_config)); @@ -233,34 +420,9 @@ impl ProxyServer { Self::handle_https_stream(tls_stream, domain_with_port, interceptor, slinger_client).await } - /// Handle HTTPS tunnel without interception (transparent proxy) - async fn handle_https_tunnel(mut client_stream: TcpStream, uri: &str) -> Result<()> { + /// Generic TCP tunnel helper. If `send_response` is true, sends HTTP/1.1 200 before tunneling. + async fn tcp_tunnel(mut client_stream: TcpStream, uri: &str, send_response: bool) -> Result<()> { let (host, port) = Self::parse_host_port(uri)?; - let target_addr = format!("{}:{}", host, port); - - // Connect to target server - let mut target_stream = TcpStream::connect(&target_addr).await.map_err(|e| { - Error::connection_error(format!("Failed to connect to {}: {}", target_addr, e)) - })?; - client_stream - .write_all(b"HTTP/1.1 200 Connection Established\r\n\r\n") - .await?; - let (mut client_read, mut client_write) = client_stream.split(); - let (mut target_read, mut target_write) = target_stream.split(); - let client_to_target = tokio::io::copy(&mut client_read, &mut target_write); - let target_to_client = tokio::io::copy(&mut target_read, &mut client_write); - tokio::select! { - _ = client_to_target => {}, - _ = target_to_client => {}, - } - Ok(()) - } - - /// Handle TCP tunnel without interception (for SOCKS5) - /// This function doesn't send any response - the SOCKS5 handshake already sent the reply - async fn handle_tcp_tunnel(mut client_stream: TcpStream, target_addr: &str) -> Result<()> { - // Parse host and port - let (host, port) = Self::parse_host_port(target_addr)?; let addr = format!("{}:{}", host, port); // Connect to target server @@ -268,7 +430,12 @@ impl ProxyServer { .await .map_err(|e| Error::connection_error(format!("Failed to connect to {}: {}", addr, e)))?; - // Bidirectional copy + if send_response { + client_stream + .write_all(b"HTTP/1.1 200 Connection Established\r\n\r\n") + .await?; + } + let (mut client_read, mut client_write) = client_stream.split(); let (mut target_read, mut target_write) = target_stream.split(); @@ -283,32 +450,42 @@ impl ProxyServer { Ok(()) } - /// Handle HTTPS CONNECT with MITM interception for SOCKS5 - /// This function doesn't send HTTP response - the SOCKS5 handshake already sent the reply - async fn handle_https_connect_socks5( - client_stream: TcpStream, - uri: &str, - cert_manager: Arc, + /// Forward a prepared `Request` through the inner `slinger::Client` and run interceptors. + /// Returns Some(response_bytes) if there is a response to write back to the caller, or None if + /// the interceptor chain dropped the request/response. + async fn forward_request_via_client( interceptor: Arc>, - slinger_client: Client, - ) -> Result<()> { - // Parse domain and port - let (domain, port) = Self::parse_host_port(uri)?; - - // Generate server certificate for this domain - let (cert_chain, key) = cert_manager.get_server_cert(&domain).await?; - - // Create TLS acceptor - let tls_config = Self::create_tls_server_config(cert_chain, key)?; - let acceptor = TlsAcceptor::from(Arc::new(tls_config)); - - // Perform TLS handshake with client - let tls_stream = acceptor - .accept(client_stream) - .await - .map_err(|e| Error::tls_error(format!("TLS handshake failed: {}", e)))?; - let domain_with_port = format!("{}:{}", domain, port); - Self::handle_https_stream(tls_stream, domain_with_port, interceptor, slinger_client).await + client: &Client, + request: Request, + ) -> Result>> { + let handler = interceptor.read().await; + if let Some(modified_req) = handler.process_request(request).await? { + let uri = modified_req.uri().clone(); + let method = modified_req.method().clone(); + let headers = modified_req.headers().clone(); + let body_data = if let Some(body) = modified_req.body() { + body.to_vec() + } else { + Vec::new() + }; + let mut req_builder = client.request(method, uri); + for (name, value) in headers.iter() { + req_builder = req_builder.header(name, value); + } + req_builder = req_builder.body(body_data); + match req_builder.send().await { + Ok(response) => { + if let Some(final_response) = handler.process_response(response).await? { + let response_bytes = Self::serialize_http_response(&final_response)?; + return Ok(Some(response_bytes)); + } + } + Err(_e) => { + return Ok(Some(b"HTTP/1.1 502 Bad Gateway\r\n\r\n".to_vec())); + } + } + } + Ok(None) } /// Handle HTTPS requests over TLS connection @@ -346,34 +523,10 @@ impl ProxyServer { // Parse request if let Ok(request) = Self::parse_http_request(&buffer, &domain) { - // Process through interceptors - let handler = interceptor.read().await; - if let Some(modified_req) = handler.process_request(request).await? { - let uri = modified_req.uri().clone(); - let method = modified_req.method().clone(); - let headers = modified_req.headers().clone(); - let body_data = if let Some(body) = modified_req.body() { - body.to_vec() - } else { - Vec::new() - }; - let mut req_builder = client.request(method, uri); - for (name, value) in headers.iter() { - req_builder = req_builder.header(name, value); - } - req_builder = req_builder.body(body_data); - match req_builder.send().await { - Ok(response) => { - if let Some(final_response) = handler.process_response(response).await? { - let response_bytes = Self::serialize_http_response(&final_response)?; - tls_stream.write_all(&response_bytes).await?; - } - } - Err(_e) => { - let error_response = b"HTTP/1.1 502 Bad Gateway\r\n\r\n"; - tls_stream.write_all(error_response).await?; - } - } + if let Some(response_bytes) = + Self::forward_request_via_client(interceptor, &client, request).await? + { + tls_stream.write_all(&response_bytes).await?; } } @@ -428,40 +581,12 @@ impl ProxyServer { let http_request = request_builder.body(Bytes::new())?; let request: Request = http_request.into(); - // Process through interceptors - let handler = interceptor.read().await; - if let Some(modified_req) = handler.process_request(request).await? { - // Forward using slinger - let uri = modified_req.uri().clone(); - let method = modified_req.method().clone(); - let headers = modified_req.headers().clone(); - let body_data = if let Some(body) = modified_req.body() { - body.to_vec() - } else { - Vec::new() - }; - - let mut req_builder = client.request(method, uri); - for (name, value) in headers.iter() { - req_builder = req_builder.header(name, value); - } - req_builder = req_builder.body(body_data); - - match req_builder.send().await { - Ok(response) => { - if let Some(final_response) = handler.process_response(response).await? { - let response_bytes = Self::serialize_http_response(&final_response)?; - let mut stream = reader.into_inner(); - stream.write_all(&response_bytes).await?; - } - } - Err(_e) => { - let mut stream = reader.into_inner(); - stream - .write_all(b"HTTP/1.1 502 Bad Gateway\r\n\r\n") - .await?; - } - } + // Process through interceptors and forward + if let Some(response_bytes) = + Self::forward_request_via_client(interceptor, &client, request).await? + { + let mut stream = reader.into_inner(); + stream.write_all(&response_bytes).await?; } Ok(()) diff --git a/slinger/examples/custom_tls.rs b/slinger/examples/custom_tls.rs index 6b37b2a..40ee967 100644 --- a/slinger/examples/custom_tls.rs +++ b/slinger/examples/custom_tls.rs @@ -10,15 +10,17 @@ //! //! Note: This example demonstrates the API with a mock TLS implementation for educational purposes. +#[cfg(feature = "tls")] use slinger::tls::CustomTlsStream; #[cfg(feature = "tls")] use slinger::tls::{CustomTlsConnector, PeerCertificate}; #[cfg(feature = "tls")] use slinger::{ConnectorBuilder, Result, Socket, StreamWrapper}; -use std::pin::Pin; #[cfg(feature = "tls")] use std::sync::Arc; +#[cfg(feature = "tls")] use std::task::{Context, Poll}; +#[cfg(feature = "tls")] use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; #[cfg(feature = "tls")]