Dump UNIX sockets to standard output

unixdump.rs

use std::env;
use std::fs;
use std::io::{self, Read, Write};
use std::os::unix::fs::PermissionsExt;
use std::os::unix::net::{UnixListener, UnixStream};
use std::path::Path;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::thread;

fn main() {
    let args: Vec<String> = env::args().collect();
    if args.len() != 2 {
        eprintln!("Usage: {} <socket-path>", args[0]);
        std::process::exit(1);
    }

    let socket_path = &args[1];
    let orig_socket_path = format!("{}.orig", socket_path);

    // Check if the socket exists
    if !Path::new(socket_path).exists() {
        eprintln!("Socket {} does not exist", socket_path);
        std::process::exit(1);
    }

    // Get original socket permissions and ownership
    let metadata = match fs::metadata(socket_path) {
        Ok(m) => m,
        Err(e) => {
            eprintln!("Failed to get socket metadata: {}", e);
            std::process::exit(1);
        }
    };
    let mode = metadata.permissions().mode();

    // Get owner/group (Unix specific)
    #[cfg(unix)]
    let (uid, gid) = {
        use std::os::unix::fs::MetadataExt;
        (metadata.uid(), metadata.gid())
    };

    // Rename the original socket
    if let Err(e) = fs::rename(socket_path, &orig_socket_path) {
        eprintln!("Failed to rename socket: {}", e);
        std::process::exit(1);
    }

    // Setup cleanup function
    let socket_path_clone = socket_path.to_string();
    let orig_socket_path_clone = orig_socket_path.clone();
    let cleanup = move || {
        eprintln!("Shutting down, restoring original socket...");
        let _ = fs::remove_file(&socket_path_clone);
        if Path::new(&orig_socket_path_clone).exists() {
            if let Err(e) = fs::rename(&orig_socket_path_clone, &socket_path_clone) {
                eprintln!("Warning: Failed to restore original socket: {}", e);
            } else {
                eprintln!("Original socket restored to {}", socket_path_clone);
            }
        }
    };

    // Create new listener on the original path
    let listener = match UnixListener::bind(socket_path) {
        Ok(l) => l,
        Err(e) => {
            cleanup();
            eprintln!("Failed to create listener: {}", e);
            std::process::exit(1);
        }
    };

    // Set permissions on the new socket
    if let Err(e) = fs::set_permissions(socket_path, fs::Permissions::from_mode(mode)) {
        eprintln!("Warning: Failed to set socket permissions: {}", e);
    }

    // Set ownership on the new socket
    #[cfg(unix)]
    {
        use std::os::unix::fs::chown;
        if let Err(e) = chown(socket_path, Some(uid), Some(gid)) {
            eprintln!("Warning: Failed to set socket ownership: {}", e);
        }
    }

    eprintln!(
        "Listening on {}, forwarding to {}",
        socket_path, orig_socket_path
    );
    eprintln!("Press Ctrl+C to stop and restore the original socket");

    // Setup signal handler
    let shutdown_flag = Arc::new(AtomicBool::new(false));
    let shutdown_flag_clone = shutdown_flag.clone();

    ctrlc::set_handler(move || {
        eprintln!("\nReceived interrupt signal");
        shutdown_flag_clone.store(true, Ordering::SeqCst);
    })
    .expect("Error setting Ctrl-C handler");

    let mut conn_id = 0;
    for stream in listener.incoming() {
        if shutdown_flag.load(Ordering::SeqCst) {
            break;
        }

        match stream {
            Ok(client) => {
                // Set non-blocking mode to false for proper handling
                let _ = client.set_nonblocking(false);

                let orig_path = orig_socket_path.clone();
                let id = conn_id;
                conn_id += 1;

                thread::spawn(move || {
                    handle_connection(client, &orig_path, id);
                });
            }
            Err(e) => {
                if shutdown_flag.load(Ordering::SeqCst) {
                    break;
                }
                eprintln!("Accept error: {}", e);
            }
        }
    }

    cleanup();
}

fn handle_connection(mut client: UnixStream, orig_socket_path: &str, conn_id: usize) {
    eprintln!("[Connection #{}] New connection", conn_id);

    // Connect to the original socket
    let mut server = match UnixStream::connect(orig_socket_path) {
        Ok(s) => s,
        Err(e) => {
            eprintln!(
                "[Connection #{}] Failed to connect to original socket: {}",
                conn_id, e
            );
            return;
        }
    };

    // Clone streams for bidirectional communication
    let mut client_read = match client.try_clone() {
        Ok(c) => c,
        Err(e) => {
            eprintln!(
                "[Connection #{}] Failed to clone client stream: {}",
                conn_id, e
            );
            return;
        }
    };
    let mut server_read = match server.try_clone() {
        Ok(s) => s,
        Err(e) => {
            eprintln!(
                "[Connection #{}] Failed to clone server stream: {}",
                conn_id, e
            );
            return;
        }
    };

    let conn_id_clone = conn_id;

    // Client -> Server
    let handle1 = thread::spawn(move || {
        copy_and_log(&mut client_read, &mut server, conn_id, "CLIENT->SERVER");
    });

    // Server -> Client
    let handle2 = thread::spawn(move || {
        copy_and_log(
            &mut server_read,
            &mut client,
            conn_id_clone,
            "SERVER->CLIENT",
        );
    });

    let _ = handle1.join();
    let _ = handle2.join();

    eprintln!("[Connection #{}] Connection closed", conn_id);
}

fn copy_and_log(src: &mut UnixStream, dst: &mut UnixStream, conn_id: usize, direction: &str) {
    let mut buffer = [0u8; 32 * 1024];
    let mut stdout = io::stdout();

    loop {
        match src.read(&mut buffer) {
            Ok(0) => break, // EOF
            Ok(n) => {
                // Write to stdout
                if let Err(e) = stdout.write_all(&buffer[..n]) {
                    eprintln!(
                        "[Connection #{}] {} stdout write error: {}",
                        conn_id, direction, e
                    );
                    break;
                }
                let _ = stdout.flush();

                // Forward to destination
                if let Err(e) = dst.write_all(&buffer[..n]) {
                    eprintln!("[Connection #{}] {} write error: {}", conn_id, direction, e);
                    break;
                }
            }
            Err(e) => {
                eprintln!("[Connection #{}] {} read error: {}", conn_id, direction, e);
                break;
            }
        }
    }
}