coreutiles 0.1.0

Core utils in Rust
Documentation
use std::net::{ToSocketAddrs, UdpSocket};
use std::{io::{self, stdin, stdout, Read, Write}, net::{IpAddr, SocketAddr, TcpListener, TcpStream}};

use clap::{ArgAction, Parser, Subcommand};

#[derive(Parser)]
#[command(version, about)]
#[command(propagate_version = true)]
pub struct Config {
    #[clap(long, short, action, action=ArgAction::SetFalse)]
    /// Use the TCP protocol
    tcp: bool,

    #[clap(long, short, action)]
    /// Use the UDP protocol
    udp: bool,

    #[command(subcommand)]
    command: Op,
}

#[derive(Subcommand)]
enum Op {
    /// Send command
    Send { host: String, port: u16 },
    /// Listen command
    Listen {
        /// Port to listen on
        port: u16,
    }
}

fn write_all(mut from: impl Read, mut to: impl Write) -> io::Result<usize> {
    let mut buf = [0_u8; 1024];
    let mut total = 0;
    while let Ok(n) = from.read(&mut buf) {
        if n == 0 { break; }
        total += n;
        to.write_all(&buf[0..n])?;
    }
    Ok(total)
}

fn tcp(conf: Config) -> crate::Result<()> {
    match conf.command {
        Op::Send { host, port } => {
            let stream = TcpStream::connect(format!("{}:{}",host,port))?;
            let stdin = stdin().lock();
            write_all(stdin, stream)?;
        }
        Op::Listen { port } => {
            let ip = IpAddr::V6([0_u8; 16].into());
            let socket = SocketAddr::new(ip, port);
            let stream = TcpListener::bind(socket)?;
            let mut stdout = stdout().lock();
            println!("Listening on {socket}");
            let stream = stream.incoming().next().unwrap()?;
            write_all(stream, &mut stdout)?;
        }
    }
    Ok(())
}

fn udp(conf: Config) -> crate::Result<()> {
    let mut buf = [0_u8; 65600];

    match conf.command {
        Op::Send { host, port } => {
            let stream = UdpSocket::bind("0.0.0.0:0")?;
            let target = format!("{}:{}",host,port).to_socket_addrs()?.next().unwrap();

            let mut stdin = stdin().lock();

            while let Ok(n) = stdin.read(&mut buf) {
                if n == 0 { break }
                let mut sent = 0;
                loop {
                    sent += stream.send_to(&buf[sent..n], target)?;
                    if sent == n { break }
                }
            }
        }
        Op::Listen { port } => {
            let ip = IpAddr::V6([0_u8; 16].into());
            let socket = SocketAddr::new(ip, port);
            let listener = UdpSocket::bind(socket)?;
            let mut stdout = stdout().lock();
            let (_,addr) = listener.peek_from(&mut buf)?;
            listener.connect(addr)?;
            while let Ok(n) = listener.recv(&mut buf) {
                if n == 0 { break }
                stdout.write_all(&buf[..n])?;
            }
        }
    }
    Ok(())
}

pub fn run() -> crate::Result<()> {
    let conf = Config::parse();

    if conf.udp {
        udp(conf)
    } else {
        tcp(conf)
    }
}