feat: add support for PROXY protocol v2

Cargo.toml

@@ -8,20 +8,59 @@ keywords = ["tcp", "ctf", "socat", "fcsc"]
 categories = ["command-line-utilities"]
 repository = "https://github.com/erdnaxe/sossette"
 authors = ["erdnaxe <erdnaxe@crans.org>"]
-license = "MIT"
-version = "0.1.1"
+version = "0.2.0"
 edition = "2024"
 
 [dependencies]
 anyhow = "1.0"
-clap = { version = "4.5", features = ["derive", "env"] }
+clap = { version = "4.6", features = ["derive", "env"] }
 clap-verbosity-flag = "3.0"
 command-group = { version = "5.0", features = ["with-tokio"] }
 env_logger = "0.11"
 log = "0.4"
 rand = "0.10"
 sha2 = "0.10"
-tokio = { version = "1.50", features = ["rt-multi-thread", "io-util", "signal", "net", "time", "process", "macros"] }
+tokio = { version = "1.5", features = [
+    "rt-multi-thread",
+    "io-util",
+    "signal",
+    "net",
+    "time",
+    "process",
+    "macros",
+] }
+
+[lints.rust]
+arithmetic_overflow = { level = "deny", priority = -1 }
+
+[lints.clippy]
+pedantic = { level = "deny", priority = -1 }
+nursery = { level = "deny", priority = -1 }
+missing-errors-doc = "allow"
+
+indexing_slicing = { level = "deny", priority = -1 }
+fallible_impl_from = { level = "deny", priority = -1 }
+wildcard_enum_match_arm = { level = "deny", priority = -1 }
+unneeded_field_pattern = { level = "deny", priority = -1 }
+fn_params_excessive_bools = { level = "deny", priority = -1 }
+must_use_candidate = { level = "deny", priority = -1 }
+checked_conversions = { level = "deny", priority = -1 }
+cast_possible_truncation = { level = "deny", priority = -1 }
+cast_sign_loss = { level = "deny", priority = -1 }
+cast_possible_wrap = { level = "deny", priority = -1 }
+cast_precision_loss = { level = "deny", priority = -1 }
+integer_division = { level = "deny", priority = -1 }
+arithmetic_side_effects = { level = "deny", priority = -1 }
+unchecked_duration_subtraction = { level = "deny", priority = -1 }
+unwrap_used = "warn"
+expect_used = "warn"
+panicking_unwrap = { level = "deny", priority = -1 }
+option_env_unwrap = { level = "deny", priority = -1 }
+join_absolute_paths = { level = "deny", priority = -1 }
+serde_api_misuse = { level = "deny", priority = -1 }
+uninit_vec = { level = "deny", priority = -1 }
+transmute_ptr_to_ref = { level = "deny", priority = -1 }
+transmute_undefined_repr = { level = "deny", priority = -1 }
 
 [profile.release]
 codegen-units = 1

README.md

