TCP/UDP open/setsockopt/bind/listen/connect/accept/send/recv/shutdown/close by libc crate

rust-libc-tcpudp.rs

use libc::*;
use std::convert::TryInto;
use structopt::StructOpt;

#[derive(StructOpt, Debug)]
struct Cli {
    #[structopt(subcommand)]
    cmd: Command,
}

#[derive(StructOpt, Debug)]
enum Command {
    Tcp(Tcp),
    Udp(Udp),
}

#[derive(StructOpt, Debug)]
struct Tcp {
    #[structopt(short, long)]
    src_port: Option<u16>,
    #[structopt(short, long)]
    dst_port: u16,
    #[structopt(short, long, default_value = "Hello, world!")]
    message: String,
    #[structopt(short, long, default_value = "100")]
    buffer_size: usize,
}

#[derive(StructOpt, Debug)]
struct Udp {
    #[structopt(short, long)]
    src_port: Option<u16>,
    #[structopt(short, long)]
    dst_port: u16,
    #[structopt(short, long, default_value = "Hello, world!")]
    message: String,
    #[structopt(short, long, default_value = "100")]
    buffer_size: usize,
}

fn main() {
    let args = Cli::from_args();
    println!("{:?}", args);

    match match args.cmd {
        Command::Tcp(opts) => unsafe {
            do_tcp_test(opts.src_port, opts.dst_port, opts.message, opts.buffer_size)
        },
        Command::Udp(opts) => unsafe {
            do_udp_test(opts.src_port, opts.dst_port, opts.message, opts.buffer_size)
        },
    } {
        Err(errno) => println!("errno: {}", errno_msg(errno)),
        Ok(()) => println!("ok"),
    };
}

unsafe fn errno() -> c_int {
    *__errno_location()
}

fn errno_msg(errno: c_int) -> String {
    let s = unsafe {
        let s = strerror(errno);
        let len = strlen(s);
        let s = s as *mut u8;
        std::slice::from_raw_parts_mut(s, len)
    };
    String::from(std::str::from_utf8_mut(s).unwrap())
}

fn htons(port: u16) -> in_port_t {
    port.swap_bytes()
}
#[allow(dead_code)]
fn ntohs(port: in_port_t) -> u16 {
    port.swap_bytes()
}

fn inet_addr(addr: &str) -> in_addr {
    in_addr {
        s_addr: u32::from(addr.parse::<std::net::Ipv4Addr>().unwrap()).to_be(),
    }
}

unsafe fn map_errno<E: std::fmt::Debug, T: TryInto<isize, Error = E> + Clone>(
    ret: T,
) -> Result<T, c_int> {
    match ret.clone().try_into().unwrap() {
        -1 => Err(errno()),
        _ => Ok(ret),
    }
}

