use std::{ process::exit, sync::Arc, thread, time::{SystemTime, UNIX_EPOCH}, }; use http_body_util::{BodyExt, Full}; use hyper::{ Request, Response, body::{Buf, Bytes, Incoming}, server::conn::http1, }; use hyper_util::rt::TokioIo; use owo_colors::OwoColorize; use serde::{Deserialize, Serialize}; use tokio::{net::TcpListener, sync::Mutex}; use crate::{ replay::start_replay_server, store::{LogStore, save_logs_to_file}, }; pub const PROXY_LOG_FILE: &str = "replay_mocks.json"; #[derive(Debug, Serialize, Deserialize, PartialEq, Clone)] pub struct RequestLog { pub timestamp: u64, pub method: String, pub path: String, pub query_params: Option, pub headers: Vec<(String, String)>, pub body: Option, } #[derive(Debug, Serialize, Deserialize, PartialEq, Clone)] pub struct ResponseLog { pub status: u16, pub headers: Vec<(String, String)>, pub body: Option, } #[derive(Debug, Serialize, Deserialize, PartialEq, Clone)] pub struct ProxyLog { pub request: RequestLog, pub response: ResponseLog, } pub async fn start_server( target: &str, listen: &str, ) -> Result<(), Box> { let target_uri = target.parse::()?; let target_authority = target_uri.authority().ok_or("Invalid target URL")?; let target_scheme = target_uri.scheme_str().ok_or("http")?; let target_host = target_authority.host(); let target_port = target_authority .port_u16() .unwrap_or(if target_scheme == "https" { 443 } else { 80 }); let logs = Arc::new(Mutex::new(Vec::::new())); let logs_for_saving = logs.clone(); tokio::spawn(async move { loop { tokio::time::sleep(tokio::time::Duration::from_secs(5)).await; save_logs_to_file(&logs_for_saving, PROXY_LOG_FILE) .await .unwrap_or_else(|e| eprintln!("Error saving logs to file: {}", e)); } }); let logs_for_replay = logs.clone(); thread::spawn(move || { let rt = tokio::runtime::Builder::new_multi_thread() .enable_all() .build() .unwrap(); rt.block_on(async { match start_replay_server(logs_for_replay, "127.0.0.1:6688").await { Ok(_) => { println!("Replay server stopped"); exit(0); } Err(e) => eprintln!("Replay server error: {}", e), } }); }); let listener = TcpListener::bind(listen).await?; println!("Target URL: {}", target.magenta()); println!("Proxy server is listening on {}", listen.green()); println!("Replay server is running on {}", "127.0.0.1:6688".green()); loop { let (stream, _) = listener.accept().await?; let io = TokioIo::new(stream); let target_host_str = target_host.to_string(); let target_scheme = target_scheme.to_string(); let logs_clone = logs.clone(); tokio::task::spawn(async move { let service = hyper::service::service_fn(move |req: Request| { let target_host = target_host_str.clone(); let scheme = target_scheme.clone(); let logs = logs_clone.clone(); async move { proxy_handler(req, &target_host, target_port, &scheme, logs).await } }); if let Err(err) = http1::Builder::new() .keep_alive(false) .max_buf_size(30 * 1024 * 1024) .serve_connection(io, service) .await { eprintln!("> Connection error: {}", err); } }); } } pub async fn proxy_handler( req: Request, target_host: &str, target_port: u16, scheme: &str, logs: LogStore, ) -> Result>, hyper::Error> { let timestamp = SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap() .as_secs(); let method = req.method().clone(); let path = req.uri().path().to_string(); let query = req.uri().query().map(|q| q.to_string()); let headers: Vec<(String, String)> = req .headers() .iter() .map(|(name, value)| (name.to_string(), value.to_str().unwrap_or("").to_string())) .collect(); let (parts, body) = req.into_parts(); let body_bytes = match body.collect().await { Ok(collected) => collected.aggregate(), Err(e) => { eprintln!("Error collecting request body: {}", e); return Ok(Response::builder() .status(500) .body(Full::new(Bytes::from("Internal Server Error"))) .unwrap()); } }; let body_vec = body_bytes.chunk().to_vec(); let body_str = String::from_utf8(body_vec.clone()).ok(); let forward_uri = if target_port != 443 && target_port != 80 { format!( "{}://{}:{}{}{}", scheme, target_host, target_port, parts.uri.path(), parts .uri .query() .map_or(String::new(), |q| format!("?{}", q)) ) } else { format!( "{}://{}{}{}", scheme, target_host, parts.uri.path(), parts .uri .query() .map_or(String::new(), |q| format!("?{}", q)) ) }; println!("{} {} {}", method.yellow(), path, forward_uri.magenta()); let client = reqwest::Client::builder() .timeout(std::time::Duration::from_secs(30)) .danger_accept_invalid_certs(true) .build() .unwrap_or_else(|_| reqwest::Client::new()); let mut req_builder = match method.as_str() { "GET" => client.get(&forward_uri), "POST" => client.post(&forward_uri), "PUT" => client.put(&forward_uri), "DELETE" => client.delete(&forward_uri), "HEAD" => client.head(&forward_uri), "OPTIONS" => client.request(reqwest::Method::OPTIONS, &forward_uri), "PATCH" => client.patch(&forward_uri), _ => { eprintln!("Unsupported method: {}", method); return Ok(Response::builder() .status(400) .body(Full::new(Bytes::from("Bad Request: Unsupported Method"))) .unwrap()); } }; for (name, value) in &headers { if name.to_lowercase() != "host" && name.to_lowercase() != "connection" { if let Ok(header_name) = reqwest::header::HeaderName::from_bytes(name.as_bytes()) { if let Ok(header_value) = reqwest::header::HeaderValue::from_str(value) { req_builder = req_builder.header(header_name, header_value); } } } } if !body_vec.is_empty() { req_builder = req_builder.body(body_vec.clone()); } let resp = match req_builder.send().await { Ok(resp) => resp, Err(e) => { eprintln!("Error sending request: {}", e); return Ok(Response::builder() .status(502) .body(Full::new(Bytes::from(format!("Bad Gateway: {}", e)))) .unwrap()); } }; let status = resp.status().as_u16(); let resp_headers: Vec<(String, String)> = resp .headers() .iter() .map(|(name, value)| (name.to_string(), value.to_str().unwrap_or("").to_string())) .collect(); let resp_bytes = match resp.bytes().await { Ok(bytes) => bytes, Err(e) => { eprintln!("Error reading response body: {}", e); return Ok(Response::builder() .status(500) .body(Full::new(Bytes::from("Internal Server Error"))) .unwrap()); } }; let resp_vec = resp_bytes.to_vec(); let resp_str = String::from_utf8(resp_vec.clone()).ok(); let log_entry = ProxyLog { request: RequestLog { timestamp, method: method.to_string(), path, query_params: query, headers, body: body_str, }, response: ResponseLog { status, headers: resp_headers.clone(), body: resp_str.clone(), }, }; { let mut logs_guard = logs.lock().await; if !logs_guard.iter().any(|log| { log.request.method == log_entry.request.method && log.request.path == log_entry.request.path && log.request.query_params == log_entry.request.query_params }) { logs_guard.push(log_entry.clone()); } } println!( "Saved request/response to log store {}", PROXY_LOG_FILE.magenta() ); let mut builder = Response::builder().status(status); for (name, value) in resp_headers { if name.to_lowercase() != "connection" && name.to_lowercase() != "transfer-encoding" { if let Ok(header_name) = hyper::header::HeaderName::from_bytes(name.as_bytes()) { if let Ok(header_value) = hyper::header::HeaderValue::from_str(&value) { builder = builder.header(header_name, header_value); } } } } builder = builder.header("content-length", resp_vec.len()); builder = builder.header("connection", "close"); Ok(builder .body(Full::new(Bytes::from(resp_vec))) .unwrap_or_else(|_| { Response::builder() .status(500) .body(Full::new(Bytes::from("Internal Server Error"))) .unwrap() })) }