~halzy/waitforit

ebdd33fd83086c2e593e39b6b9ab89b8081c17ba — Benjamin Halsted a month ago 8eadd3a master
Added tcp timeout settings.
2 files changed, 26 insertions(+), 4 deletions(-)

M src/arguments.rs
M src/main.rs
M src/arguments.rs => src/arguments.rs +14 -1
@@ 7,6 7,7 @@ pub struct Arguments {
    pub listen: SocketAddr,
    pub connect: SocketAddr,
    pub delay: Duration,
    pub timeout: Duration,
}

impl Arguments {


@@ 39,19 40,31 @@ impl Arguments {
                    .takes_value(true)
                    .help("How many seconds the data is delayed for"),
            )
            .arg(
                Arg::with_name("timeout")
                    .long("timeout")
                    .short("t")
                    .required(true)
                    .takes_value(true)
                    .help("How many seconds the sockets will wait before dropping."),
            )
            .get_matches();

        // unwraps are never hit because get_matches() above will exit for these required values
        let listen: SocketAddr = matches.value_of("listen").unwrap().parse()?;
        let connect: SocketAddr = matches.value_of("connect").unwrap().parse()?;
        let delay: u64 = matches.value_of("delay").unwrap().parse()?;

        let delay: u64 = matches.value_of("delay").unwrap().parse()?;
        let delay = Duration::from_secs(delay);

        let timeout: u64 = matches.value_of("timeout").unwrap().parse()?;
        let timeout = Duration::from_secs(timeout);

        Ok(Arguments {
            listen,
            connect,
            delay,
            timeout,
        })
    }
}

M src/main.rs => src/main.rs +12 -3
@@ 20,14 20,22 @@ use std::{
};
use std::{net::SocketAddr, sync::atomic::AtomicU32};

fn set_delay(stream: TcpStream, timeout: Duration) -> IoResult<TcpStream> {
    let stream = stream.into_std()?;
    stream.set_read_timeout(Some(timeout))?;
    stream.set_write_timeout(Some(timeout))?;
    TcpStream::from_std(stream)
}

#[tokio::main(flavor = "current_thread")]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let args = Arguments::parse()?;

    let (tx, rx) = unbounded_channel::<(u32, Bytes)>();

    let incoming = accept(args.listen).await?;
    let incoming = accept(args.listen, args.timeout).await?;
    let outgoing = TcpStream::connect(args.connect).await?;
    let outgoing = set_delay(outgoing, args.timeout)?;

    let buffer_packets = Arc::new(AtomicU32::new(0));



@@ 115,7 123,8 @@ fn read_incoming(
    })
}

async fn accept(address: SocketAddr) -> IoResult<TcpStream> {
async fn accept(address: SocketAddr, timeout: Duration) -> IoResult<TcpStream> {
    let listener = TcpListener::bind(address).await?;
    listener.accept().await.map(|(t, _)| t)
    let (socket, _) = listener.accept().await?;
    set_delay(socket, timeout)
}