unsafe fn do_tcp_test(src: Option<u16>, dst: u16, msg: String, buflen: usize) -> Result<(), c_int> {
    let sk_listen = map_errno(socket(AF_INET, SOCK_STREAM, 0))?;

    let yes = 1u32;
    map_errno(setsockopt(
        sk_listen,
        SOL_SOCKET,
        SO_REUSEADDR,
        &yes as *const u32 as *const c_void,
        std::mem::size_of::<u32>() as socklen_t,
    ))?;

    let mut sin = sockaddr_in {
        sin_family: AF_INET as sa_family_t,
        sin_port: htons(dst),
        sin_addr: inet_addr("127.0.0.1"),
        sin_zero: Default::default(),
    };
    map_errno(bind(
        sk_listen,
        &sin as *const sockaddr_in as *const sockaddr,
        std::mem::size_of::<sockaddr_in>() as socklen_t,
    ))?;

    map_errno(listen(sk_listen, SOMAXCONN))?;

    let sk_connect = map_errno(socket(AF_INET, SOCK_STREAM, 0))?;

    if let Some(src) = src {
        let yes = 1u32;
        map_errno(setsockopt(
            sk_connect,
            SOL_SOCKET,
            SO_REUSEADDR,
            &yes as *const u32 as *const c_void,
            std::mem::size_of::<u32>() as socklen_t,
        ))?;

        let sin = sockaddr_in {
            sin_family: AF_INET as u16,
            sin_port: htons(src),
            sin_addr: inet_addr("127.0.0.1"),
            sin_zero: Default::default(),
        };
        map_errno(bind(
            sk_connect,
            &sin as *const sockaddr_in as *const sockaddr,
            std::mem::size_of::<sockaddr_in>() as socklen_t,
        ))?;
    }

    map_errno(connect(
        sk_connect,
        &sin as *const sockaddr_in as *const sockaddr,
        std::mem::size_of::<sockaddr_in>() as u32,
    ))?;

    let mut sin_len = std::mem::size_of::<sockaddr_in>() as socklen_t;
    let sk_accept = map_errno(accept(
        sk_listen,
        &mut sin as *mut sockaddr_in as *mut sockaddr,
        &mut sin_len as *mut socklen_t,
    ))?;

    let sent = map_errno(send(
        sk_connect,
        msg.as_ptr() as *const c_void,
        msg.len(),
        0,
    ))?;
    assert_eq!(sent as usize, msg.len());

    let mut buf = vec![0u8; buflen];
    let received = map_errno(recv(
        sk_accept,
        buf.as_mut_ptr() as *mut c_void,
        buf.len(),
        0,
    ))?;
    assert_eq!(received as usize, msg.len().min(buflen));

    let msg = String::from_utf8(buf[..received as usize].into()).unwrap();
    println!("Received: {}", msg);

    map_errno(shutdown(sk_connect, SHUT_RDWR))?;
    map_errno(shutdown(sk_accept, SHUT_RDWR))?;
    map_errno(shutdown(sk_listen, SHUT_RDWR))?;

    map_errno(close(sk_connect))?;
    map_errno(close(sk_accept))?;
    map_errno(close(sk_listen))?;

    println!("TCP OK");
    Ok(())
}

unsafe fn do_udp_test(src: Option<u16>, dst: u16, msg: String, buflen: usize) -> Result<(), c_int> {
    let sk_dst = map_errno(socket(AF_INET, SOCK_DGRAM, 0))?;

    let sin = sockaddr_in {
        sin_family: AF_INET as sa_family_t,
        sin_port: htons(dst),
        sin_addr: inet_addr("127.0.0.1"),
        sin_zero: Default::default(),
    };
    map_errno(bind(
        sk_dst,
        &sin as *const sockaddr_in as *const sockaddr,
        std::mem::size_of::<sockaddr_in>() as socklen_t,
    ))?;

    let sk_src = map_errno(socket(AF_INET, SOCK_DGRAM, 0))?;

    if let Some(src) = src {
        let sin = sockaddr_in {
            sin_family: AF_INET as u16,
            sin_port: htons(src),
            sin_addr: inet_addr("127.0.0.1"),
            sin_zero: Default::default(),
        };
        map_errno(bind(
            sk_src,
            &sin as *const sockaddr_in as *const sockaddr,
            std::mem::size_of::<sockaddr_in>() as socklen_t,
        ))?;
    }

    map_errno(connect(
        sk_src,
        &sin as *const sockaddr_in as *const sockaddr,
        std::mem::size_of::<sockaddr_in>() as u32,
    ))?;

    let sent = map_errno(send(sk_src, msg.as_ptr() as *const c_void, msg.len(), 0))?;
    assert_eq!(sent as usize, msg.len());

    let mut buf = vec![0u8; buflen];
    let received = map_errno(recv(
        sk_dst,
        buf.as_mut_slice() as *mut [u8] as *mut c_void,
        buf.len(),
        0,
    ))?;
    assert_eq!(received as usize, msg.len().min(buf.len()));

    let msg = String::from_utf8(buf[..received as usize].into()).unwrap();
    println!("Received: {}", msg);

    map_errno(close(sk_src))?;
    map_errno(close(sk_dst))?;

    println!("UDP OK");
    Ok(())
}