diff --git a/src/main.rs b/src/main.rs index 8b55279..9d7faae 100644 --- a/src/main.rs +++ b/src/main.rs @@ -32,6 +32,10 @@ struct Args { /// The path on the server this server will save requests and responses #[arg(short, long, default_value = "storage")] storage_path: String, + + /// Only forward requests that match a given prefix + #[arg(short, long)] + prefix: Option, } /// Get the current time since epoch @@ -62,7 +66,10 @@ async fn main() -> Result<(), Box> { let (mut client_socket, _) = listener.accept().await?; tokio::spawn(async move { - log::info!("Start new connection from {}", client_socket.peer_addr().unwrap()); + log::info!( + "Start new connection from {}", + client_socket.peer_addr().unwrap() + ); let args = Args::parse(); @@ -126,7 +133,17 @@ async fn main() -> Result<(), Box> { } // We need to modify some headers (if not done already) to adapt the request to the server - let buff = if !modified_headers { + let buff = if !modified_headers { + + // Check for URL prefix + if let Some(prefix) = &args.prefix { + if !String::from_utf8_lossy(&buf_client[..count]).split_once('\n').map(|l|l.0).unwrap_or("").contains(&format!(" {prefix}")) { + client_write.write_all(b"HTTP/1.1 401 Forbidden\r\n\r\nNot proxifiable.\r\n").await.expect("Failed to respond to client"); + client_write.flush().await.expect("Failed to flush response to client!"); + return; + } + } + modified_headers = true; manipulate_headers(&buf_client[..count], &args.upstream_dns) } @@ -165,8 +182,6 @@ async fn main() -> Result<(), Box> { } fn manipulate_headers(buff: &[u8], host: &str) -> Vec { - // return buff.to_vec(); - let mut out = Vec::with_capacity(buff.len()); let mut i = 0;