Update Rust crate tokio to 1.50.0

Cargo.toml

@@ -1,14 +1,14 @@
 [package]
-name = "proxy_saver"
+name = "header_proxy"
 version = "0.1.0"
 edition = "2024"
 
 [dependencies]
-log = "0.4.28"
-env_logger = "0.11.8"
-clap = { version = "4.5.50", features = ["env", "derive"] }
-tokio = { version = "1.48.0", features = ["full"] }
-tokio-rustls = { version = "0.26.4", features = ["ring", "tls12"], default-features = false }
-rustls-pki-types = "1.12.0"
-webpki-roots = "1.0.3"
-rand = "0.9.2"
+log = "0.4.29"
+env_logger = "0.11.9"
+clap = { version = "4.5.60", features = ["env", "derive"] }
+tokio = { version = "1.50.0", features = ["full"] }
+rand = "0.10.0"
+lazy_static = "1.5.0"
+httparse = "1.10.1"
+anyhow = "1.0.102"

Dockerfile

@@ -0,0 +1,10 @@
+FROM debian:bookworm-slim
+
+RUN apt-get update \
+  && apt-get install -y libcurl4 \
+  && rm -rf /var/lib/apt/lists/*
+
+COPY header_proxy /usr/local/bin/header_proxy
+
+ENTRYPOINT ["/usr/local/bin/header_proxy"]
+

README.md

@@ -1,2 +1,2 @@
-# Proxy saver
-Proxify HTTP request to HTTPS upstream, saving requests and responses.
+# HTTP header proxy
+Proxify HTTP request to another HTTP upstream, optionally manipulating path.

build_docker_image.sh

@@ -0,0 +1,10 @@
+#!/bin/bash
+cargo build --release
+
+TEMP_DIR=$(mktemp -d)
+cp target/release/header_proxy "$TEMP_DIR"
+
+docker build -f Dockerfile "$TEMP_DIR" -t pierre42100/header_proxy
+
+rm -r $TEMP_DIR
+

src/main.rs

@@ -1,51 +1,38 @@
 use clap::Parser;
 use rand::distr::{Alphanumeric, SampleString};
-use rustls_pki_types::ServerName;
 use std::error::Error;
-use std::fs::OpenOptions;
-use std::io::Write;
-use std::path::Path;
-use std::sync::Arc;
-use std::time::{SystemTime, UNIX_EPOCH};
 use tokio::io::{AsyncReadExt, AsyncWriteExt};
 use tokio::net::TcpListener;
 use tokio::net::TcpStream;
-use tokio_rustls::TlsConnector;
-use tokio_rustls::rustls::{ClientConfig, RootCertStore};
 
-/// Simple program that proxify requests and save responses
+/// Simple program that redirect requests to a given server in request
 #[derive(Parser, Debug)]
 #[command(version, about, long_about = None)]
 struct Args {
     /// The address the server will listen to
-    #[arg(short, long, default_value = "0.0.0.0:8000")]
+    #[arg(short, long, env, default_value = "0.0.0.0:8000")]
     listen_address: String,
 
-    /// Upstream address this server will connect to
-    #[arg(short, long, default_value = "communiquons.org")]
-    upstream_dns: String,
+    /// The name of the header that contain target host and port
+    #[arg(short, long, env, default_value = "x-target-host")]
+    target_host_port_header: String,
 
-    /// Upstream address this server will connect to
-    #[arg(short('I'), long, default_value = "10.0.1.10:443")]
-    upstream_ip: String,
-
-    /// 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>,
+    /// Name of optional header that contains path to add to the request.
+    ///
+    /// If this value is defined, all clients packets are inspected in research for path to
+    /// manipulate
+    #[arg(short, long, env)]
+    path_prefix_header: Option<String>,
 }
 
-/// Get the current time since epoch
-pub fn time() -> u64 {
-    SystemTime::now()
-        .duration_since(UNIX_EPOCH)
-        .unwrap()
-        .as_secs()
+lazy_static::lazy_static! {
+    static ref ARGS: Args = {
+        Args::parse()
+    };
 }
 
+const ERR_NOT_PROXIFIALBE: &[u8; 44] = b"HTTP/1.1 400 Forbidden\r\n\r\nNot proxifiable.\r\n";
+
 pub fn rand_str(len: usize) -> String {
     Alphanumeric.sample_string(&mut rand::rng(), len)
 }
@@ -54,105 +41,105 @@ pub fn rand_str(len: usize) -> String {
 async fn main() -> Result<(), Box<dyn Error>> {
     env_logger::init_from_env(env_logger::Env::new().default_filter_or("info"));
 
-    let args = Args::parse();
-
-    log::info!("Will start to listen on {}", args.listen_address);
-    let listener = TcpListener::bind(&args.listen_address).await?;
-
-    std::fs::create_dir_all(Path::new(args.storage_path.as_str())).unwrap();
+    log::info!("Will start to listen on {}", ARGS.listen_address);
+    let listener = TcpListener::bind(&ARGS.listen_address).await?;
 
     loop {
         // Asynchronously wait for an inbound socket.
         let (mut client_socket, _) = listener.accept().await?;
 
         tokio::spawn(async move {
-            let conn_id = rand_str(10);
+            let conn_id = rand_str(5);
 
             log::info!(
-                "[{conn_id}] Start new connection from {}",
+                "[{conn_id}] Handle new connection from {}",
                 client_socket.peer_addr().unwrap()
             );
 
-            let args = Args::parse();
-
-            let base_file_name = format!(
-                "{}-{}-{}",
-                client_socket.peer_addr().unwrap().ip(),
-                time(),
-                conn_id
-            );
-
-            let mut req_file = OpenOptions::new()
-                .create(true)
-                .truncate(true)
-                .write(true)
-                .open(Path::new(&args.storage_path).join(format!("{base_file_name}-req")))
-                .expect("Failed to create req file");
-            let mut res_file = OpenOptions::new()
-                .create(true)
-                .truncate(true)
-                .write(true)
-                .open(Path::new(&args.storage_path).join(format!("{base_file_name}-res")))
-                .expect("Failed to create req file");
-
-            let mut root_cert_store = RootCertStore::empty();
-            root_cert_store.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
-            let config = ClientConfig::builder()
-                .with_root_certificates(root_cert_store)
-                .with_no_client_auth();
-            let connector = TlsConnector::from(Arc::new(config));
-            let dnsname = ServerName::try_from(args.upstream_dns.to_string()).unwrap();
-
-            let stream = TcpStream::connect(args.upstream_ip)
-                .await
-                .expect("Failed to connect to upstream");
-            let mut upstream = connector
-                .connect(dnsname, stream)
-                .await
-                .expect("Failed to establish TLS connection");
-
             let (mut client_read, mut client_write) = client_socket.split();
-
-            let mut buf_client = [0u8; 1024];
+            let mut buf_client = [0u8; 10000];
             let mut buf_server = [0u8; 1024];
 
-            let mut modified_headers = false;
+            // Perform first read operation manually to manipulate path and determine target server
+            let mut total = 0;
+            let headers_processed = loop {
+                let count = match client_read.read(&mut buf_client[total..]).await {
+                    Ok(count) => count,
+                    Err(e) => {
+                        log::error!(
+                            "[{conn_id}]  Failed to read initial data from client, closing connection! {e}"
+                        );
+                        return;
+                    }
+                };
 
+                if count == 0 {
+                    log::error!("[{conn_id}] read from client count is null, cannot continue!");
+                    let _ = client_write.write_all(ERR_NOT_PROXIFIALBE).await;
+                    return;
+                }
+
+                total += count;
+
+                match process_headers(&buf_client[..total]) {
+                    Ok(None) => {
+                        log::debug!("[{conn_id}] Insufficient amount of data, need to continue");
+                        continue;
+                    }
+                    Ok(Some(res)) => break res,
+                    Err(e) => {
+                        log::error!("[{conn_id}] failed to parse initial request headers! {e}");
+                        let _ = client_write.write_all(ERR_NOT_PROXIFIALBE).await;
+                        return;
+                    }
+                }
+            };
+
+            // Connect to upstream
+            let mut upstream = match TcpStream::connect(headers_processed.target_host).await {
+                Ok(upstream) => upstream,
+                Err(e) => {
+                    log::error!("[{conn_id}] Could not connect to upstream! {e}");
+                    let _ = client_write.write_all(ERR_NOT_PROXIFIALBE).await;
+                    return;
+                }
+            };
+
+            // Transfer modified headers to upstream
+            if let Err(e) = upstream.write_all(&headers_processed.buff).await {
+                log::error!("[{conn_id}] Could not forward initial bytes to upstream! {e}");
+                let _ = client_write.write_all(ERR_NOT_PROXIFIALBE).await;
+                return;
+            }
+
+            // Enter in loop to forward remaining data
             loop {
                 tokio::select! {
                     count = client_read.read(&mut buf_client) => {
                         let  count = match count{  Ok(count) => count, Err(e) => {
                             log::error!("[{conn_id}]  Failed to read data from client, closing connection! {e}");
-                            return;
+                            break;
                         }};
-                        log::info!("[{conn_id}]  Got a new client read {count} - {base_file_name}");
+                        log::debug!("[{conn_id}]  Got a new client read {count}");
 
                         if count == 0 {
                             log::warn!("[{conn_id}] infinite loop (client), closing connection");
                             drop(upstream);
-                            return;
+                            break;
                         }
 
-                        // We need to modify some headers (if not done already) to adapt the request to the server
-                        let buff = if !modified_headers {
-
-                            // Check for URL prefix
-                            if let Some(prefix) = &args.prefix
-                                && !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)
+                        // In case of connection reuse, we need to reanalyze data
+                        if ARGS.path_prefix_header.is_some() &&
+                            let Ok(Some(res))= process_headers(&buf_client[..count])
+                            && let Err(e) = upstream.write_all(&res.buff).await {
+                                log::error!("[{conn_id}] Failed to write to upstream! {e}");
+                                break;
                         }
-                        else {
-                            buf_client[..count].to_vec()
-                        };
 
-                        upstream.write_all(&buff).await.unwrap_or_else(|_| panic!("[{conn_id}] Failed to write to upstream"));
-                        req_file.write_all(&buff).unwrap_or_else(|_| panic!("[{conn_id}] Failed to write to req"));
+                        if let Err(e)=upstream.write_all(&buf_client[..count]).await {
+                            log::error!("[{conn_id}] Failed to write to upstream! {e}");
+                            break;
+                        }
                     }
 
                     count = upstream.read(&mut buf_server) => {
@@ -160,48 +147,88 @@ async fn main() -> Result<(), Box<dyn Error>> {
                             Ok(count) => count,
                             Err(e) => {
                                 log::error!("[{conn_id}] Failed to read from upstream! {e}");
-                                return;
+                                break;
                             }
                         };
 
                         if count == 0 {
                             log::warn!("[{conn_id}] infinite loop (upstream), closing connection");
                             drop(upstream);
-                            return;
+                            break;
                         }
 
-                        log::info!("[{conn_id}] Got a new upstream read {count} - {base_file_name}");
-                        client_write.write_all(&buf_server[..count]).await.expect("Failed to write to client");
-                        res_file.write_all(&buf_server[..count]).expect("Failed to write to res");
+                        log::debug!("[{conn_id}] Got a new upstream read {count}");
+                        if let Err(e) = client_write.write_all(&buf_server[..count]).await {
+                            log::error!("Failed to write to upstream! {e}");
+                            break;
+                        };
                     }
                 }
             }
+
+            log::info!("[{conn_id}] Connection finished.");
         });
     }
 }
 
-fn manipulate_headers(buff: &[u8], host: &str) -> Vec<u8> {
-    let mut out = Vec::with_capacity(buff.len());
+struct ProcessHeadersResult {
+    buff: Vec<u8>,
+    target_host: String,
+}
 
-    let mut i = 0;
-    while i < buff.len() {
-        if buff[i] != b'\n' || i + 1 == buff.len() || !buff[i + 1..].starts_with(b"Host:") {
-            out.push(buff[i]);
-            i += 1;
-            continue;
+fn process_headers(buff: &[u8]) -> anyhow::Result<Option<ProcessHeadersResult>> {
+    let mut headers = [httparse::EMPTY_HEADER; 64];
+    let mut req = httparse::Request::new(&mut headers);
+    let parsing_res = req.parse(buff)?;
+
+    let target_host = headers
+        .iter()
+        .find(|h| h.name.eq_ignore_ascii_case(&ARGS.target_host_port_header))
+        .map(|h| String::from_utf8_lossy(h.value));
+
+    log::debug!("Request headers: {:?}", headers);
+
+    let Some(target_host) = target_host else {
+        if parsing_res.is_partial() {
+            return Ok(None);
+        } else {
+            anyhow::bail!(
+                "Request is complete without header {}",
+                ARGS.target_host_port_header
+            )
         }
+    };
 
-        i += 1;
-        out.push(b'\n');
+    // Check if path needs to be prefixed
+    let prefix_path = if let Some(hname) = &ARGS.path_prefix_header {
+        headers
+            .iter()
+            .find(|h| h.name.eq_ignore_ascii_case(hname))
+            .map(|h| String::from_utf8_lossy(h.value).replace(['\n', '\r', '\t', ' '], ""))
+    } else {
+        None
+    };
 
-        out.append(&mut format!("Host: {host}").as_bytes().to_vec());
-
-        while buff[i] != b'\r' && buff[i] != b'\n' {
-            i += 1;
+    // Perform prefix injection
+    let mut buff = buff.to_vec();
+    if let Some(prefix) = prefix_path {
+        let pos = buff.iter().position(|c| c == &b' ');
+        log::debug!("Add path prefix to request {prefix}");
+        if let Some(pos) = pos {
+            for (num, c) in prefix.as_bytes().iter().enumerate() {
+                buff.insert(pos + 1 + num, *c);
+            }
+        } else {
+            log::warn!("Unable to inject prefix!");
         }
     }
 
-    out
+    log::trace!("Final request: {}", String::from_utf8_lossy(&buff));
+
+    Ok(Some(ProcessHeadersResult {
+        target_host: target_host.to_string(),
+        buff,
+    }))
 }
 
 #[cfg(test)]