Enforce prefix

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<String>,
 }
 
 /// Get the current time since epoch
@@ -62,7 +66,10 @@ async fn main() -> Result<(), Box<dyn Error>> {
         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<dyn Error>> {
                         }
 
                         // 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<dyn Error>> {
 }
 
 fn manipulate_headers(buff: &[u8], host: &str) -> Vec<u8> {
-    // return buff.to_vec();
-
     let mut out = Vec::with_capacity(buff.len());
 
     let mut i = 0;