@@ -66,6 +66,71 @@ world
 ^C
 ```
 
+## PROXY protocol support
+
+Sossette supports the [PROXY protocol v2](https://github.com/haproxy/haproxy/blob/master/doc/proxy-protocol.txt) to preserve client IP addresses when running behind a load balancer or reverse proxy.
+
+### Usage
+
+Enable PROXY protocol v2 with the `--proxy-protocol` flag. When enabled, a valid PROXY protocol v2 header is **required** and connections without one are rejected:
+
+```bash
+$ sossette --proxy-protocol -l 0.0.0.0:4000 cat
+```
+
+Or using the environment variable:
+
+```bash
+$ WRAPPER_PROXY_PROTOCOL=true sossette -l 0.0.0.0:4000 cat
+```
+
+### Accessing client information
+
+When PROXY protocol is enabled and a valid header is received, sossette:
+
+1. **Logs the real client IP** instead of the proxy's IP:
+   ```
+   [2024-03-09T10:15:23Z INFO  sossette] Client [::1]:55438 connected
+   [2024-03-09T10:15:23Z INFO  sossette] Real client: 192.0.2.123:54321 (via proxy [::1]:55438)
+   ```
+
+### Load balancer configuration
+
+#### HAProxy
+
+Configure HAProxy to send PROXY protocol v2 headers:
+
+```haproxy
+frontend tcp_front
+    bind *:443
+    mode tcp
+    default_backend tcp_back
+
+backend tcp_back
+    mode tcp
+    server sossette 127.0.0.1:4000 send-proxy-v2
+```
+
+#### nginx
+
+Configure nginx stream module with PROXY protocol:
+
+```nginx
+stream {
+    upstream sossette {
+        server 127.0.0.1:4000;
+    }
+
+    server {
+        listen 443;
+        proxy_pass sossette;
+        proxy_protocol on;
+    }
+}
+```
+
+**Security note**: When using PROXY protocol, ensure that only trusted load balancers can connect to sossette (e.g., using firewall rules). Otherwise, clients could spoof their IP addresses by sending fake PROXY protocol headers.
+
 ## Applying transformations to stdin
 
 `process_stdin` in [src/main.rs](./src/main.rs) can be easily patched to apply

src/handler.rs

@@ -1,15 +1,17 @@
 // SPDX-FileCopyrightText: 2023-2025 erdnaxe
 // SPDX-License-Identifier: MIT
 
-use crate::pow;
 use crate::Args;
+use crate::pow;
+use crate::proxy;
 
+use std::net::SocketAddr;
 use std::process::Stdio;
 use std::time::Duration;
 
 use anyhow::{Context, Result};
 use command_group::AsyncCommandGroup;
-use log::debug;
+use log::{debug, info, warn};
 use tokio::io::{AsyncReadExt, AsyncWriteExt};
 use tokio::net::TcpStream;
 use tokio::process::Command;
@@ -21,19 +23,16 @@ async fn process_stdin<R: AsyncReadExt + Unpin, W: AsyncWriteExt + Unpin>(
     mut socket: R,
     mut child_stdin: W,
 ) -> Result<()> {
-    let mut in_buf = [0; 1024];
+    let mut in_buf = [0u8; 1024];
     loop {
         let n = socket.read(&mut in_buf).await?;
         if n == 0 {
             return Ok(()); // socket closed
         }
-        if in_buf[0] == 3 {
-            debug!("Client sent Ctrl-C");
-            return Ok(());
-        }
-        debug!("Writting to stdin: {:?}", &in_buf[0..n]);
+        let data = in_buf.get(..n).context("stdin read index out of bounds")?;
+        debug!("Writting to stdin: {data:?}");
         child_stdin
-            .write_all(&in_buf[0..n])
+            .write_all(data)
             .await
             .context("Failed to write to stdin")?;
     }
@@ -44,14 +43,18 @@ async fn process_stdout<R: AsyncReadExt + Unpin, W: AsyncWriteExt + Unpin>(
     mut socket: W,
     mut child_stdout: R,
 ) -> Result<()> {
-    let mut out_buf = [0; 1024];
+    let mut out_buf = [0u8; 1024];
     loop {
         let n = child_stdout.read(&mut out_buf).await?;
         if n == 0 {
             return Ok(()); // process closed
         }
+        let data = out_buf
+            .get(..n)
+            .context("stdout read index out of bounds")?;
+        debug!("Reading from stdout: {data:?}");
         socket
-            .write_all(&out_buf[0..n])
+            .write_all(data)
             .await
             .context("Failed to write to socket")?;
     }
@@ -61,49 +64,94 @@ async fn process_stdout<R: AsyncReadExt + Unpin, W: AsyncWriteExt + Unpin>(
 ///
 /// Spawn one process and then spawn 3 tasks to manage input, output and
 /// timeout. If one of these tasks reach its end, kill the process.
-pub async fn handle_client(mut socket: TcpStream, args: Args) -> Result<()> {
-    // Send message of the day
+pub async fn handle_client(
+    mut socket: TcpStream,
+    peer_addr: SocketAddr,
+    args: Args,
+) -> Result<Option<proxy::ProxyInfo>> {
+    // Parse PROXY protocol header if enabled
+    let proxy_info = if args.proxy_protocol {
+        match proxy::parse_proxy_v2_header(&mut socket).await {
+            Ok(proxy::ProxyHeader::Proxied(info)) => {
+                info!(
+                    "Client: {}:{} -> {}:{} (via proxy {}) connected",
+                    info.src_addr, info.src_port, info.dst_addr, info.dst_port, peer_addr
+                );
+                Some(info)
+            }
+            Ok(proxy::ProxyHeader::Local) => {
+                debug!("PROXY protocol LOCAL command");
+                None
+            }
+            Err(e) => {
+                warn!("Rejecting connection from {peer_addr} due to PROXY protocol error: {e:?}");
+                return Err(e);
+            }
+        }
+    } else {
+        None
+    };
+
+    // MOTD
     if let Some(motd) = &args.motd {
         socket.write_all(motd.as_bytes()).await?;
-        socket.write_all(b"\r\n").await?;
+        socket.write_all(&b"\r\n"[..]).await?;
     }
 
-    // Proof-of-work prompt
+    // Proof-of-work
     if args.pow > 0 {
-        let valid = pow::proof_of_work_prompt(&mut socket, args.pow, args.pow_backdoor).await?;
+        let valid =
+            pow::proof_of_work_prompt(&mut socket, args.pow, args.pow_backdoor.as_ref()).await?;
         if !valid {
-            return Ok(());
+            return Ok(proxy_info);
         }
     }
 
     // Start command
     let mut command = Command::new(&args.command);
     command.args(&args.arguments);
     command.stdin(Stdio::piped()).stdout(Stdio::piped());
+
     let mut child = command.group_spawn().context("Failed to run command")?;
+
     let child_stdin = child.inner().stdin.take().context("Failed to open stdin")?;
     let child_stdout = child
         .inner()
         .stdout
         .take()
         .context("Failed to open stdout")?;
 
-    // Start tasks
-    let mut set = JoinSet::new();
+    // Split socket
     let (read_half, write_half) = socket.into_split();
+
+    let mut set = JoinSet::new();
+
     set.spawn(async move { process_stdin(read_half, child_stdin).await });
+
     set.spawn(async move { process_stdout(write_half, child_stdout).await });
-    if let Some(timeout) = args.timeout {
+
+    let session_timeout = args.timeout.map(Duration::from_secs);
+
+    if let Some(timeout) = session_timeout {
         set.spawn(async move {
-            sleep(Duration::from_secs(timeout)).await;
+            sleep(timeout).await;
             debug!("Timeout reached");
             Ok(())
         });
     }
 
-    // If one task exits, drop the others
-    // Child group should always be killed before dropping child handle.
+    // Wait for first task to finish
     let res = set.join_next().await;
+
+    // Cancel remaining tasks immediately
+    set.abort_all();
+
+    // Kill the process group
     child.kill().await.context("Failed to kill process group")?;
-    res.unwrap_or(Ok(Ok(())))?
+
+    // Await child to avoid zombie process
+    let _ = child.wait().await;
+
+    res.unwrap_or(Ok(Ok(())))??;
+    Ok(proxy_info)
 }

src/main.rs

@@ -3,6 +3,7 @@
 
 mod handler;
 mod pow;
+mod proxy;
 
 use anyhow::{Context, Result};
 use clap::Parser;
@@ -35,6 +36,10 @@ struct Args {
     #[arg(long, value_name = "STRING", env = "WRAPPER_POW_BACKDOOR")]
     pow_backdoor: Option<String>,
 
+    /// Require PROXY protocol v2 header, reject connections without it
+    #[arg(long, env = "WRAPPER_PROXY_PROTOCOL")]
+    proxy_protocol: bool,
+
     #[command(flatten)]
     verbose: Verbosity<InfoLevel>,
 
@@ -56,16 +61,24 @@ async fn serve(args: Args) -> Result<()> {
         match listener.accept().await {
             Ok((socket, peer_addr)) => {
                 info!("Client {peer_addr:?} connected");
+                info!("Client {peer_addr:?} connected");
 
                 // Spawn task to handle this client
                 let my_args = args.clone();
                 tokio::spawn(async move {
-                    match handler::handle_client(socket, my_args).await {
-                        Ok(()) => {
+                    match handler::handle_client(socket, peer_addr, my_args).await {
+                        Ok(Some(proxy_info)) => {
+                            info!(
+                                "Client: {}:{} (via proxy {}) disconnected",
+                                proxy_info.src_addr, proxy_info.src_port, peer_addr
+                            );
+                        }
+                        Ok(None) => {
                             info!("Client {peer_addr:?} disconnected");
                         }
                         Err(e) => {
                             warn!("Handling client {peer_addr:?} failed: {e:?}");
+                            warn!("Handling client {peer_addr:?} failed: {e:?}");
                         }
                     }
                 });
@@ -95,6 +108,7 @@ async fn main() {
         Ok(()) => {}
         Err(err) => {
             warn!("Unable to listen for shutdown signal: {err}");
+            warn!("Unable to listen for shutdown signal: {err}");
         }
     }
 }
@@ -104,6 +118,6 @@ mod tests {
     #[test]
     fn verify_cli() {
         use clap::CommandFactory;
-        crate::Args::command().debug_assert()
+        crate::Args::command().debug_assert();
     }
 }